Shuffle in YTsaurus Flow (Java)

The pipeline reads the stream of events, groups them by key, and counts the number of unique events using an external state (ExternalStateAccessor). This example shows how to configure a companion with Spring Boot.

Source code (Java)

Source code (Kotlin)

Companion components

EventMapper

This is a process function for the source-computation reader. It parses and transforms the input data. The version below is simplified. In the actual code, you also parse the data field as JSON using Jackson’s ObjectMapper:

    @SuppressWarnings("unchecked")
    @Override
    public void onMessage(ExtendedMessage message, OutputCollector output, RuntimeContext ctx) {
        var messageBuilder = ctx.createMessageBuilder("event");
        String data = message.get("data", String.class);
        try {
            Map<String, Object> map = ysonMapper.readValue(data, Map.class);
            messageBuilder.set("value", map.get("value"));
            messageBuilder.set("key_a", map.get("key_a"));
            messageBuilder.set("key_b", map.get("key_b"));
            messageBuilder.set("key_c", map.get("key_c"));
            messageBuilder.set("key_d", map.get("key_d"));
            output.addMessage(messageBuilder.finish());
        } catch (JsonProcessingException e) {
            log.error("Error parsing json from data filed", e);
            throw new RuntimeException(e);
        }
    }

    @Suppress("UNCHECKED_CAST")
    override fun onMessage(message: ExtendedMessage, output: OutputCollector, ctx: RuntimeContext) {
        val messageBuilder = ctx.createMessageBuilder("event")
        val data = message.get("data", String::class.java)
        try {
            val map = ysonMapper.readValue(data, Map::class.java) as Map<String, Any?>
            messageBuilder.set("value", map["value"])
            messageBuilder.set("key_a", map["key_a"])
            messageBuilder.set("key_b", map["key_b"])
            messageBuilder.set("key_c", map["key_c"])
            messageBuilder.set("key_d", map["key_d"])
            output.addMessage(messageBuilder.finish())
        } catch (e: JsonProcessingException) {
            log.error("Error parsing json from data filed", e)
            throw RuntimeException(e)
        }
    }

EventReducer

This is a process function that uses ExternalStateAccessor to count the number of events:

    @Override
    public void onMessage(ExtendedMessage message, OutputCollector output, RuntimeContext ctx) {
        ExternalStateAccessor externalStateAccessor = ctx.getState(SHUFFLE_STATE, message);
        Payload state = externalStateAccessor.getOrDefault();
        PayloadBuilder stateBuilder = state.toBuilder();
        if (state.get("count", Long.class) == null) {
            stateBuilder.set("count", 1L);
        } else {
            stateBuilder.set("count", state.get("count", Long.class) + 1);
        }
        externalStateAccessor.set(stateBuilder.finish());
    }

    override fun onMessage(message: ExtendedMessage, output: OutputCollector, ctx: RuntimeContext) {
        val externalStateAccessor = ctx.getState(SHUFFLE_STATE, message)
        val state = externalStateAccessor.getOrDefault()
        val stateBuilder = state.toBuilder()
        val count = state.get("count", Long::class.java) ?: 0L
        stateBuilder.set("count", count + 1)
        externalStateAccessor.set(stateBuilder.finish())
    }

How it works:

  • You get an ExternalStateAccessor for the state "shuffle-state", which is bound to the key of the current message.
  • You extract the current state value. If there is no state, getOrDefault() returns an empty Payload.
  • You create a PayloadBuilder from the current state and increment the counter.
  • You save the updated state.

Registering computations

You register the source computation reader with the @FlowSourceComputation annotation, and the transformation reducer with the @FlowComputation annotation:

@FlowSourceComputation(id = "reader")

@FlowComputation(id = "reducer")

@FlowSourceComputation(id = "reader")

@FlowComputation(id = "reducer")

PipelineMain

This single entry point starts the pipeline or serves as its companion, depending on YT_FLOW_MODE.

    public static void main(String[] args) throws Exception {
        new SpringApplicationBuilder(PipelineMain.class).run(args);
    }

@SpringBootApplication
open class PipelineMain {
    companion object {
        @JvmStatic
        fun main(args: Array<String>) {
            SpringApplicationBuilder(PipelineMain::class.java)
                .run(*args)
        }
    }
}

Key patterns

  • Configuration via Spring Boot: you register computations with the @FlowSourceComputation / @FlowComputation annotations; flow-spring-boot-starter manages the gRPC server lifecycle.
  • ExternalStateAccessor: you work with the external state using Payload and PayloadBuilder.
  • SourceComputation with ProcessFunction: the reader uses EventMapper to transform input data on the companion side.

See also

Previous
Word Count
Next
Async Request