From 4ca27aec303f54dc9ee670c67c2fefa80dd004a1 Mon Sep 17 00:00:00 2001
From: SangTran-127 <tranquangsang12.7@gmail.com>
Date: Sun, 9 Aug 2026 18:32:44 +0700
Subject: feat: add KeyByOperator and wordCounter for stateful processing

---
 runtime/task_test.go | 116 ++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 111 insertions(+), 5 deletions(-)

(limited to 'runtime/task_test.go')

diff --git a/runtime/task_test.go b/runtime/task_test.go
index f27a5eb..ae72c14 100644
--- a/runtime/task_test.go
+++ b/runtime/task_test.go
@@ -6,6 +6,8 @@ import (
 	"testing"
 
 	"goflink/core"
+	"goflink/core/operator"
+	"goflink/state"
 	"goflink/transport"
 )
 
@@ -21,12 +23,12 @@ func TestPipeline(t *testing.T) {
 	mapped := RunTask(ctx,
 		transport.NewLocalReceiveChannel(source),
 		transport.NewLocalEmitChannel(mid),
-		core.NewMapOperator(func(in int) (int, error) { return in * 10, nil }),
+		operator.NewMapOperator(func(in int) (int, error) { return in * 10, nil }),
 	)
 	filtered := RunTask(ctx,
 		transport.NewLocalReceiveChannel(mid),
 		transport.NewLocalEmitChannel(sink),
-		core.NewFilterOperator(func(in int) bool { return in > 20 }),
+		operator.NewFilterOperator(func(in int) bool { return in > 20 }),
 	)
 
 	go func() {
@@ -58,6 +60,110 @@ func TestPipeline(t *testing.T) {
 	}
 }
 
+// wordCounter is the canonical stateful user func: read the old count for the
+// current key, +1, store it back.
+type wordCounter struct {
+	count core.ValueState[any]
+}
+
+func (w *wordCounter) Open(ctx context.Context) error {
+	be, ok := core.ExtractStateBackend(ctx)
+	if !ok {
+		return errors.New("no state backend in context")
+	}
+	s, err := be.GetValueState(ctx, "word-count")
+	if err != nil {
+		return err
+	}
+	w.count = s
+	return nil
+}
+
+func (w *wordCounter) Map(ctx context.Context, in string) (int, error) {
+	v, err := w.count.Value(ctx)
+	if err != nil {
+		return 0, err
+	}
+	n, _ := v.(int) // unseen key -> nil -> 0
+	n++
+	_, err = w.count.Update(ctx, n)
+	return n, err
+}
+
+func (w *wordCounter) Close(ctx context.Context) error { return nil }
+
+// keyBy -> statefulMap: state must be per key and survive across records.
+func TestStatefulMapCountsPerKey(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	ctx = core.InjectStateBackend(ctx, state.NewMemoryStateBackend(core.DefaultMaxParallelism))
+
+	source := make(chan core.StreamRecord[string])
+	mid := make(chan core.StreamRecord[string])
+	sink := make(chan core.StreamRecord[int])
+
+	keyed := RunTask(ctx,
+		transport.NewLocalReceiveChannel(source),
+		transport.NewLocalEmitChannel(mid),
+		operator.NewKeyByOperator(func(in string) string { return in }),
+	)
+	counted := RunTask(ctx,
+		transport.NewLocalReceiveChannel(mid),
+		transport.NewLocalEmitChannel(sink),
+		operator.NewStatefulMapOperator[string, int](&wordCounter{}),
+	)
+
+	go func() {
+		defer close(source)
+		for _, name := range []string{"Sang", "Tran", "Sang", "Sang"} {
+			source <- core.NewDataRecord("", name, 0)
+		}
+	}()
+
+	var got []int
+	for record := range sink {
+		got = append(got, record.Value)
+	}
+
+	want := []int{1, 1, 2, 3}
+	if len(got) != len(want) {
+		t.Fatalf("got %v, want %v", got, want)
+	}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Fatalf("got %v, want %v", got, want)
+		}
+	}
+
+	for _, errCh := range []<-chan error{keyed, counted} {
+		if err := <-errCh; err != nil {
+			t.Fatalf("task failed: %v", err)
+		}
+	}
+}
+
+// An unkeyed record must fail loudly instead of silently sharing one state slot.
+func TestStatefulMapRejectsUnkeyed(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	ctx = core.InjectStateBackend(ctx, state.NewMemoryStateBackend(core.DefaultMaxParallelism))
+
+	source := make(chan core.StreamRecord[string], 1)
+	sink := make(chan core.StreamRecord[int], 1)
+
+	done := RunTask(ctx,
+		transport.NewLocalReceiveChannel(source),
+		transport.NewLocalEmitChannel(sink),
+		operator.NewStatefulMapOperator[string, int](&wordCounter{}),
+	)
+
+	source <- core.NewDataRecord("", "Sang", 0)
+
+	if err := <-done; err == nil {
+		t.Fatal("unkeyed record was accepted, want an error")
+	}
+}
+
 // A cancelled ctx must unblock a task parked in Receive.
 func TestCancelStopsTask(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
@@ -68,7 +174,7 @@ func TestCancelStopsTask(t *testing.T) {
 	done := RunTask(ctx,
 		transport.NewLocalReceiveChannel(source),
 		transport.NewLocalEmitChannel(sink),
-		core.NewMapOperator(func(in int) (int, error) { return in, nil }),
+		operator.NewMapOperator(func(in int) (int, error) { return in, nil }),
 	)
 
 	cancel()
@@ -90,7 +196,7 @@ func TestProcessErrorPropagates(t *testing.T) {
 	done := RunTask(ctx,
 		transport.NewLocalReceiveChannel(source),
 		transport.NewLocalEmitChannel(sink),
-		core.NewMapOperator(func(in int) (int, error) { return 0, boom }),
+		operator.NewMapOperator(func(in int) (int, error) { return 0, boom }),
 	)
 
 	source <- core.NewDataRecord("k", 1, 0)
@@ -107,7 +213,7 @@ type closeProbe struct {
 
 func (o *closeProbe) Open(ctx context.Context) error { return nil }
 
-func (o *closeProbe) Process(ctx context.Context, r core.StreamRecord[int], out core.Emitter[int]) error {
+func (o *closeProbe) Process(ctx context.Context, r core.StreamRecord[int], out operator.Emitter[int]) error {
 	return out.Emit(ctx, r)
 }
 
-- 
cgit v1.2.3