Working with states in YTsaurus Flow (Go)

A state accessor is opened by the free function flow.OpenXxxState(rt, name, input) with three arguments:

• rt — flow.Runtime, the second argument of any handler (see Process Function).
• name — the state name declared in the computation spec.
• input — the input to whose key the accessor is bound: flow.ExtendedMessage, flow.Timer, or flow.Visit. All three implement the flow.Input interface, so working with state looks the same in a message, timer, and visit handler.

The opening functions are free functions rather than methods of Runtime: the YSON and Proto accessors are parameterized by the state type, and methods in Go have no type parameters of their own.

An accessor addresses the state of exactly one key and lives within a single request. Only the records the computation wrote to travel to the worker: reading a state doesn’t send it back.

YSON State

The simplest way to work with state is the YSON format. The state is described by an ordinary Go structure with yson tags:

state, err := flow.OpenYSONState[wordCountState](rt, "word-state", msg)
if err != nil {
    return err
}

if state.Empty() {
    state.Value().Word = word
}
state.Value().Count++
return nil

flow.YSONState[T] provides Empty(), Get() (*T, bool), Value() *T, Clear(), and ReadOnly(). Value() returns a mutable value and creates a zero value for a missing state. A separate Set isn’t needed: after the batch is processed successfully, the SDK serializes the changes automatically; if the handler fails, they are discarded.

type wordCountState struct {
	Word  string `yson:"word"`
	Count int64  `yson:"count"`
}

type wordCountMapper struct{}

var _ flow.RowFunction = (*wordCountMapper)(nil)

func (*wordCountMapper) OnMessage(
	ctx context.Context,
	rt flow.Runtime,
	msg flow.ExtendedMessage,
	out flow.OutputCollector,
) error {
	var input wordMessage
	if err := msg.ConvertTo(&input); err != nil {
		return err
	}

	state, err := flow.OpenYSONState[wordCountState](rt, wordStateName, msg)
	if err != nil {
		return err
	}

	fresh := state.Empty()
	counter := state.Value()
	if fresh {
		counter.Word = input.Word
	}
	counter.Count++
	return nil
}

The state is bound to the message key, which is defined through group_by_schema in the computation spec. An independent value is stored for every unique key.

Raw State

To store the state as raw bytes:

state, err := flow.OpenRawState(rt, "raw-state", msg)

Returns flow.RawStateAccessor with the methods:

• Get() ([]byte, bool) — get the value and the indication of its presence.
• Or(fallback []byte) []byte — get the value or fallback.
• Set(data []byte) error — save the value.
• Clear() error — delete the state.

The YSON and Proto states are wrappers over the raw accessor: RawStateAccessor itself is needed when the computation performs the serialization on its own.

Proto State

To store the state as a Protobuf message:

state, err := flow.OpenProtoState[TJoinState](rt, "join-state", msg)

The state type is named by its value form (TJoinState), and the state works with a pointer to it (*TJoinState) — it is on the pointer that the generated code implements proto.Message.

Returns *flow.ProtoState[T, PT] with the methods:

• Empty() bool — check whether the value is missing.
• Get() (PT, bool) — return the mutable message.
• Or(fallback PT) PT — return the saved message, or write and return fallback.
• Set(value PT) error — replace the whole value.
• Clear() — delete the state.
• ReadOnly() flow.ReadOnlyProtoState[T, PT] — a read-only view.

The message is changed in place: the changes are written at the end of the batch without a Set, see Changing the value in place.

state, err := flow.OpenProtoState[TJoinState](rt, "join-state", msg)
if err != nil {
    return err
}

window := state.Or(&TJoinState{})
window.HitPayload = payload

return nil

External State

An external state is a row of your own dynamic table. In user code it is represented by an ordinary Go structure with yson tags:

state, err := flow.OpenExternalState(rt, "/shuffle-state", msg)

The state name is an absolute path matching the key in the external_state_managers section of the static spec. A name without a leading / is rejected with the flow.ErrInvalidStateName error, an undeclared name with the flow.ErrUnknownState error.

Returns flow.ExternalStateAccessor. The main operations:

• ConvertTo(&value) (bool, error) — fill the structure with the saved row; bool distinguishes a missing row.
• ConvertFrom(&value) error — save the fields of the structure into the state row.
• Clear() error — delete the row.

The low-level Get, Or, Builder, Set, and Schema are needed only for dynamic schemas and column-wise processing.

An example from Shuffle:

type shuffleState struct {
	Count int64 `yson:"count"`
}

type eventReducer struct{}

var _ flow.RowFunction = (*eventReducer)(nil)

func (*eventReducer) OnMessage(
	ctx context.Context,
	rt flow.Runtime,
	msg flow.ExtendedMessage,
	out flow.OutputCollector,
) error {
	state, err := flow.OpenExternalState(rt, shuffleStateName, msg)
	if err != nil {
		return err
	}

	var counter shuffleState
	if _, err := state.ConvertTo(&counter); err != nil {
		return err
	}
	counter.Count++
	return state.ConvertFrom(&counter)
}

The pattern for working with external state:

1. Open the state through flow.OpenExternalState(...).
2. Convert the row into a structure through state.ConvertTo(&value).
3. Change the fields of the structure and save it through state.ConvertFrom(&value).

For more details, see the External State (Go) section.

Joined External State

If an external state is needed for reading only — for example, to enrich messages with a reference dataset — the computation declares it in the external_state_joiners section and opens it with a separate function:

reference, err := flow.OpenJoinedExternalState(rt, "/reference", msg)

Returns flow.JoinedExternalStateAccessor. The row is read into a structure through ConvertTo(&value); the accessor has no writing by design, because a joined state is filled from the request and never travels back. The low-level Get, Or, and Schema are available for dynamic schemas.

The namespaces don’t intersect: a state the computation owns isn’t available through flow.OpenJoinedExternalState, and vice versa.

type referenceState struct {
	Name *string `yson:"name"`
}

type lookupJoin struct{}

var _ flow.RowFunction = (*lookupJoin)(nil)

func (*lookupJoin) OnMessage(
	ctx context.Context,
	rt flow.Runtime,
	msg flow.ExtendedMessage,
	out flow.OutputCollector,
) error {
	var event eventMessage
	if err := msg.ConvertTo(&event); err != nil {
		return err
	}

	name, ok, err := referenceName(rt, msg)
	if err != nil || !ok {
		return err
	}

	enriched := flow.NewYSONMessage[enrichedMessage](enrichedStreamID)
	enriched.Key = event.Key
	enriched.Name = name
	encoded, err := flow.ConvertFrom(rt, enriched)
	if err != nil {
		return err
	}
	out.AddMessage(encoded)
	return nil
}

func referenceName(rt flow.Runtime, msg flow.ExtendedMessage) (string, bool, error) {
	reference, err := flow.OpenJoinedExternalState(rt, referenceStateName, msg)
	if errors.Is(err, flow.ErrStateNotRead) {
		return "", false, nil
	}
	if err != nil {
		return "", false, err
	}

	var state referenceState
	exists, err := reference.ConvertTo(&state)
	if err != nil || !exists || state.Name == nil {
		return "", false, err
	}
	return *state.Name, true, nil
}

Source code of the example

State in timers

The API for working with state in a timer handler is identical — a timer is passed instead of a message:

state, err := flow.OpenExternalState(rt, "/join-state", timer)

Here is an example from WaitClickJoin — the window is closed by a timer, and the state is cleared right after the result is published:

func (*joinFunction) OnTimer(
	ctx context.Context,
	rt flow.Runtime,
	timer flow.Timer,
	out flow.OutputCollector,
) error {
	if timer.StreamID != timerStream {
		return fmt.Errorf("unhandled timer stream %q", timer.StreamID)
	}

	state, err := flow.OpenExternalState(rt, joinStateName, timer)
	if err != nil {
		return err
	}
	var window joinState
	if _, err := state.ConvertTo(&window); err != nil {
		return err
	}

	if window.ShowTime != nil && *window.ShowTime != 0 && window.HitPayload != nil {
		var key joinKey
		if err := timer.Key.ConvertTo(&key); err != nil {
			return err
		}
		joined, err := joinedAction(rt, key, window)
		if err != nil {
			return err
		}
		out.AddMessage(joined)
	}

	return state.Clear()
}

A state cleared in this request is read as missing from then on: the computation sees the state as it will be after the response to the worker.

Binding the state to the key

In TTransformCompanionComputation, the state is bound to the message key, which is defined through group_by_schema in the computation spec. All the messages with the same key share one state. The key the accessor is bound to is taken from the input — that is why one handler can’t accidentally write the state of a foreign key.

In TTransformOrderedSourceCompanionComputation, the group_by_schema field isn’t supported. The key of the internal state is the source partition key, so all the messages of one partition share the state. For more on choosing the class for a SourceComputation, see the Computation (Go) section.

For more on configuring keys, see Stateful processing.

The key schema is available in the handler through rt.KeySchema(), and the key itself through the Key field of the input (msg.Key, timer.Key, visit.Key).

Configuring states in the spec

A state that isn’t declared in the spec can’t be opened: the opening function returns flow.ErrUnknownState with the list of the declared names. States are declared in three different places of the computation description:

• Internal states (YSON, Raw, Proto) — as a list of names in parameters.internal_states. The names are arbitrary, without a leading /.
• External states the computation owns — in the external_state_managers section. The section key sets the state name (an absolute path), and the external_state_manager_class_name field sets the registered manager class; for the typical scenario it is "NYT::NFlow::TSimpleExternalStateManager".
• External states the computation only reads — in the external_state_joiners section, next to external_state_managers.

"mapper" = {
    "computation_class_name" = "NYT::NFlow::NCompanion::TTransformCompanionComputation";
    "external_state_managers" = {
        "/shuffle-state" = {
            "external_state_manager_class_name" = "NYT::NFlow::TSimpleExternalStateManager";
            "parameters" = {
                "path" = "<cluster=cluster_name>//path/to/state";
            };
        };
    };
    "external_state_joiners" = {
        "/reference" = {
            "external_state_joiner_class_name" = "NYT::NFlow::TSimpleExternalStateJoiner";
            "parameters" = {
                "path" = "//path/to/current";
            };
        };
    };
    "parameters" = {
        "internal_states" = ["word-state"];
    };
};

The tables of internal states are created and maintained by Flow, while the external state table is created by you. For more details, see the Internal State (Go) and External State (Go) sections.

See also

• State Accessor (Go)
• Internal State (Go)
• External State (Go)
• Computation (Go)
• Stateful processing