Computation in YTsaurus Flow (Python)

Computation Types

In Flow, there are two types of Computation: Swift and Transform. Your choice determines how exactly-once guarantees are provided and what transformations you can implement.

What determinism means and what breaks if you violate it is explained in the Determinism requirement section.

When using a companion, you select Swift or Transform by specifying computation_class_name in the static spec:

• NYT::NFlow::NCompanion::TTransformCompanionComputation — for Transform.
• NYT::NFlow::NCompanion::TSwiftMapCompanionComputation — for Swift.
• NYT::NFlow::NCompanion::TTransformOrderedSourceCompanionComputation — for a Transform source.
• NYT::NFlow::NCompanion::TSwiftOrderedSourceCompanionComputation — for a Swift source.

Creating a Computation

In Python, you create a computation with Pipeline.add(), and it registers automatically in PipelineContext. See an example from WordCount:

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

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

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

In the static spec, you create a Computation with the same id (in this example, 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" = {
        ...
    };
};

For more on specs, see the Spec, DynamicSpec, and Config section.

SourceComputation

SourceComputation is the top node in the pipeline graph that reads data from external sources. On the worker side, it corresponds to TSwiftOrderedSourceComputation or TTransformOrderedSourceComputation.

In Python, you create a SourceComputation by passing source=True to Pipeline.add(). You filter messages inside the Process Function using the distribute flag.

For deterministic processing without user state, set TSwiftOrderedSourceCompanionComputation in the static spec. If the source uses internal state or nondeterministic logic, use TTransformOrderedSourceCompanionComputation: the worker materializes its output and commits that output atomically with state and the source offset. The internal-state key is the source partition key.

Parameters

Creating a SourceComputation

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

For a passthrough Source, don’t use Python. Instead, specify NYT::NFlow::TSwiftPassthroughOrderedSourceComputation in computation_class_name in the spec, and leave the computation unregistered in the Python companion. See Passthrough Computation for more details.

Interaction with Worker

When initializing, the Worker requests information about registered Computation and SourceComputation objects from the Python companion. The source computation on the worker side sends input messages to the Python companion, which applies the ProcessFunction and returns the result.

Process Function

You implement the business logic for data processing in a Process Function. To do this, choose one of two base classes: RowFunction or BatchFunction.

RowFunction

RowFunction receives messages and timers one at a time. The class provides two methods:

• on_message(message, output, ctx) — called for each input message.
• on_timer(timer, output, ctx) — called when a timer fires (optional).

Example of a stateless function

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

Let’s go line by line:

1. ctx.message_builder("x2_numbers") — creates a MessageBuilder for the output stream with id = x2_numbers. A stream with this identifier must be present in the output_stream_ids list in the static spec of the computation.
2. message.payload["number"] — gets the value of the number field from the incoming message. The payload supports dict-like access to fields.
3. builder.set("number_x2", number * 2) — writes the value to the number_x2 field. This field must exist in the schema of the x2_numbers stream in the static spec.
4. output.add_message(builder.finish()) — the finish() method returns the completed message and resets the builder for reuse. The message is added to the OutputCollector.

BatchFunction

BatchFunction receives the entire list of messages and timers that came from the worker. The class provides two methods:

• on_messages(messages, output, ctx) — called for a batch of messages.
• on_timers(timers, output, ctx) — called for a batch of timers (optional).

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])
        ...  # Process the group, emit via out.

Example of a batch function

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

The key difference from RowFunction is:

• You create the MessageBuilder once for the entire batch (line 1).
• The finish() method returns the completed message and resets the MessageBuilder to its initial state, so you can reuse it for the next message (line 5).

Message Filtering

To filter messages in source computations, use the per-message distribute flag: you emit the message from the Process Function with distribute=False, and it isn’t published further along the graph, but it’s still considered when evaluating the watermark.

Registration in Pipeline

You register all computations with Pipeline.add(). Inside, Pipeline uses PipelineContext to store and manage registered objects.

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

pipeline = Pipeline()

# Transform computation
pipeline.add("computation_id", my_function)

# Source computation
pipeline.add("reader", my_function, source=True)

You can also use the @pipeline.computation decorator:

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

Source code

RuntimeContext (ctx) gives you access to the computation’s runtime context. Key methods:

For more on working with states, see the Working with States (Python) section.

MessageBuilder

You use MessageBuilder to create output messages:

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

The finish() method returns the completed Message object and resets the builder for reuse. The stream_id field must be present in the output_stream_ids list in the static spec of the computation.

Computation Parameters

wait_for_actions = ctx.parameters["wait_for_actions"]

The ctx.parameters map contains the parameters specified in the static spec of the computation.

Watermarks

# Minimum watermark across all input streams
min_wm = ctx.min_watermark

# Watermark of a specific stream (int or None)
stream_wm = ctx.watermark("stream_id")

OutputCollector

Source code

Use OutputCollector to send processing results:

For a Swift computation, MessageIdSuffix.payload_hash() or MessageIdSuffix.user_defined(value) makes message identity independent of emission order:

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

Default options use the message sequence number. Payload-hash and user-defined suffixes are supported only by Swift computations.

Example of creating an output message and a timer:

def on_message(self, message, output, ctx):
    # Create an output message
    builder = ctx.message_builder("output_stream")
    builder.set("field", value)
    output.add_message(builder.finish())

    # Create a timer
    output.add_timer(trigger_timestamp=1000, event_timestamp=500)

ExtendedMessage

An incoming message (ExtendedMessage) contains:

• message.payload — Payload with dict-like access to fields: message.payload["field"].
• message.stream_id — The ID of the input stream (str).
• message.key — The key of the message (Payload) from group_by_schema: message.key["field"].
• message.event_timestamp — The event timestamp of the message (int).

Timer

A timer object (Timer) contains:

• timer.key — The key of the timer (Payload): timer.key["field"].
• timer.stream_id — The ID of the timer’s stream (str).
• timer.trigger_timestamp — The trigger time (int).
• timer.event_timestamp — The event timestamp (int).

CompanionManager resource configuration

To run a Python companion, you need to declare the CompanionManager resource in a static spec:

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

The resource_class_name parameter specifies the resource class that will launch the companion.
For a Python companion, resource_class_name must always be NYT::NFlow::NCompanion::TCompanionManager.

The companion process is described by the entrypoint parameter (executable, args, env). When you launch a pipeline from a host via pipeline.run(), you don’t need to fill in entrypoint manually: the Python binary sets entrypoint = {"executable" = "./py_companion"} for every TCompanionManager resource in the spec, and flow_server delivers the binary to the job under that name.

The key difference from the Java configuration is that Java has a separate resource class NYT::NFlow::NCompanion::TJavaCompanionManager with parameters jdk_bin_path, classpath, and main_class, while the Python companion uses the shared TCompanionManager with the entrypoint parameter.

For more details about the spec, see the section Spec, DynamicSpec and Config.

See also

• Computation (concept)
• Working with states (Python)
• Quick start (Python)
• Companion