7 files changed, 144 insertions, 19 deletions
diff --git a/core/operator.go b/core/operator.go
deleted file mode 100644
index 7785d44..0000000
--- a/core/operator.go
+++ /dev/null
@@ -1,13 +0,0 @@
-package core
-
-import "context"
-
-type Emitter[OUT any] interface {
-	Emit(ctx context.Context, record StreamRecord[OUT]) error
-}
-
-type Operator[IN, OUT any] interface {
-	Open(ctx context.Context) error
-	Process(ctx context.Context, record StreamRecord[IN], out Emitter[OUT]) error
-	Close(ctx context.Context) error
-}
diff --git a/core/operator/keyby.go b/core/operator/keyby.go
new file mode 100644
index 0000000..1296df0
--- /dev/null
+++ b/core/operator/keyby.go
@@ -0,0 +1,32 @@
+package operator
+
+import (
+	"context"
+	"goflink/core"
+)
+
+type KeySelector[IN any] func(in IN) string
+type KeyByOperator[IN any] struct {
+	selector KeySelector[IN]
+}
+
+func NewKeyByOperator[IN any](selector KeySelector[IN]) *KeyByOperator[IN] {
+	return &KeyByOperator[IN]{selector: selector}
+}
+
+func (k *KeyByOperator[IN]) Open(ctx context.Context) error {
+	return nil
+}
+
+func (k *KeyByOperator[IN]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[IN]) error {
+	if record.Type != core.RecordData {
+		return out.Emit(ctx, record)
+	}
+
+	record.Key = k.selector(record.Value)
+	return out.Emit(ctx, record)
+}
+
+func (k *KeyByOperator[IN]) Close(ctx context.Context) error {
+	return nil
+}
diff --git a/core/map.go b/core/operator/map.go
index bbb8dcd..ef6912e 100644
--- a/core/map.go
+++ b/core/operator/map.go
@@ -1,8 +1,9 @@
-package core
+package operator
 
 import (
 	"context"
 	"fmt"
+	"goflink/core"
 )
 
 // MapFunc is business logic for user implement
@@ -22,10 +23,10 @@ func (fo *MapOperator[IN, OUT]) Open(ctx context.Context) error {
 	return nil
 }
 
-func (fo *MapOperator[IN, OUT]) Process(ctx context.Context, record StreamRecord[IN], out Emitter[OUT]) error {
-	if record.Type != RecordData {
+func (fo *MapOperator[IN, OUT]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error {
+	if record.Type != core.RecordData {
 		//	 watermark and barrier
-		return out.Emit(ctx, StreamRecord[OUT]{
+		return out.Emit(ctx, core.StreamRecord[OUT]{
 			Type:      record.Type,
 			Key:       record.Key,
 			Timestamp: record.Timestamp,
@@ -39,8 +40,8 @@ func (fo *MapOperator[IN, OUT]) Process(ctx context.Context, record StreamRecord
 		return fmt.Errorf("cannot process value %v: %w", record.Value, err)
 	}
 
-	return out.Emit(ctx, StreamRecord[OUT]{
-		Type:      RecordData,
+	return out.Emit(ctx, core.StreamRecord[OUT]{
+		Type:      core.RecordData,
 		Key:       record.Key,
 		Timestamp: record.Timestamp,
 		Value:     outValue,
diff --git a/core/operator/operator.go b/core/operator/operator.go
new file mode 100644
index 0000000..d4bf21b
--- /dev/null
+++ b/core/operator/operator.go
@@ -0,0 +1,16 @@
+package operator
+
+import (
+	"context"
+	"goflink/core"
+)
+
+type Emitter[OUT any] interface {
+	Emit(ctx context.Context, record core.StreamRecord[OUT]) error
+}
+
+type Operator[IN, OUT any] interface {
+	Open(ctx context.Context) error
+	Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error
+	Close(ctx context.Context) error
+}
diff --git a/core/operator/stateful_map.go b/core/operator/stateful_map.go
new file mode 100644
index 0000000..7263170
--- /dev/null
+++ b/core/operator/stateful_map.go
@@ -0,0 +1,75 @@
+package operator
+
+import (
+	"context"
+	"fmt"
+
+	"goflink/core"
+)
+
+// RichMapFunc is a stateful map on the user side: Open grabs the state handles
+// off ctx, Map runs per record with the backend already pointed at that
+// record's key.
+type RichMapFunc[IN, OUT any] interface {
+	Open(ctx context.Context) error
+	Map(ctx context.Context, in IN) (OUT, error)
+	Close(ctx context.Context) error
+}
+
+type StatefulMapOperator[IN, OUT any] struct {
+	userFunc RichMapFunc[IN, OUT]
+	backend  core.StateBackend
+}
+
+func NewStatefulMapOperator[IN, OUT any](fn RichMapFunc[IN, OUT]) *StatefulMapOperator[IN, OUT] {
+	return &StatefulMapOperator[IN, OUT]{userFunc: fn}
+}
+
+func (s *StatefulMapOperator[IN, OUT]) Open(ctx context.Context) error {
+	// Get backend state from context via using context KV
+	be, ok := core.ExtractStateBackend(ctx)
+	if !ok || be == nil {
+		return fmt.Errorf("stateful map: no state backend in context")
+	}
+
+	s.backend = be
+	// The user func opens its ValueState off the same ctx.
+	return s.userFunc.Open(ctx)
+}
+
+func (s *StatefulMapOperator[IN, OUT]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error {
+	if record.Type != core.RecordData {
+		return out.Emit(ctx, core.StreamRecord[OUT]{
+			Type:      record.Type,
+			Timestamp: record.Timestamp,
+			Key:       record.Key,
+			BarrierID: record.BarrierID,
+		})
+	}
+
+	// Without a key every record would silently share the "" state slot.
+	if record.Key == "" {
+		return fmt.Errorf("stateful map: record %v has no key, keyBy it first", record.Value)
+	}
+
+	if err := s.backend.SetCurrentKey(record.Key); err != nil {
+		return fmt.Errorf("set current key %q: %w", record.Key, err)
+	}
+
+	// run the user transform
+	res, err := s.userFunc.Map(ctx, record.Value)
+	if err != nil {
+		return fmt.Errorf("cannot process value %v: %w", record.Value, err)
+	}
+
+	return out.Emit(ctx, core.StreamRecord[OUT]{
+		Type:      core.RecordData,
+		Timestamp: record.Timestamp,
+		Key:       record.Key,
+		Value:     res,
+	})
+}
+
+func (s *StatefulMapOperator[IN, OUT]) Close(ctx context.Context) error {
+	return s.userFunc.Close(ctx)
+}
diff --git a/core/rich_map.go b/core/rich_map.go
new file mode 100644
index 0000000..9a8bc95
--- /dev/null
+++ b/core/rich_map.go
@@ -0,0 +1 @@
+package core
diff --git a/core/state_backend.go b/core/state_backend.go
index 219bb63..f8b83fb 100644
--- a/core/state_backend.go
+++ b/core/state_backend.go
@@ -11,3 +11,16 @@ type StateBackend interface {
 	// GetValueState will perform in disk
 	GetValueState(ctx context.Context, key string) (ValueState[any], error)
 }
+
+type stateBackendKeyType struct{}
+
+var stateBackendKey stateBackendKeyType
+
+func InjectStateBackend(ctx context.Context, state StateBackend) context.Context {
+	return context.WithValue(ctx, stateBackendKey, state)
+}
+
+func ExtractStateBackend(ctx context.Context) (StateBackend, bool) {
+	backend, ok := ctx.Value(stateBackendKey).(StateBackend)
+	return backend, ok
+}