Computation в YTsaurus Flow (Python)

Типы Computation

Во Flow есть два вида Computation: Swift и Transform. От их выбора зависит способ обеспечения exactly-once гарантий и то, какие преобразования возможно реализовать с их применением.

Что именно означает детерминированность и что ломается при её нарушении — в разделе Требование детерминированности.

При использовании компаньона выбор Swift или Transform осуществляется через указание computation_class_name в статической спеке:

• NYT::NFlow::NCompanion::TTransformCompanionComputation — для Transform.
• NYT::NFlow::NCompanion::TSwiftMapCompanionComputation — для Swift.
• NYT::NFlow::NCompanion::TTransformOrderedSourceCompanionComputation — для Transform-сорса.
• NYT::NFlow::NCompanion::TSwiftOrderedSourceCompanionComputation — для Swift-сорса.

Создание Computation

В Python компьютейшен создаётся через Pipeline.add() и регистрируется в PipelineContext автоматически. Пример из WordCount:

def main():
    log.info("Starting companion execution")

    pipeline = Pipeline()
    pipeline.add("mapper", WordCountMapper())

    log.info("Starting pipeline...")
    pipeline.run()
    log.info("Pipeline completed")

В статической спеке создаётся Computation с таким же id (в данном примере mapper):

"mapper" = {
    "computation_class_name" = "NYT::NFlow::NCompanion::TTransformCompanionComputation";
    "group_by_schema" = [
        ...
    ];
    "input_stream_ids" = [...];
    "output_stream_ids" = [...];
    "required_resource_ids" = {
        "CompanionManager" = {
            "worker" = true;
            "controller" = false;
        };
    };
    "parameters" = {
        ...
    };
};

Подробнее про спеку в разделе Spec, DynamicSpec и Config.

SourceComputation

SourceComputation — вершина в графе пайплайна, осуществляющая чтение данных из внешних источников. На стороне воркера ей соответствует TSwiftOrderedSourceComputation или TTransformOrderedSourceComputation.

В Python SourceComputation создаётся через передачу source=True в Pipeline.add(). Фильтрация сообщений выполняется внутри Process Function через флаг distribute.

В статической спеке для детерминированной обработки без пользовательского стейта указывают TSwiftOrderedSourceCompanionComputation. Если SourceComputation использует внутренний стейт или недетерминированную логику, указывают TTransformOrderedSourceCompanionComputation: воркер материализует выход и фиксирует его вместе со стейтом и смещением источника. Ключ внутреннего стейта в таком компьютейшене — ключ партиции источника.

Параметры

Создание SourceComputation

pipeline.add("reader", MyParsingFunction(), source=True)

Для passthrough Source не используйте Python — укажите в спеке NYT::NFlow::TSwiftPassthroughOrderedSourceComputation в computation_class_name и оставьте компьютейшен незарегистрированным в Python-компаньоне. Подробнее — Passthrough Computation.

Взаимодействие с Worker

При инициализации Worker запрашивает у Python-компаньона информацию о зарегистрированных объектах Computation и SourceComputation. Source-компьютейшен на стороне воркера отправляет входные сообщения в Python-компаньон, который применяет к ним ProcessFunction и возвращает результат.

Process Function

Бизнес-логика обработки данных реализуется через Process Function. Для реализации необходимо выбрать один из двух базовых классов: RowFunction или BatchFunction.

RowFunction

RowFunction получает сообщения и таймеры по одному. Класс предоставляет два метода:

• on_message(message, output, ctx) — вызывается для каждого входного сообщения.
• on_timer(timer, output, ctx) — вызывается при срабатывании таймера (опционально).

Пример stateless-функции

from yt.yt.flow.library.python.companion.computation import RowFunction


class X2Mapper(RowFunction):
    def on_message(self, message, output, ctx):
        builder = ctx.message_builder("x2_numbers")        # 1
        number = message.payload["number"]                  # 2
        builder.set("number_x2", number * 2)                # 3
        output.add_message(builder.finish())                # 4

Разберём построчно:

1. ctx.message_builder("x2_numbers") — создаётся MessageBuilder для output-стрима с id = x2_numbers. Стрим с таким идентификатором должен присутствовать в списке output_stream_ids в статической спеке компьютейшена.

2. message.payload["number"] — получаем значение поля number из входящего сообщения. Payload поддерживает dict-like доступ к полям.

3. builder.set("number_x2", number * 2) — записываем значение в поле number_x2. Это поле должно присутствовать в схеме стрима x2_numbers в статической спеке.

4. output.add_message(builder.finish()) — метод finish() возвращает готовое сообщение и сбрасывает билдер для повторного использования. Сообщение добавляется в OutputCollector.

BatchFunction

BatchFunction получает весь список сообщений и таймеров, пришедших от воркера. Класс предоставляет два метода:

• on_messages(messages, output, ctx) — вызывается для батча сообщений.
• on_timers(timers, output, ctx) — вызывается для батча таймеров (опционально).

def on_messages(self, messages, output, ctx):
    groups = {}
    for message in messages:
        groups.setdefault(message.key["user_id"], []).append(message)
    for key, group in groups.items():
        out = output.set_parent_ids([m.message_id for m in group])
        ...  # Обработка группы, вывод через out.

Пример batch-функции

from yt.yt.flow.library.python.companion.computation import BatchFunction


class X2BatchMapper(BatchFunction):
    def on_messages(self, messages, output, ctx):
        builder = ctx.message_builder("x2_numbers")         # 1
        for message in messages:                             # 2
            number = message.payload["number"]               # 3
            builder.set("number_x2", number * 2)             # 4
            output.add_message(builder.finish())             # 5

Ключевое отличие от RowFunction:

• MessageBuilder создаётся один раз за весь батч (строка 1).
• Метод finish() одновременно с возвратом готового сообщения сбрасывает MessageBuilder в исходное состояние, после чего его можно переиспользовать для следующего сообщения (строка 5).

Фильтрация сообщений

Для фильтрации сообщений в source-компьютейшенах используется per-message-флаг distribute: сообщение эмитится из Process Function с distribute=False и не публикуется дальше по графу, но учитывается при оценке watermark.

Регистрация в Pipeline

Все компьютейшены регистрируются через Pipeline.add(). Внутри Pipeline использует PipelineContext для хранения и управления зарегистрированными объектами.

from yt.yt.flow.library.python.companion import Pipeline

pipeline = Pipeline()

# Transform-компьютейшен
pipeline.add("computation_id", my_function)

# Source-компьютейшен
pipeline.add("reader", my_function, source=True)

Также доступен декоратор @pipeline.computation:

pipeline = Pipeline()

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

RuntimeContext

Исходный код

RuntimeContext (ctx) предоставляет доступ к контексту выполнения компьютейшена. Основные методы:

Подробнее про работу со стейтами — в разделе Работа со стейтами (Python).

MessageBuilder

Для создания выходных сообщений используется MessageBuilder:

builder = ctx.message_builder("stream_id")
builder.set("field_name", value)
message = builder.finish()
output.add_message(message)

Метод finish() возвращает готовый объект Message и сбрасывает билдер для повторного использования. Поле stream_id должно присутствовать в списке output_stream_ids в статической спеке компьютейшена.

Параметры компьютейшена

wait_for_actions = ctx.parameters["wait_for_actions"]

Словарь ctx.parameters содержит параметры, указанные в статической спеке компьютейшена.

Вотермарки

# Минимальный вотермарк по всем входным стримам
min_wm = ctx.min_watermark

# Вотермарк конкретного стрима (int или None)
stream_wm = ctx.watermark("stream_id")

OutputCollector

Исходный код

OutputCollector используется для отправки результатов обработки:

В Swift-компьютейшене MessageIdSuffix.payload_hash() или MessageIdSuffix.user_defined(value) позволяет отвязать идентичность сообщения от порядка эмита:

output.add_message(
    message,
    AddMessageOptions(message_id_suffix=MessageIdSuffix.payload_hash()),
)

Options по умолчанию используют порядковый номер сообщения. Хеш payload и пользовательский суффикс поддерживаются только Swift-компьютейшенами.

Пример создания выходного сообщения и таймера:

def on_message(self, message, output, ctx):
    # Создание выходного сообщения
    builder = ctx.message_builder("output_stream")
    builder.set("field", value)
    output.add_message(builder.finish())

    # Создание таймера
    output.add_timer(trigger_timestamp=1000, event_timestamp=500)

ExtendedMessage

Входящее сообщение (ExtendedMessage) содержит:

• message.payload — Payload с dict-like доступом к полям: message.payload["field"].
• message.stream_id — идентификатор входного стрима (str).
• message.key — ключ сообщения (Payload) из group_by_schema: message.key["field"].
• message.event_timestamp — event timestamp сообщения (int).

Timer

Объект таймера (Timer) содержит:

• timer.key — ключ таймера (Payload): timer.key["field"].
• timer.stream_id — идентификатор стрима таймера (str).
• timer.trigger_timestamp — время срабатывания (int).
• timer.event_timestamp — event timestamp (int).

Конфигурация ресурса CompanionManager

Для запуска Python-компаньона необходимо объявить ресурс CompanionManager в статической спеке:

"CompanionManager" = {
    "resource_class_name" = "NYT::NFlow::NCompanion::TCompanionManager";
    "parameters" = {
        "entrypoint" = {
            "executable" = "./py_companion";
        };
    };
    "dependencies" = {};
};

Параметр resource_class_name указывает на класс ресурса, который будет осуществлять запуск компаньона.
В случае Python-компаньона resource_class_name всегда должен быть NYT::NFlow::NCompanion::TCompanionManager.

Процесс компаньона описывается параметром entrypoint (executable, args, env); воркер сам запускает компаньон и следит за его жизненным циклом. При запуске пайплайна с хоста через pipeline.run() заполнять entrypoint вручную не нужно: Python-бинарь сам прописывает entrypoint = {"executable" = "./py_companion"}, а flow_server доставляет бинарь в джобу под этим именем.

Ключевое отличие от Java-конфигурации: для Java существует отдельный класс ресурса NYT::NFlow::NCompanion::TJavaCompanionManager с параметрами jdk_bin_path, classpath и main_class, тогда как Python-компаньон использует общий TCompanionManager с параметром entrypoint.

Подробнее про спеку в разделе Spec, DynamicSpec и Config.

См. также

• Computation (концепция)
• Работа со стейтами (Python)
• Быстрый старт (Python)
• Companion