Internal State in YTsaurus Flow (Java)

Internal State is a mechanism for working with internal state stored in Flow’s internal tables. Unlike External State, you don’t need to create the tables yourself — Flow manages them automatically.

Learn more about StateAccessor and working with state: State Accessor.

General information about stateful processing is in the Stateful processing section.

Overview

The Java SDK Flow (Java and Kotlin) provides several types of state accessors for working with Internal State; they differ by serialization format:

Accessor Format Description
YsonStateAccessor YSON Serialization via @YTreeObject annotations
ProtoStateAccessor Protobuf Serialization via Protobuf
DefaultStateAccessor Custom User-defined serializer/deserializer
RawStateAccessor byte[] No serialization (raw bytes)
NoOpStateAccessor — Stores only the fact that the state exists

All accessors implement the common StateAccessor<T> interface.

StateAccessor interface

public interface StateAccessor<T> {
    /** Get the state value. */
    @Nullable
    T get();

    /** Get the state value or a default value. */
    default T getOrDefault(T defaultValue);

    /** Set the state value. */
    void set(T value);

    /** Clear/delete the state for the key. */
    void clear();

    /** Get the state class. */
    Class<T> getStateClass();
}
interface StateAccessor<T> {
    /** Get the state value. */
    fun get(): T?

    /** Get the state value or a default value. */
    fun getOrDefault(defaultValue: T): T

    /** Set the state value. */
    fun set(value: T)

    /** Clear/delete the state for the key. */
    fun clear()

    /** Get the state class. */
    fun getStateClass(): Class<T>
}

Changing the value in place

The value returned by get() and getOrDefault() is live: it is decoded once per key and batch, every accessor for that key returns the same object, and the changes made to it are written to the state at the end of the batch without a set() call. Nothing is written when the value did not change. The default from getOrDefault() is attached to the key but not written: it can be changed right away, and it becomes the state value only once it is changed — an untouched default creates no row in the state table:

ctx.getState(COUNTER, message).getOrDefault(new CounterState()).count += 1;
ctx.getState(COUNTER, message).getOrDefault(CounterState()).count += 1

set() still replaces the whole value, and clear() removes the state. Protobuf messages are immutable, so their new value goes through set().

To detect the changes, a value that was read is re-encoded once at the end of the batch — even when the computation only inspects the state. readOnly() takes that work away: the accessor returns the same value but does not track it, so the state is neither re-encoded at the end of the batch nor sent to the worker. Use it wherever a state is only read. Its getOrDefault() does not create the state, and set() and clear() throw:

long count = ctx.getState(COUNTER, message).readOnly().getOrDefault().count;
val count = ctx.getState(COUNTER, message).readOnly().getOrDefault().count

The same intent can be declared once on the descriptor: InternalStateDescriptor.readOnly() returns a descriptor of the same state whose accessors are read-only, so no access site needs a readOnly() call of its own. Declare it once as a constant next to the original:

private static final InternalStateDescriptor<CounterState> COUNTER_READ_ONLY = COUNTER.readOnly();

long count = ctx.getState(COUNTER_READ_ONLY, message).getOrDefault().count;
private val COUNTER_READ_ONLY: InternalStateDescriptor<CounterState> = COUNTER.readOnly()

val count = ctx.getState(COUNTER_READ_ONLY, message).getOrDefault().count

An internal state lives within a single computation: the worker serves only the names listed in the internal_states of its parameters. A read-only descriptor therefore reads the state of the computation it is used in: it is not a way to read a state that another computation writes — that is what External State joiners are for.

Read-only is a discipline of access, not a property of the state: the tracking lives on the state itself. If the value for the same key was already read through a writable accessor in this batch, the state is already tracked, and a change made in place afterwards reaches the worker whichever accessor handed the value out. The descriptor guarantees "this accessor does not write", not "this state is not tracked".

YsonStateAccessor

YsonStateAccessor uses YSON serialization. The state class must be annotated with @YTreeObject.

Getting the accessor

// For a message
YsonStateAccessor<MyState> stateAccessor =
        ctx.getYsonStateAccessor("state-name", message, MyState.class);

// For a timer
YsonStateAccessor<MyState> stateAccessor =
        ctx.getYsonStateAccessor("state-name", timer, MyState.class);
// For a message
val stateAccessor: YsonStateAccessor<MyState> =
        ctx.getYsonStateAccessor("state-name", message, MyState::class.java)

// For a timer
val stateAccessor: YsonStateAccessor<MyState> =
        ctx.getYsonStateAccessor("state-name", timer, MyState::class.java)

Example state class

import ru.yandex.inside.yt.kosher.impl.ytree.object.annotation.YTreeObject;
import ru.yandex.inside.yt.kosher.impl.ytree.object.annotation.YTreeField;

@YTreeObject
public class CounterState {
    @YTreeField(key = "count")
    private long count;

    @YTreeField(key = "last_update")
    private long lastUpdate;

    // A default constructor is required
    public CounterState() {}

    // Getters and setters...
    public long getCount() { return count; }
    public void setCount(long count) { this.count = count; }
    public long getLastUpdate() { return lastUpdate; }
    public void setLastUpdate(long lastUpdate) { this.lastUpdate = lastUpdate; }
}
import ru.yandex.inside.yt.kosher.impl.ytree.object.annotation.YTreeObject
import ru.yandex.inside.yt.kosher.impl.ytree.object.annotation.YTreeField

@YTreeObject
class CounterState {
    @YTreeField(key = "count")
    var count: Long = 0

    @YTreeField(key = "last_update")
    var lastUpdate: Long = 0

    // A default constructor is required
    constructor()
}

Example usage

public class CounterFunction implements RowFunction {
    @Override
    public void onMessage(ExtendedMessage message, OutputCollector output, RuntimeContext ctx) {
        YsonStateAccessor<CounterState> stateAccessor =
                ctx.getYsonStateAccessor("counter", message, CounterState.class);

        // Get the current state or create a new one
        CounterState state = stateAccessor.getOrDefault(new CounterState());

        // Modify the state
        state.setCount(state.getCount() + 1);
        state.setLastUpdate(message.getEventTimestamp());
    }
}
class CounterFunction : RowFunction {
    override fun onMessage(message: ExtendedMessage, output: OutputCollector, ctx: RuntimeContext) {
        val stateAccessor: YsonStateAccessor<CounterState> =
                ctx.getYsonStateAccessor("counter", message, CounterState::class.java)

        // Get the current state or create a new one
        val state: CounterState = stateAccessor.getOrDefault(CounterState())

        // Modify the state
        state.count = state.count + 1
        state.lastUpdate = message.getEventTimestamp()
    }
}

ProtoStateAccessor

Source code

ProtoStateAccessor uses Protobuf serialization. Your state class must inherit from com.google.protobuf.MessageLite.

Get the accessor

// For a message
ProtoStateAccessor<MyProtoState> stateAccessor =
        ctx.getProtoStateAccessor("state-name", message, MyProtoState.class);

// For a timer
ProtoStateAccessor<MyProtoState> stateAccessor =
        ctx.getProtoStateAccessor("state-name", timer, MyProtoState.class);
// For a message
val stateAccessor: ProtoStateAccessor<MyProtoState> =
        ctx.getProtoStateAccessor("state-name", message, MyProtoState::class.java)

// For a timer
val stateAccessor: ProtoStateAccessor<MyProtoState> =
        ctx.getProtoStateAccessor("state-name", timer, MyProtoState::class.java)

The getOrDefault method

ProtoStateAccessor provides a parameterless getOrDefault() method that returns a default Protobuf object:

ProtoStateAccessor<MyProtoState> stateAccessor =
        ctx.getProtoStateAccessor("state-name", message, MyProtoState.class);

// Get the state or a default Protobuf object
MyProtoState state = stateAccessor.getOrDefault();
val stateAccessor: ProtoStateAccessor<MyProtoState> =
        ctx.getProtoStateAccessor("state-name", message, MyProtoState::class.java)

// Get the state or a default Protobuf object
val state: MyProtoState = stateAccessor.getOrDefault()

Usage example

public class ProtoCounterFunction implements RowFunction {
    @Override
    public void onMessage(ExtendedMessage message, OutputCollector output, RuntimeContext ctx) {
        ProtoStateAccessor<CounterProto> stateAccessor =
                ctx.getProtoStateAccessor("counter", message, CounterProto.class);

        CounterProto state = stateAccessor.getOrDefault();

        // Modify using the Protobuf builder
        CounterProto updatedState = state.toBuilder()
                .setCount(state.getCount() + 1)
                .setLastUpdate(message.getEventTimestamp())
                .build();

        stateAccessor.set(updatedState);
    }
}
class ProtoCounterFunction : RowFunction {
    override fun onMessage(message: ExtendedMessage, output: OutputCollector, ctx: RuntimeContext) {
        val stateAccessor: ProtoStateAccessor<CounterProto> =
                ctx.getProtoStateAccessor("counter", message, CounterProto::class.java)

        val state: CounterProto = stateAccessor.getOrDefault()

        // Modify using the Protobuf builder
        val updatedState: CounterProto = state.toBuilder()
                .setCount(state.getCount() + 1)
                .setLastUpdate(message.getEventTimestamp())
                .build()

        stateAccessor.set(updatedState)
    }
}

DefaultStateAccessor

Source code

DefaultStateAccessor lets you use custom serialization and deserialization functions.

Get the accessor

// For a message
DefaultStateAccessor<MyState> stateAccessor = ctx.getStateAccessor(
        "state-name",
        message,
        MyState.class,
        state -> serialize(state),      // Function<MyState, byte[]>
        bytes -> deserialize(bytes)     // Function<byte[], MyState>
);

// For a timer
DefaultStateAccessor<MyState> stateAccessor = ctx.getStateAccessor(
        "state-name",
        timer,
        MyState.class,
        state -> serialize(state),
        bytes -> deserialize(bytes)
);
// For a message
val stateAccessor: DefaultStateAccessor<MyState> = ctx.getStateAccessor(
        "state-name",
        message,
        MyState::class.java,
        { state -> serialize(state) },      // (MyState) -> ByteArray
        { bytes -> deserialize(bytes) }     // (ByteArray) -> MyState
)

// For a timer
val stateAccessor: DefaultStateAccessor<MyState> = ctx.getStateAccessor(
        "state-name",
        timer,
        MyState::class.java,
        { state -> serialize(state) },
        { bytes -> deserialize(bytes) }
)

Example with Jackson

public class JsonCounterFunction implements RowFunction {
    private static final ObjectMapper mapper = new ObjectMapper();

    @Override
    public void onMessage(ExtendedMessage message, OutputCollector output, RuntimeContext ctx) {
        DefaultStateAccessor<CounterState> stateAccessor = ctx.getStateAccessor(
                "counter",
                message,
                CounterState.class,
                state -> {
                    try { return mapper.writeValueAsBytes(state); }
                    catch (Exception e) { throw new RuntimeException(e); }
                },
                bytes -> {
                    try { return mapper.readValue(bytes, CounterState.class); }
                    catch (Exception e) { throw new RuntimeException(e); }
                }
        );

        CounterState state = stateAccessor.getOrDefault(new CounterState());
        state.setCount(state.getCount() + 1);
    }
}
class JsonCounterFunction : RowFunction {
    companion object {
        private val mapper = ObjectMapper()
    }

    override fun onMessage(message: ExtendedMessage, output: OutputCollector, ctx: RuntimeContext) {
        val stateAccessor: DefaultStateAccessor<CounterState> = ctx.getStateAccessor(
                "counter",
                message,
                CounterState::class.java,
                { state ->
                    try { mapper.writeValueAsBytes(state) }
                    catch (e: Exception) { throw RuntimeException(e) }
                },
                { bytes ->
                    try { mapper.readValue(bytes, CounterState::class.java) }
                    catch (e: Exception) { throw RuntimeException(e) }
                }
        )

        val state: CounterState = stateAccessor.getOrDefault(CounterState())
        state.count = state.count + 1
    }
}

RawStateAccessor

Source code

Use RawStateAccessor to work with raw bytes without serialization or deserialization.

Get the accessor

RawStateAccessor stateAccessor = ctx.getRawStateAccessor("state-name", message);
val stateAccessor: RawStateAccessor = ctx.getRawStateAccessor("state-name", message)

Example usage

RawStateAccessor stateAccessor = ctx.getRawStateAccessor("raw-state", message);

byte[] data = stateAccessor.get();
if (data != null) {
    // Process raw data...
}

// Write raw data
stateAccessor.set(new byte[]{0x01, 0x02, 0x03});

// Clear the data
stateAccessor.clear();
val stateAccessor: RawStateAccessor = ctx.getRawStateAccessor("raw-state", message)

val data: ByteArray? = stateAccessor.get()
if (data != null) {
    // Process raw data...
}

// Write raw data
stateAccessor.set(byteArrayOf(0x01, 0x02, 0x03))

// Clear the data
stateAccessor.clear()

NoOpStateAccessor

Source code

NoOpStateAccessor stores only the fact that a state exists for a key, with no payload. Use it to track already processed keys (for deduplication).

Get the accessor

NoOpStateAccessor stateAccessor = ctx.getNoOpStateAccessor("seen-keys", message);
val stateAccessor: NoOpStateAccessor = ctx.getNoOpStateAccessor("seen-keys", message)

Example usage

public class DeduplicationFunction implements RowFunction {
    @Override
    public void onMessage(ExtendedMessage message, OutputCollector output, RuntimeContext ctx) {
        NoOpStateAccessor stateAccessor = ctx.getNoOpStateAccessor("seen-keys", message);

        // Check if the key was already processed
        if (stateAccessor.get() != null) {
            // The key is already processed, skip it
            return;
        }

        // Mark the key as processed
        stateAccessor.touch();

        // Process the message...
        output.addMessage(new Message("output", message.getPayload()));
    }
}
class DeduplicationFunction : RowFunction {
    override fun onMessage(message: ExtendedMessage, output: OutputCollector, ctx: RuntimeContext) {
        val stateAccessor: NoOpStateAccessor = ctx.getNoOpStateAccessor("seen-keys", message)

        // Check if the key was already processed
        if (stateAccessor.get() != null) {
            // The key is already processed, skip it
            return
        }

        // Mark the key as processed
        stateAccessor.touch()

        // Process the message...
        output.addMessage(Message("output", message.getPayload()))
    }
}

Configuration in the static spec

Internal State doesn’t require you to create external tables. States are automatically stored in Flow’s internal tables (states and partition_states).

You must declare the names of internal states in the internal_states section of the computation parameters in the static spec:

"computations" = {
    "counter" = {
        "computation_class_name" = "NYT::NFlow::NCompanion::TTransformCompanionComputation";
        "group_by_schema" = [
            {"name" = "hash"; "expression" = "farm_hash(key)"; "type" = "uint64"};
            {"name" = "key"; "type" = "string"};
        ];
        "input_stream_ids" = ["input"];
        "output_stream_ids" = ["output"];
        "parameters" = {
            "internal_states" = ["counter"];
        };
    };
};

The state name in your code (the first argument of ctx.getYsonStateAccessor(...), ctx.getProtoStateAccessor(...), and so on) must match the name declared in internal_states.