Shuffle в YTsaurus Flow (Python)

Пример пайплайна из двух компьютейшенов: source-компьютейшен парсит JSON-данные и отправляет типизированные сообщения, transform-компьютейшен ведёт подсчёт событий во внешнем стейте.

Исходный код

Структура

reader (source) -- EventMapper: парсит JSON из поля data и отправляет типизированные сообщения в стрим event.
reducer (transform) -- EventReducer: подсчитывает количество событий по ключу с использованием external state.

`main.py`

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

    pipeline = Pipeline()
    pipeline.add("reader", EventMapper(), source=True)
    pipeline.add("reducer", EventReducer())

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

`event_mapper.py`

Source-функция, которая парсит JSON из поля data входного сообщения и создаёт типизированное сообщение через ctx.message_builder():

class EventMapper(RowFunction):
    """Parses JSON from data field and emits event messages with key fields."""

    def on_message(self, message, output, ctx):
        data = message.payload["data"]
        try:
            parsed = json.loads(data)
        except (json.JSONDecodeError, TypeError) as e:
            log.error("Error parsing json from data field: %s", e)
            raise

        builder = ctx.message_builder("event")
        builder.set("value", parsed.get("value"))
        builder.set("key_a", parsed.get("key_a"))
        builder.set("key_b", parsed.get("key_b"))
        builder.set("key_c", parsed.get("key_c"))
        builder.set("key_d", parsed.get("key_d"))
        output.add_message(builder.finish())

`event_reducer.py`

Transform-функция с external state для подсчёта событий:

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())

Ключевые паттерны

Пайплайн из нескольких компьютейшенов: source + transform.
Source-компьютейшен с source=True для чтения из внешнего источника.
Парсинг JSON и создание типизированных сообщений через ctx.message_builder().
External state с паттерном to_builder() / set() / finish().

См. также

Word Count

Async Request