Word Count in YTsaurus Flow (Python)

Use this simple stateful-pipeline example in Python to count how many times each word appears. You’ll use an internal YSON-state for the counts.

Source code

Structure

The pipeline includes a single transform-computation called mapper. You use it to read words from the input stream and update the counter in the state.

__main__.py

This is the entry point. You create the pipeline and register the only computation.

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

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

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

word_count_mapper.py

This RowFunction uses ctx.state() to work with the YSON-state. For each message key, you store a Map with the word and count fields.

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)

Key patterns

  • A simple stateful pipeline with a single computation.
  • Internal YSON-state via ctx.state() with get_or_default and set.
  • The state key is defined by group_by_schema from the spec (in this case, by the word field).

See also

Previous
Testing
Next
Shuffle