1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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)
}
|