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

---
 main.go | 75 +++++++++++++++++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 62 insertions(+), 13 deletions(-)

(limited to 'main.go')

diff --git a/main.go b/main.go
index 806c1ed..1d86d66 100644
--- a/main.go
+++ b/main.go
@@ -2,46 +2,95 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
-	"strconv"
 
 	"goflink/core"
+	"goflink/core/operator"
 	"goflink/runtime"
+	"goflink/state"
 	"goflink/transport"
 )
 
+// wordCounter counts how many times each key has been seen, in keyed state.
+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 }
+
 func main() {
+
+	/*
+			Bây giờ bạn đã có đầy đủ đồ chơi: MemoryStateBackend, KeyByOperator, StatefulMapOperator.
+
+		Trong file main.go, bạn hãy thiết lập một Pipeline có dạng: Source ➡️ KeyBy (Trích xuất Tên) ➡️ StatefulMap (Đếm số lần Tên xuất hiện) ➡️ Sink
+
+		Yêu cầu Logic của hàm StatefulMap:
+
+		Trong hàm Open, bạn hãy khởi tạo một ValueState[int] tên là "word-count".
+		Trong hàm xử lý (User Func), bạn hãy đọc giá trị cũ lên, cộng thêm 1, lưu lại vào State, và trả về con số mới đó.
+		Hãy chạy thử xem nếu bạn truyền vào ["Sang", "Tran", "Sang", "Sang"], nó có in ra 1, 1, 2, 3 không nhé. Nếu thành công, bạn đã tự tay chế tạo xong một Streaming Engine Stateful đích thực!
+	*/
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	source := make(chan core.StreamRecord[int])
-	mid := make(chan core.StreamRecord[int])
-	sink := make(chan core.StreamRecord[string])
+	// ponytail: one backend for the whole job. Each stateful task needs its own
+	// once tasks run on separate key-group shards.
+	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])
 
-	// x*10, then render as text
-	mapped := runtime.RunTask(ctx,
+	// keyBy(name) -> count occurrences per name
+	keyed := runtime.RunTask(ctx,
 		transport.NewLocalReceiveChannel(source),
 		transport.NewLocalEmitChannel(mid),
-		core.NewMapOperator(func(in int) (int, error) { return in * 10, nil }),
+		operator.NewKeyByOperator(func(in string) string { return in }),
 	)
-	printed := runtime.RunTask(ctx,
+	counted := runtime.RunTask(ctx,
 		transport.NewLocalReceiveChannel(mid),
 		transport.NewLocalEmitChannel(sink),
-		core.NewMapOperator(func(in int) (string, error) { return "v=" + strconv.Itoa(in), nil }),
+		operator.NewStatefulMapOperator[string, int](&wordCounter{}),
 	)
 
 	go func() {
 		defer close(source)
-		for i := 1; i <= 5; i++ {
-			source <- core.NewDataRecord("k", i, 0)
+		for _, name := range []string{"Sang", "Tran", "Sang", "Sang"} {
+			source <- core.NewDataRecord("", name, 0)
 		}
 	}()
 
 	for record := range sink {
-		fmt.Println(record.Value)
+		fmt.Println(record.Key, record.Value)
 	}
 
-	for _, errCh := range []<-chan error{mapped, printed} {
+	for _, errCh := range []<-chan error{keyed, counted} {
 		if err := <-errCh; err != nil {
 			fmt.Println("task failed:", err)
 		}
-- 
cgit v1.2.3