Работа со стейтами в YTsaurus Flow (Python)

Примечание

Данная страница описывает Python API для работы со стейтами. Общие концепции стейтов описаны в разделе Stateful-вычисления.

YSON State

Самый простой способ работы со стейтом -- YSON-формат. Стейт хранится как словарь Python, автоматически сериализуемый в YSON:

state = ctx.state("state-name", message)

Возвращает YsonStateAccessor с методами:

  • get() -- получить текущее значение (dict или None).
  • set(dict) -- сохранить значение.
  • clear() -- удалить стейт.
  • get_or_default(dict) -- получить текущее значение или вернуть значение по умолчанию.

Пример из WordCount:

class WordCountMapper(RowFunction):
    """Processes Word messages and maintains word count state."""

    def on_message(self, message, output, ctx):
        state = ctx.state("word-state", message)
        data = state.get_or_default({"word": message.payload["word"], "count": 0})
        data["count"] += 1
        state.set(data)

Здесь стейт привязан к ключу сообщения (определяемому через group_by_schema в спеке). Для каждого уникального ключа хранится свой независимый стейт.

Raw State

Для хранения стейта в виде сырых байтов:

state = ctx.raw_state("state-name", message)

Возвращает RawStateAccessor с методами:

  • get() -- получить значение (bytes или None).
  • set(bytes) -- сохранить значение.
  • clear() -- удалить стейт.
  • get_or_default(bytes) -- получить значение или вернуть значение по умолчанию.

Proto State

Для хранения стейта в виде Protobuf-сообщения:

state_accessor = ctx.proto_state("state-name", message, TJoinState)

Возвращает ProtoStateAccessor с методами:

  • get() -- десериализовать и вернуть Protobuf-объект (или None).
  • set(proto_message) -- сериализовать и сохранить.
  • clear() -- удалить стейт.
  • get_or_default(default=None) -- получить значение или вернуть default. Если default не указан, возвращает пустой экземпляр Proto-класса.

External State

Внешний стейт работает как Payload -- предоставляет dict-like доступ к полям:

state = ctx.external_state("/state-name", message)

Имя стейта обязательно начинается с / и должно совпадать с ключом в external_state_managers статической спеки. Вызов ctx.external_state("state-name", message) (без ведущего /) бросит ValueError.

Возвращает ExternalStateAccessor, который одновременно является Payload:

  • state.get("field") -- чтение поля.
  • state["field"] -- dict-like чтение поля.
  • state.to_builder() -- получить PayloadBuilder с текущими значениями.
  • state.set(payload) -- сохранить новое значение (принимает Payload от builder.finish()).
  • state.clear() -- удалить стейт.

Пример из Shuffle (EventReducer):

class EventReducer(RowFunction):
    """Counts events per key using external state."""

    def on_message(self, message, output, ctx):
        state = ctx.external_state("/shuffle-state", message)
        builder = state.to_builder()
        builder.set("count", (state.get("count") or 0) + 1)
        state.set(builder.finish())

Паттерн работы с external state:

  1. Получить текущее состояние через ctx.external_state(...).
  2. Создать билдер через state.to_builder().
  3. Обновить нужные поля через builder.set(...).
  4. Сохранить через state.set(builder.finish()).

Стейт в таймерах

API для работы со стейтом в обработчике таймеров идентичен -- вместо message передаётся объект timer:

def on_timer(self, timer, output, ctx):
    state = ctx.external_state("/join-state", timer)
    # Чтение стейта
    show_time = state.get("show_time")
    hit_payload = state.get("hit_payload")
    # Очистка стейта после обработки
    state.clear()

Пример из WaitClickJoin (JoinProcessFunction):

    def on_timer(self, timer, output, ctx):
        hit_id = timer.key["hit_id"]
        log.debug("Process timer (KeyHitId: %s, Timer: %s)", hit_id, timer)

        state = ctx.external_state("/join-state", timer)

        log.debug("Process state (KeyHitId: %s, State: %s)", hit_id, state)

        if timer.stream_id == "timer":
            show_time = state.get("show_time")
            hit_payload = state.get("hit_payload")

            if show_time is not None and hit_payload is not None and show_time != 0:
                log.debug("Click join: %s", state)
                hit_time = timer.key["hit_time"]
                click_time = state.get("click_time")

                builder = ctx.message_builder("joined_action")
                builder.set("hit_id", hit_id)
                builder.set("hit_time", hit_time)
                builder.set("show_time", show_time)
                builder.set("hit_payload", hit_payload)

                if click_time is not None and click_time != 0:
                    builder.set("is_click", True)
                    builder.set("click_time", click_time)
                else:
                    builder.set("is_click", False)
                    builder.set("click_time", 0)

                output.add_message(builder.finish())

            state.clear()
        else:
            raise ValueError(f"Unknown streamId: {timer.stream_id}")

Привязка стейта к ключу

В TTransformCompanionComputation стейт привязывается к ключу сообщения, определяемому через group_by_schema в спеке компьютейшена. Все сообщения с одинаковым ключом разделяют один стейт.

В TTransformOrderedSourceCompanionComputation поле group_by_schema не поддерживается. Ключом внутреннего стейта служит ключ партиции источника, поэтому все сообщения одной партиции разделяют стейт. Подробнее о выборе класса для SourceComputation см. в разделе Computation (Python).

Подробнее о конфигурации ключей см. Stateful-вычисления.

Конфигурация стейтов в спеке

Внутренние стейты (YSON, Raw, Proto) должны быть объявлены в параметрах компьютейшена в секции internal_states. Имя стейта в коде (первый аргумент ctx.state(...)) должно совпадать с именем, объявленным в спеке.

Внешние стейты (External) конфигурируются через секцию external_state_managers в спеке компьютейшена и имеют собственную схему, описывающую доступные поля. Ключ внутри external_state_managers (например, "/shuffle-state") задаёт имя стейта, начинающееся с /; в поле external_state_manager_class_name указывается зарегистрированный класс менеджера (для типового сценария — "NYT::NFlow::TSimpleExternalStateManager"). Подробнее про спеку и доступные менеджеры см. в External State и в C++ документации.

См. также

Предыдущая
Флаг distribute
Следующая
State Accessor