Shuffle in YTsaurus Flow (Python)

Use this pipeline example with two computations: the source computation parses JSON data and sends typed messages, and the transform computation counts events in an external state.

Source code

Structure

  • reader (source) -- EventMapper: parses JSON from the data field and sends typed messages to the event stream.
  • reducer (transform) -- EventReducer: counts the number of events by key using 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

The source function parses JSON from the data field of the input message and creates a typed message using 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

The transform function uses external state to count events:

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

Key patterns

  • A pipeline with multiple computations: source + transform.
  • A source computation with source=True to read from an external source.
  • Parsing JSON and creating typed messages using ctx.message_builder().
  • External state with the to_builder() / set() / finish() pattern.

See also

Previous
Word Count
Next
Async Request