diff --git a/CHANGES.md b/CHANGES.md index f0d5d06b9d9f..8ea41fe6e13e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -79,6 +79,7 @@ * [Flink Runner] Added opt-in static round-robin split assignment for small bounded sources via the new `sourceStaticSplitThresholdMb` pipeline option. The default of 0 keeps the existing lazy pull-based assignment ([#39873](https://github.com/apache/beam/issues/39873)). * Added automatic caching of bounded, single-pane side-input views for classic Java Flink DataStream execution ([#39866](https://github.com/apache/beam/issues/39866)). * (Python) Added `Sample.Any`, the Python equivalent of Java's `Sample.any`, which returns up to n arbitrary elements from a PCollection ([#18552](https://github.com/apache/beam/issues/18552)). +* (Go) Added support for custom WindowFns ([#20627](https://github.com/apache/beam/issues/20627)). ## Breaking Changes diff --git a/sdks/go/pkg/beam/core/graph/window/fn.go b/sdks/go/pkg/beam/core/graph/window/fn.go index df32a97b89c2..a212637919ba 100644 --- a/sdks/go/pkg/beam/core/graph/window/fn.go +++ b/sdks/go/pkg/beam/core/graph/window/fn.go @@ -17,6 +17,7 @@ package window import ( "fmt" + "reflect" "time" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" @@ -30,6 +31,7 @@ const ( FixedWindows Kind = "FIX" SlidingWindows Kind = "SLI" Sessions Kind = "SES" + CustomWindows Kind = "CUS" // User-defined custom WindowFn ) // NewGlobalWindows returns the default WindowFn, which places all elements @@ -53,6 +55,28 @@ func NewSessions(gap time.Duration) *Fn { return &Fn{Kind: Sessions, Gap: gap} } +// NewCustom returns a WindowFn backed by a user-defined custom window +// function. The fn value must be a pointer-to-struct whose concrete type +// has been registered with [RegisterWindowFn] during init. Custom window +// functions are non-merging and must return [IntervalWindow] values from +// their AssignWindows method. +// +// NewCustom panics if fn is nil or its type was not registered. +func NewCustom(fn any) *Fn { + if fn == nil { + panic("window.NewCustom: fn must not be nil") + } + t := reflect.TypeOf(fn) + st := t + if t.Kind() == reflect.Pointer { + st = t.Elem() + } + if _, ok := LookupWindowFn(st); !ok { + panic(fmt.Sprintf("window.NewCustom: type %v is not registered; call window.RegisterWindowFn during init()", t)) + } + return &Fn{Kind: CustomWindows, CustomFn: fn} +} + // Fn defines the window fn. type Fn struct { Kind Kind @@ -60,6 +84,22 @@ type Fn struct { Size time.Duration // FixedWindows, SlidingWindows Period time.Duration // SlidingWindows Gap time.Duration // Sessions + + CustomFn any // CustomWindows (nil for built-in kinds) +} + +// NeedsElement reports whether a CustomWindows Fn has an element-aware +// AssignWindows signature. Returns false for all built-in window kinds. +func (w *Fn) NeedsElement() bool { + if w.Kind != CustomWindows || w.CustomFn == nil { + return false + } + t := reflect.TypeOf(w.CustomFn) + if t.Kind() == reflect.Pointer { + t = t.Elem() + } + elems, _ := LookupWindowFn(t) + return len(elems) > 0 } // TODO(herohde) 4/17/2018: do we need to expose the window type as well? @@ -82,6 +122,8 @@ func (w *Fn) String() string { return fmt.Sprintf("%v[%v@%v]", w.Kind, w.Size, w.Period) case Sessions: return fmt.Sprintf("%v[%v]", w.Kind, w.Gap) + case CustomWindows: + return fmt.Sprintf("%v[%T %+v]", w.Kind, w.CustomFn, w.CustomFn) default: return string(w.Kind) } @@ -105,6 +147,8 @@ func (w *Fn) Equals(o *Fn) bool { return w.Period == o.Period && w.Size == o.Size case Sessions: return w.Gap == o.Gap + case CustomWindows: + return reflect.DeepEqual(w.CustomFn, o.CustomFn) default: panic(fmt.Sprintf("unknown window type: %v", w)) } diff --git a/sdks/go/pkg/beam/core/graph/window/fn_test.go b/sdks/go/pkg/beam/core/graph/window/fn_test.go index 6baa37c41ea0..55a3bf697258 100644 --- a/sdks/go/pkg/beam/core/graph/window/fn_test.go +++ b/sdks/go/pkg/beam/core/graph/window/fn_test.go @@ -16,10 +16,18 @@ package window import ( + "strings" "testing" "time" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" ) +func init() { + RegisterWindowFn[*testWindowFn]() +} + func TestEquals(t *testing.T) { tests := []struct { name string @@ -81,6 +89,24 @@ func TestEquals(t *testing.T) { NewSlidingWindows(10*time.Millisecond, 100*time.Millisecond), false, }, + { + "custom equal", + NewCustom(&testWindowFn{BucketSize: 3000}), + NewCustom(&testWindowFn{BucketSize: 3000}), + true, + }, + { + "custom inequal", + NewCustom(&testWindowFn{BucketSize: 3000}), + NewCustom(&testWindowFn{BucketSize: 5000}), + false, + }, + { + "custom vs fixed", + NewCustom(&testWindowFn{BucketSize: 3000}), + NewFixedWindows(3 * time.Second), + false, + }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { @@ -90,3 +116,52 @@ func TestEquals(t *testing.T) { }) } } + +// testWindowFn is a minimal custom WindowFn for testing. +type testWindowFn struct { + BucketSize int64 +} + +func (f *testWindowFn) AssignWindows(ts typex.EventTime) []typex.Window { + bucket := typex.EventTime(f.BucketSize) + // Euclidean remainder; correct floor for negative ts. + start := ts - ((ts%bucket)+bucket)%bucket + end := start + bucket + return []typex.Window{IntervalWindow{Start: start, End: end}} +} + +func TestNewCustom(t *testing.T) { + fn := NewCustom(&testWindowFn{BucketSize: 3000}) + if fn.Kind != CustomWindows { + t.Errorf("NewCustom().Kind = %v, want %v", fn.Kind, CustomWindows) + } + if fn.CustomFn == nil { + t.Fatal("NewCustom().CustomFn is nil") + } +} + +func TestNewCustomPanicsOnNil(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Error("NewCustom(nil) did not panic") + } + }() + NewCustom(nil) +} + +func TestCustomCoder(t *testing.T) { + fn := NewCustom(&testWindowFn{BucketSize: 3000}) + got := fn.Coder() + want := coder.NewIntervalWindow() + if got.Kind != want.Kind { + t.Errorf("Coder().Kind = %v, want %v", got.Kind, want.Kind) + } +} + +func TestCustomString(t *testing.T) { + fn := NewCustom(&testWindowFn{BucketSize: 3000}) + s := fn.String() + if !strings.HasPrefix(s, "CUS[") { + t.Errorf("String() = %q, want prefix CUS[", s) + } +} diff --git a/sdks/go/pkg/beam/core/graph/window/invoke.go b/sdks/go/pkg/beam/core/graph/window/invoke.go new file mode 100644 index 000000000000..715906488da9 --- /dev/null +++ b/sdks/go/pkg/beam/core/graph/window/invoke.go @@ -0,0 +1,128 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package window + +import ( + "reflect" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" + "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" +) + +// tsOnlyAssigner is the fast-path interface for timestamp-only custom WindowFns. +type tsOnlyAssigner interface { + AssignWindows(typex.EventTime) []typex.Window +} + +// anyElemAssigner is the fast-path interface for element-aware WindowFns whose +// element parameter is typed as any. +type anyElemAssigner interface { + AssignWindows(typex.EventTime, any) []typex.Window +} + +// anyKVAssigner is the fast-path interface for KV-aware WindowFns whose key and +// value parameters are both typed as any. +type anyKVAssigner interface { + AssignWindows(typex.EventTime, any, any) []typex.Window +} + +// WindowFnInvoker wraps a custom WindowFn instance and dispatches AssignWindows +// through a path chosen once at construction: a typed interface assertion where +// the signature permits one, otherwise reflect.Value.Call, which costs a small +// per-call allocation. +// +// Element parameters follow the convention the DoFn invoker uses for main +// input: a KV element arrives as a key and a value, anything else as a single +// element. +type WindowFnInvoker struct { + call func(ts typex.EventTime, elm, elm2 any) []typex.Window +} + +// NewWindowFnInvoker builds an invoker for fn. The concrete type of fn must +// have been previously registered via RegisterWindowFn. +func NewWindowFnInvoker(fn any) (*WindowFnInvoker, error) { + t := reflect.TypeOf(fn) + if t == nil { + return nil, errors.New("window.NewWindowFnInvoker: fn must not be nil") + } + structType := t + if t.Kind() == reflect.Pointer { + structType = t.Elem() + } + + elems, ok := LookupWindowFn(structType) + if !ok { + return nil, errors.Errorf("window.NewWindowFnInvoker: type %v is not registered; call window.RegisterWindowFn during init()", t) + } + + inv := &WindowFnInvoker{} + + switch len(elems) { + case 0: + if a, ok := fn.(tsOnlyAssigner); ok { + inv.call = func(ts typex.EventTime, _, _ any) []typex.Window { + return a.AssignWindows(ts) + } + return inv, nil + } + case 1: + if a, ok := fn.(anyElemAssigner); ok { + inv.call = func(ts typex.EventTime, elm, _ any) []typex.Window { + return a.AssignWindows(ts, elm) + } + return inv, nil + } + default: + if a, ok := fn.(anyKVAssigner); ok { + inv.call = func(ts typex.EventTime, elm, elm2 any) []typex.Window { + return a.AssignWindows(ts, elm, elm2) + } + return inv, nil + } + } + + // Concrete element types cannot be reached through an interface assertion. + m := reflect.ValueOf(fn).MethodByName("AssignWindows") + if !m.IsValid() { + return nil, errors.Errorf("window.NewWindowFnInvoker: %v has no AssignWindows method", t) + } + + switch len(elems) { + case 0: + inv.call = func(ts typex.EventTime, _, _ any) []typex.Window { + out := m.Call([]reflect.Value{reflect.ValueOf(ts)}) + return out[0].Interface().([]typex.Window) + } + case 1: + inv.call = func(ts typex.EventTime, elm, _ any) []typex.Window { + out := m.Call([]reflect.Value{reflect.ValueOf(ts), reflect.ValueOf(elm)}) + return out[0].Interface().([]typex.Window) + } + default: + inv.call = func(ts typex.EventTime, elm, elm2 any) []typex.Window { + out := m.Call([]reflect.Value{reflect.ValueOf(ts), reflect.ValueOf(elm), reflect.ValueOf(elm2)}) + return out[0].Interface().([]typex.Window) + } + } + return inv, nil +} + +// Invoke calls AssignWindows on the underlying WindowFn. Signatures that do not +// take element parameters ignore elm and elm2; elm2 carries the value of a KV +// element. +func (inv *WindowFnInvoker) Invoke(ts typex.EventTime, elm, elm2 any) []typex.Window { + return inv.call(ts, elm, elm2) +} diff --git a/sdks/go/pkg/beam/core/graph/window/invoke_test.go b/sdks/go/pkg/beam/core/graph/window/invoke_test.go new file mode 100644 index 000000000000..3f54187e69da --- /dev/null +++ b/sdks/go/pkg/beam/core/graph/window/invoke_test.go @@ -0,0 +1,175 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package window + +import ( + "testing" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" +) + +// elemAwareAnyWindowFn accepts an element typed as any, so the invoker can +// reach it through an interface assertion. +type elemAwareAnyWindowFn struct { + SizeMs int64 +} + +func (f *elemAwareAnyWindowFn) AssignWindows(ts typex.EventTime, _ any) []typex.Window { + size := typex.EventTime(f.SizeMs) + start := ts - ((ts%size)+size)%size + return []typex.Window{IntervalWindow{Start: start, End: start + size}} +} + +// elemAwareConcreteWindowFn accepts a concrete element type, so the invoker +// has to dispatch through reflect. +type elemAwareConcreteWindowFn struct { + DefaultSizeMs int64 +} + +func (f *elemAwareConcreteWindowFn) AssignWindows(ts typex.EventTime, elem int64) []typex.Window { + size := typex.EventTime(elem) + if size <= 0 { + size = typex.EventTime(f.DefaultSizeMs) + } + start := ts - ((ts%size)+size)%size + return []typex.Window{IntervalWindow{Start: start, End: start + size}} +} + +// kvConcreteWindowFn takes a KV element as concrete key and value types. +type kvConcreteWindowFn struct{} + +func (f *kvConcreteWindowFn) AssignWindows(ts typex.EventTime, k string, v int64) []typex.Window { + size := typex.EventTime(v) + start := ts - ((ts%size)+size)%size + return []typex.Window{IntervalWindow{Start: start, End: start + size}} +} + +// kvAnyWindowFn takes a KV element as any, reaching the interface fast path. +type kvAnyWindowFn struct{} + +func (f *kvAnyWindowFn) AssignWindows(ts typex.EventTime, k, v any) []typex.Window { + size := typex.EventTime(v.(int64)) + start := ts - ((ts%size)+size)%size + return []typex.Window{IntervalWindow{Start: start, End: start + size}} +} + +func init() { + RegisterWindowFn[*elemAwareAnyWindowFn]() + RegisterWindowFn[*elemAwareConcreteWindowFn]() + RegisterWindowFn[*kvConcreteWindowFn]() + RegisterWindowFn[*kvAnyWindowFn]() +} + +// TestWindowFnInvoker_KV checks that a KV element reaches AssignWindows as a +// separate key and value, the way a DoFn receives its main input. +func TestWindowFnInvoker_KV(t *testing.T) { + tests := []struct { + name string + fn any + }{ + {"concrete key and value", &kvConcreteWindowFn{}}, + {"key and value as any", &kvAnyWindowFn{}}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + inv := mustInvoker(t, tc.fn) + + windows := inv.Invoke(7500, "key", int64(5000)) + if len(windows) != 1 { + t.Fatalf("got %d windows, want 1", len(windows)) + } + want := IntervalWindow{Start: 5000, End: 10000} + if !windows[0].Equals(want) { + t.Errorf("Invoke(7500, key, 5000) = %v, want %v", windows[0], want) + } + }) + } +} + +func TestWindowFnInvoker_TimestampOnly(t *testing.T) { + fn := &testWindowFn{BucketSize: 3000} + inv := mustInvoker(t, fn) + + windows := inv.Invoke(1500, nil, nil) + if len(windows) != 1 { + t.Fatalf("got %d windows, want 1", len(windows)) + } + want := IntervalWindow{Start: 0, End: 3000} + if !windows[0].Equals(want) { + t.Errorf("Invoke(1500, nil) = %v, want %v", windows[0], want) + } +} + +func TestWindowFnInvoker_AnyElem(t *testing.T) { + fn := &elemAwareAnyWindowFn{SizeMs: 5000} + inv := mustInvoker(t, fn) + + windows := inv.Invoke(7500, "ignored", nil) + if len(windows) != 1 { + t.Fatalf("got %d windows, want 1", len(windows)) + } + want := IntervalWindow{Start: 5000, End: 10000} + if !windows[0].Equals(want) { + t.Errorf("Invoke(7500, ignored) = %v, want %v", windows[0], want) + } +} + +func TestWindowFnInvoker_ConcreteElem(t *testing.T) { + fn := &elemAwareConcreteWindowFn{DefaultSizeMs: 1000} + inv := mustInvoker(t, fn) + + // Element provides window size of 5000ms. + windows := inv.Invoke(7500, int64(5000), nil) + if len(windows) != 1 { + t.Fatalf("got %d windows, want 1", len(windows)) + } + want := IntervalWindow{Start: 5000, End: 10000} + if !windows[0].Equals(want) { + t.Errorf("Invoke(7500, 5000) = %v, want %v", windows[0], want) + } + + // Element <= 0: falls back to default. + windows = inv.Invoke(1500, int64(0), nil) + if len(windows) != 1 { + t.Fatalf("got %d windows, want 1", len(windows)) + } + want = IntervalWindow{Start: 1000, End: 2000} + if !windows[0].Equals(want) { + t.Errorf("Invoke(1500, 0) = %v, want %v", windows[0], want) + } +} + +func TestWindowFnInvoker_ErrorOnUnregistered(t *testing.T) { + type unregisteredFn struct{} + if _, err := NewWindowFnInvoker(&unregisteredFn{}); err == nil { + t.Error("NewWindowFnInvoker succeeded on an unregistered type, want error") + } +} + +func TestWindowFnInvoker_ErrorOnNil(t *testing.T) { + if _, err := NewWindowFnInvoker(nil); err == nil { + t.Error("NewWindowFnInvoker(nil) succeeded, want error") + } +} + +func mustInvoker(t *testing.T, fn any) *WindowFnInvoker { + t.Helper() + inv, err := NewWindowFnInvoker(fn) + if err != nil { + t.Fatalf("NewWindowFnInvoker(%T) failed: %v", fn, err) + } + return inv +} diff --git a/sdks/go/pkg/beam/core/graph/window/register.go b/sdks/go/pkg/beam/core/graph/window/register.go new file mode 100644 index 000000000000..f599b5b83c2a --- /dev/null +++ b/sdks/go/pkg/beam/core/graph/window/register.go @@ -0,0 +1,131 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package window + +import ( + "fmt" + "reflect" + "sync" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" +) + +// windowFnRegistry maps the struct type of a registered custom WindowFn, with +// the pointer stripped, to the element parameter types of its AssignWindows +// method. Presence in the map rather than the value is what says a type is +// registered. +var ( + windowFnRegistryMu sync.RWMutex + windowFnRegistry = map[reflect.Type][]reflect.Type{} +) + +// LookupWindowFn reports whether t is a registered custom WindowFn struct type +// and returns the element parameter types of its AssignWindows method. elems is +// empty for a timestamp-only signature, holds one type for a single element, +// and two for a KV element taken as a key and a value. +func LookupWindowFn(t reflect.Type) (elems []reflect.Type, ok bool) { + windowFnRegistryMu.RLock() + defer windowFnRegistryMu.RUnlock() + elems, ok = windowFnRegistry[t] + return elems, ok +} + +var ( + eventTimeType = reflect.TypeFor[typex.EventTime]() + windowSliceType = reflect.TypeFor[[]typex.Window]() +) + +// RegisterWindowFn registers a custom WindowFn type so it can be serialized +// and deserialized across process boundaries. Call RegisterWindowFn during +// init for every custom window function type used in the pipeline. +// +// The type parameter T must be a pointer-to-struct type with an +// AssignWindows method of one of the following shapes: +// +// func (f *MyFn) AssignWindows(ts typex.EventTime) []typex.Window +// func (f *MyFn) AssignWindows(ts typex.EventTime, elem T) []typex.Window +// func (f *MyFn) AssignWindows(ts typex.EventTime, k K, v V) []typex.Window +// +// The element parameters mirror how a DoFn receives its main input: a KV +// PCollection arrives as two parameters, anything else as one. +// +// RegisterWindowFn panics if the type is invalid. Registering the same type +// more than once is allowed. +// +// Example: +// +// func init() { +// window.RegisterWindowFn[*myWindowFn]() +// } +func RegisterWindowFn[T any]() { + var v T + t := reflect.TypeOf(v) + if t == nil { + panic("window.RegisterWindowFn: T must not be an untyped nil interface") + } + if t.Kind() != reflect.Pointer || t.Elem().Kind() != reflect.Struct { + panic(fmt.Sprintf("window.RegisterWindowFn: T must be a pointer to struct, got %v", t)) + } + + structType := t.Elem() + + m, ok := t.MethodByName("AssignWindows") + if !ok { + panic(fmt.Sprintf("window.RegisterWindowFn: %v has no AssignWindows method", t)) + } + + elems := validateAssignWindows(t, m) + + windowFnRegistryMu.Lock() + defer windowFnRegistryMu.Unlock() + + // Re-registering a type is harmless: the element types are derived from its + // AssignWindows method, so the value cannot differ. + windowFnRegistry[structType] = elems + + runtime.RegisterType(reflect.TypeOf(v)) +} + +// validateAssignWindows checks that the method has a valid signature and +// returns its element parameter types, empty if it takes only a timestamp. +func validateAssignWindows(ptrType reflect.Type, m reflect.Method) []reflect.Type { + // The method type counts the receiver, so NumIn is one more than the + // parameter list the user wrote: 2, 3 and 4 are the accepted shapes. + mt := m.Type + + if mt.NumOut() != 1 || mt.Out(0) != windowSliceType { + panic(fmt.Sprintf( + "window.RegisterWindowFn: %v.AssignWindows must return []typex.Window, got %v", + ptrType, mt)) + } + if mt.NumIn() < 2 || mt.NumIn() > 4 { + panic(fmt.Sprintf( + "window.RegisterWindowFn: %v.AssignWindows must take (typex.EventTime), (typex.EventTime, T) or (typex.EventTime, K, V), got %d params (excluding receiver)", + ptrType, mt.NumIn()-1)) + } + if mt.In(1) != eventTimeType { + panic(fmt.Sprintf( + "window.RegisterWindowFn: %v.AssignWindows first param must be typex.EventTime, got %v", + ptrType, mt.In(1))) + } + + var elems []reflect.Type + for i := 2; i < mt.NumIn(); i++ { + elems = append(elems, mt.In(i)) + } + return elems +} diff --git a/sdks/go/pkg/beam/core/graph/window/register_test.go b/sdks/go/pkg/beam/core/graph/window/register_test.go new file mode 100644 index 000000000000..67645e6f1178 --- /dev/null +++ b/sdks/go/pkg/beam/core/graph/window/register_test.go @@ -0,0 +1,136 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package window + +import ( + "reflect" + "testing" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" +) + +type regTsOnly struct{} + +func (f *regTsOnly) AssignWindows(typex.EventTime) []typex.Window { return nil } + +type regSingleElem struct{} + +func (f *regSingleElem) AssignWindows(typex.EventTime, int64) []typex.Window { return nil } + +type regKV struct{} + +func (f *regKV) AssignWindows(typex.EventTime, string, int64) []typex.Window { return nil } + +type regTooManyParams struct{} + +func (f *regTooManyParams) AssignWindows(typex.EventTime, string, int64, bool) []typex.Window { + return nil +} + +type regWrongReturn struct{} + +func (f *regWrongReturn) AssignWindows(typex.EventTime) typex.Window { return nil } + +type regWrongFirstParam struct{} + +func (f *regWrongFirstParam) AssignWindows(int64) []typex.Window { return nil } + +type regNoMethod struct{} + +func TestRegisterWindowFn(t *testing.T) { + tests := []struct { + name string + register func() + structType reflect.Type + wantElems []reflect.Type + wantPanic bool + }{ + { + name: "timestamp only", + register: RegisterWindowFn[*regTsOnly], + structType: reflect.TypeFor[regTsOnly](), + wantElems: nil, + }, + { + name: "single element", + register: RegisterWindowFn[*regSingleElem], + structType: reflect.TypeFor[regSingleElem](), + wantElems: []reflect.Type{reflect.TypeFor[int64]()}, + }, + { + name: "kv element", + register: RegisterWindowFn[*regKV], + structType: reflect.TypeFor[regKV](), + wantElems: []reflect.Type{reflect.TypeFor[string](), reflect.TypeFor[int64]()}, + }, + { + name: "too many params", + register: RegisterWindowFn[*regTooManyParams], + wantPanic: true, + }, + { + name: "return is not a window slice", + register: RegisterWindowFn[*regWrongReturn], + wantPanic: true, + }, + { + name: "first param is not an event time", + register: RegisterWindowFn[*regWrongFirstParam], + wantPanic: true, + }, + { + name: "no AssignWindows method", + register: RegisterWindowFn[*regNoMethod], + wantPanic: true, + }, + { + name: "not a pointer to struct", + register: RegisterWindowFn[regTsOnly], + wantPanic: true, + }, + { + // Matches runtime.RegisterFunction and runtime.RegisterType, which + // both tolerate registering the same thing twice. + name: "already registered", + register: RegisterWindowFn[*regTsOnly], + structType: reflect.TypeFor[regTsOnly](), + wantElems: nil, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + panicked := func() (p bool) { + defer func() { p = recover() != nil }() + tc.register() + return + }() + + if panicked != tc.wantPanic { + t.Fatalf("RegisterWindowFn panicked = %v, want %v", panicked, tc.wantPanic) + } + if tc.wantPanic { + return + } + got, ok := LookupWindowFn(tc.structType) + if !ok { + t.Fatal("LookupWindowFn reports the type is not registered") + } + if !reflect.DeepEqual(got, tc.wantElems) { + t.Errorf("LookupWindowFn elements = %v, want %v", got, tc.wantElems) + } + }) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/exec/coder.go b/sdks/go/pkg/beam/core/runtime/exec/coder.go index b68943355383..cac087c0ede0 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/coder.go +++ b/sdks/go/pkg/beam/core/runtime/exec/coder.go @@ -1149,7 +1149,7 @@ func (enc *intervalWindowEncoder) Encode(ws []typex.Window, w io.Writer) error { } for _, elm := range ws { if err := enc.EncodeSingle(elm, w); err != nil { - return nil + return err } } return nil @@ -1157,7 +1157,10 @@ func (enc *intervalWindowEncoder) Encode(ws []typex.Window, w io.Writer) error { func (*intervalWindowEncoder) EncodeSingle(elm typex.Window, w io.Writer) error { // Encoding: upper bound and duration - iw := elm.(window.IntervalWindow) + iw, ok := elm.(window.IntervalWindow) + if !ok { + return errors.Errorf("cannot encode %T with the interval window coder; a custom WindowFn must return window.IntervalWindow values from AssignWindows", elm) + } if err := coder.EncodeEventTime(iw.End, w); err != nil { return err } @@ -1203,7 +1206,11 @@ type intervalWindowValueEncoder struct { } func (e *intervalWindowValueEncoder) Encode(v *FullValue, w io.Writer) error { - return e.EncodeSingle(v.Elm.(window.IntervalWindow), w) + iw, ok := v.Elm.(typex.Window) + if !ok { + return errors.Errorf("cannot encode %T with the interval window coder: not a window", v.Elm) + } + return e.EncodeSingle(iw, w) } type intervalWindowValueDecoder struct { diff --git a/sdks/go/pkg/beam/core/runtime/exec/coder_test.go b/sdks/go/pkg/beam/core/runtime/exec/coder_test.go index 155fd72776e1..2da038f83cef 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/coder_test.go +++ b/sdks/go/pkg/beam/core/runtime/exec/coder_test.go @@ -345,3 +345,33 @@ func TestPaneCoder(t *testing.T) { t.Errorf("got pane non-speculative index %v, want %v", got, want) } } + +// TestIntervalWindowEncoderRejectsOtherWindows checks that a window the +// interval coder cannot represent, which a custom WindowFn is able to return, +// reports an error instead of panicking or truncating the stream. +func TestIntervalWindowEncoderRejectsOtherWindows(t *testing.T) { + enc := MakeWindowEncoder(coder.NewIntervalWindow()) + + tests := []struct { + name string + windows []typex.Window + wantErr bool + }{ + {"interval", []typex.Window{window.IntervalWindow{Start: 0, End: 1000}}, false}, + {"global", []typex.Window{window.GlobalWindow{}}, true}, + { + "bad window after a good one", + []typex.Window{window.IntervalWindow{Start: 0, End: 1000}, window.GlobalWindow{}}, + true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var buf bytes.Buffer + err := enc.Encode(tc.windows, &buf) + if (err != nil) != tc.wantErr { + t.Errorf("Encode(%v) error = %v, want error presence %v", tc.windows, err, tc.wantErr) + } + }) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/exec/translate.go b/sdks/go/pkg/beam/core/runtime/exec/translate.go index 13b40ea0d1c6..22f0eaecaa92 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/translate.go +++ b/sdks/go/pkg/beam/core/runtime/exec/translate.go @@ -16,8 +16,10 @@ package exec import ( + "encoding/json" "fmt" "math/rand" + "reflect" "strconv" "strings" @@ -25,10 +27,12 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx" v1pb "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/v1" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/timers" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/jsonx" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/protox" "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1" @@ -275,6 +279,27 @@ func unmarshalWindowFn(wfn *pipepb.FunctionSpec) (*window.Fn, error) { gap := gapPB.AsDuration() return window.NewSessions(gap), nil + case graphx.URNCustomWindowFn: + var envelope struct { + Type string `json:"type"` + Payload json.RawMessage `json:"payload"` + } + if err := json.Unmarshal(wfn.GetPayload(), &envelope); err != nil { + return nil, errors.Wrapf(err, "unmarshaling custom WindowFn envelope") + } + t, ok := runtime.LookupType(envelope.Type) + if !ok { + return nil, errors.Errorf("custom WindowFn type key %q not found in registry", envelope.Type) + } + if _, ok := window.LookupWindowFn(t); !ok { + return nil, errors.Errorf("type %v is not registered via window.RegisterWindowFn", t) + } + val := reflect.New(t) + if err := jsonx.Unmarshal(val.Interface(), envelope.Payload); err != nil { + return nil, errors.Wrapf(err, "unmarshaling custom WindowFn %v", t) + } + return window.NewCustom(val.Interface()), nil + default: return nil, errors.Errorf("unsupported window type: %v", urn) } @@ -283,7 +308,7 @@ func unmarshalWindowFn(wfn *pipepb.FunctionSpec) (*window.Fn, error) { func unmarshalAndMakeWindowMapping(wmfn *pipepb.FunctionSpec) (WindowMapper, error) { switch urn := wmfn.GetUrn(); urn { case graphx.URNWindowMappingGlobal: - return &windowMapper{wfn: window.NewGlobalWindows()}, nil + return newWindowMapper(window.NewGlobalWindows()) case graphx.URNWindowMappingFixed: var payload pipepb.FixedWindowsPayload if err := proto.Unmarshal(wmfn.GetPayload(), &payload); err != nil { @@ -294,7 +319,7 @@ func unmarshalAndMakeWindowMapping(wmfn *pipepb.FunctionSpec) (WindowMapper, err return nil, err } size := sizePB.AsDuration() - return &windowMapper{wfn: window.NewFixedWindows(size)}, nil + return newWindowMapper(window.NewFixedWindows(size)) case graphx.URNWindowMappingSliding: var payload pipepb.SlidingWindowsPayload if err := proto.Unmarshal(wmfn.GetPayload(), &payload); err != nil { @@ -311,7 +336,31 @@ func unmarshalAndMakeWindowMapping(wmfn *pipepb.FunctionSpec) (WindowMapper, err return nil, err } size := sizePB.AsDuration() - return &windowMapper{wfn: window.NewSlidingWindows(period, size)}, nil + return newWindowMapper(window.NewSlidingWindows(period, size)) + case graphx.URNWindowMappingCustom: + var envelope struct { + Type string `json:"type"` + Payload json.RawMessage `json:"payload"` + } + if err := json.Unmarshal(wmfn.GetPayload(), &envelope); err != nil { + return nil, errors.Wrapf(err, "unmarshaling custom window mapping envelope") + } + t, ok := runtime.LookupType(envelope.Type) + if !ok { + return nil, errors.Errorf("custom WindowFn type key %q not found in registry", envelope.Type) + } + elems, ok := window.LookupWindowFn(t) + if !ok { + return nil, errors.Errorf("type %v is not registered via window.RegisterWindowFn", t) + } + if len(elems) > 0 { + return nil, errors.Errorf("element-aware custom WindowFn %v cannot be used for side input window mapping", t) + } + val := reflect.New(t) + if err := jsonx.Unmarshal(val.Interface(), envelope.Payload); err != nil { + return nil, errors.Wrapf(err, "unmarshaling custom WindowFn %v for window mapping", t) + } + return newWindowMapper(window.NewCustom(val.Interface())) default: return nil, fmt.Errorf("unsupported window mapping fn URN %v", urn) } diff --git a/sdks/go/pkg/beam/core/runtime/exec/translate_test.go b/sdks/go/pkg/beam/core/runtime/exec/translate_test.go index a9917ec456fe..79974218e260 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/translate_test.go +++ b/sdks/go/pkg/beam/core/runtime/exec/translate_test.go @@ -16,6 +16,7 @@ package exec import ( + "encoding/json" "fmt" fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1" "reflect" @@ -26,6 +27,7 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/protox" "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" @@ -514,3 +516,88 @@ func TestNewBuilder(t *testing.T) { }) } } + +// unregisteredWindowFn is known to the type registry but never passed to +// window.RegisterWindowFn. +type unregisteredWindowFn struct{} + +func init() { + runtime.RegisterType(reflect.TypeOf(unregisteredWindowFn{})) +} + +// customWindowFnEnvelope mirrors the JSON envelope makeWindowFn writes for +// URNCustomWindowFn. +type customWindowFnEnvelope struct { + Type string `json:"type"` + Payload json.RawMessage `json:"payload"` +} + +func mustEnvelope(t *testing.T, typeKey string, payload string) []byte { + t.Helper() + b, err := json.Marshal(customWindowFnEnvelope{Type: typeKey, Payload: json.RawMessage(payload)}) + if err != nil { + t.Fatalf("failed to build envelope: %v", err) + } + return b +} + +func TestUnmarshalWindowFnCustom(t *testing.T) { + registered, ok := runtime.TypeKey(reflect.TypeOf(fixedCustomWindowFn{})) + if !ok { + t.Fatal("fixedCustomWindowFn has no type key") + } + unregistered, ok := runtime.TypeKey(reflect.TypeOf(unregisteredWindowFn{})) + if !ok { + t.Fatal("unregisteredWindowFn has no type key") + } + + tests := []struct { + name string + payload []byte + want *window.Fn + wantErr bool + }{ + { + name: "registered type", + payload: mustEnvelope(t, registered, `{"SizeMs":3000}`), + want: window.NewCustom(&fixedCustomWindowFn{SizeMs: 3000}), + }, + { + name: "malformed envelope", + payload: []byte(`{"type":`), + wantErr: true, + }, + { + name: "unknown type key", + payload: mustEnvelope(t, "example.com/nope.MissingFn", `{}`), + wantErr: true, + }, + { + name: "type not registered as a WindowFn", + payload: mustEnvelope(t, unregistered, `{}`), + wantErr: true, + }, + { + name: "malformed inner payload", + payload: mustEnvelope(t, registered, `{"SizeMs":"three thousand"}`), + wantErr: true, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := unmarshalWindowFn(&pipepb.FunctionSpec{ + Urn: graphx.URNCustomWindowFn, + Payload: test.payload, + }) + if (err != nil) != test.wantErr { + t.Fatalf("unmarshalWindowFn error = %v, want error presence %v", err, test.wantErr) + } + if test.wantErr { + return + } + if !got.Equals(test.want) { + t.Errorf("unmarshalWindowFn = %v, want %v", got, test.want) + } + }) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/exec/window.go b/sdks/go/pkg/beam/core/runtime/exec/window.go index fabd6af933d8..3857e5675372 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window.go @@ -31,6 +31,8 @@ type WindowInto struct { UID UnitID Fn *window.Fn Out Node + + invoker *window.WindowFnInvoker // non-nil for CustomWindows } // ID returns the UnitID for this unit. @@ -39,7 +41,19 @@ func (w *WindowInto) ID() UnitID { } func (w *WindowInto) Up(ctx context.Context) error { - return nil + var err error + w.invoker, err = invokerFor(w.Fn) + return err +} + +// invokerFor returns the invoker for a custom WindowFn, or nil for the +// built-in kinds. Callers cache the result rather than rebuilding it per +// element: construction costs a registry lookup and a closure allocation. +func invokerFor(wfn *window.Fn) (*window.WindowFnInvoker, error) { + if wfn.Kind != window.CustomWindows { + return nil, nil + } + return window.NewWindowFnInvoker(wfn.CustomFn) } func (w *WindowInto) StartBundle(ctx context.Context, id string, data DataContext) error { @@ -48,7 +62,7 @@ func (w *WindowInto) StartBundle(ctx context.Context, id string, data DataContex func (w *WindowInto) ProcessElement(ctx context.Context, elm *FullValue, values ...ReStream) error { windowed := &FullValue{ - Windows: assignWindows(w.Fn, elm.Timestamp), + Windows: assignWindows(w.Fn, w.invoker, elm.Timestamp, elm.Elm, elm.Elm2), Timestamp: elm.Timestamp, Elm: elm.Elm, Elm2: elm.Elm2, @@ -57,7 +71,9 @@ func (w *WindowInto) ProcessElement(ctx context.Context, elm *FullValue, values return w.Out.ProcessElement(ctx, windowed, values...) } -func assignWindows(wfn *window.Fn, ts typex.EventTime) []typex.Window { +// assignWindows assigns windows for ts. inv is the cached invoker for +// CustomWindows and is unused for the built-in kinds. +func assignWindows(wfn *window.Fn, inv *window.WindowFnInvoker, ts typex.EventTime, elm, elm2 any) []typex.Window { switch wfn.Kind { case window.GlobalWindows: return window.SingleGlobalWindow @@ -82,6 +98,9 @@ func assignWindows(wfn *window.Fn, ts typex.EventTime) []typex.Window { // each other) will be merged. return []typex.Window{window.IntervalWindow{Start: ts, End: ts.Add(wfn.Gap)}} + case window.CustomWindows: + return inv.Invoke(ts, elm, elm2) + default: panic(fmt.Sprintf("Unexpected window fn: %v", wfn)) } @@ -170,13 +189,30 @@ type WindowMapper interface { type windowMapper struct { wfn *window.Fn + inv *window.WindowFnInvoker // non-nil for CustomWindows +} + +func newWindowMapper(wfn *window.Fn) (*windowMapper, error) { + inv, err := invokerFor(wfn) + if err != nil { + return nil, err + } + return &windowMapper{wfn: wfn, inv: inv}, nil } func (f *windowMapper) MapWindow(w typex.Window) (typex.Window, error) { - candidates := assignWindows(f.wfn, w.MaxTimestamp()) + candidates := assignWindows(f.wfn, f.inv, w.MaxTimestamp(), nil, nil) if len(candidates) == 0 { return nil, fmt.Errorf("failed to map main input window to side input window with WindowFn %v", f.wfn.String()) } + // Picking the last candidate relies on sliding windows appending the + // latest window first, which is an invariant of the built-in kinds only. + // Java offers this maximum-timestamp mapping solely for + // PartitioningWindowFn, which assigns to exactly one window, so require + // that of custom WindowFns rather than picking one arbitrarily. + if f.wfn.Kind == window.CustomWindows && len(candidates) != 1 { + return nil, errors.Errorf("custom WindowFn %v assigned %v windows to the side input window for %v; side input mapping requires exactly one", f.wfn.String(), len(candidates), w) + } // Return earliest candidate window in terms of event time (only relevant for sliding windows) // Sliding windows append the latest window first in assignWindows. return candidates[len(candidates)-1], nil diff --git a/sdks/go/pkg/beam/core/runtime/exec/window_test.go b/sdks/go/pkg/beam/core/runtime/exec/window_test.go index e0bca2a74f4d..20226e60683f 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window_test.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window_test.go @@ -30,9 +30,9 @@ import ( // correct windows for a given timestamp. func TestAssignWindow(t *testing.T) { tests := []struct { - fn *window.Fn - in typex.EventTime - out []typex.Window + fn *window.Fn + in typex.EventTime + want []typex.Window }{ { window.NewGlobalWindows(), @@ -113,22 +113,44 @@ func TestAssignWindow(t *testing.T) { window.IntervalWindow{Start: 60000, End: 120000}, }, }, + { + // Custom window that mimics 3-second fixed windows. + window.NewCustom(&fixedCustomWindowFn{SizeMs: 3000}), + 0, + []typex.Window{ + window.IntervalWindow{Start: 0, End: 3000}, + }, + }, + { + window.NewCustom(&fixedCustomWindowFn{SizeMs: 3000}), + 2999, + []typex.Window{ + window.IntervalWindow{Start: 0, End: 3000}, + }, + }, + { + window.NewCustom(&fixedCustomWindowFn{SizeMs: 3000}), + 3000, + []typex.Window{ + window.IntervalWindow{Start: 3000, End: 6000}, + }, + }, } for _, test := range tests { - out := assignWindows(test.fn, test.in) - if !window.IsEqualList(out, test.out) { - t.Errorf("assignWindows(%v, %v) = %v, want %v", test.fn, test.in, out, test.out) + out := assignWindows(test.fn, mustInvokerFor(t, test.fn), test.in, nil, nil) + if !window.IsEqualList(out, test.want) { + t.Errorf("assignWindows(%v, %v) = %v, want %v", test.fn, test.in, out, test.want) } } } func TestMapWindow(t *testing.T) { tests := []struct { - name string - wfn *window.Fn - in typex.Window - expected typex.Window + name string + wfn *window.Fn + in typex.Window + want typex.Window }{ { "interval to global", @@ -162,23 +184,23 @@ func TestMapWindow(t *testing.T) { }, } for _, test := range tests { - mapper := &windowMapper{wfn: test.wfn} + mapper := mustWindowMapper(t, test.wfn) outputWin, err := mapper.MapWindow(test.in) if err != nil { t.Fatalf("MapWindow for test %v failed, got %v", test.name, err) } - if !outputWin.Equals(test.expected) { - t.Errorf("test %v failed: expected window %v, got %v", test.name, test.expected, outputWin) + if !outputWin.Equals(test.want) { + t.Errorf("test %v failed: got window %v, want %v", test.name, outputWin, test.want) } } } func TestMapWindows(t *testing.T) { tests := []struct { - name string - wFn *window.Fn - in []typex.Window - expect []typex.Window + name string + wFn *window.Fn + in []typex.Window + want []typex.Window }{ { "fixed2fixed", @@ -195,10 +217,10 @@ func TestMapWindows(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - inV, expected := makeNoncedWindowValues(tc.in, tc.expect) + inV, wantFVs := makeNoncedWindowValues(tc.in, tc.want) out := &CaptureNode{UID: 1} - unit := &MapWindows{UID: 2, Fn: &windowMapper{wfn: tc.wFn}, Out: out} + unit := &MapWindows{UID: 2, Fn: mustWindowMapper(t, tc.wFn), Out: out} a := &FixedRoot{UID: 3, Elements: inV, Out: unit} p, err := NewPlan(tc.name, []Unit{a, unit, out}) @@ -212,13 +234,207 @@ func TestMapWindows(t *testing.T) { if err := p.Down(ctx); err != nil { t.Fatalf("down failed: %s", err) } - if !equalList(out.Elements, expected) { - t.Errorf("map_windows returned %v, want %v", extractValues(out.Elements...), extractValues(expected...)) + if !equalList(out.Elements, wantFVs) { + t.Errorf("map_windows returned %v, want %v", extractValues(out.Elements...), extractValues(wantFVs...)) + } + }) + } +} + +func init() { + window.RegisterWindowFn[*fixedCustomWindowFn]() + window.RegisterWindowFn[*elemSizedWindowFn]() + window.RegisterWindowFn[*multiWindowFn]() + window.RegisterWindowFn[*kvSizedWindowFn]() +} + +// kvSizedWindowFn derives the window size from a KV element's value. +type kvSizedWindowFn struct{} + +func (f *kvSizedWindowFn) AssignWindows(ts typex.EventTime, k string, v int64) []typex.Window { + size := typex.EventTime(v) + start := ts - ((ts%size)+size)%size + return []typex.Window{window.IntervalWindow{Start: start, End: start + size}} +} + +// TestWindowIntoKV checks that a KV element reaches AssignWindows as a +// separate key and value rather than the key alone. +func TestWindowIntoKV(t *testing.T) { + tests := []struct { + name string + ts typex.EventTime + key string + val int64 + want typex.Window + }{ + {"value sets 3s size", 1500, "a", 3000, window.IntervalWindow{Start: 0, End: 3000}}, + {"value sets 6s size", 1500, "b", 6000, window.IntervalWindow{Start: 0, End: 6000}}, + {"value selects later window", 4500, "c", 3000, window.IntervalWindow{Start: 3000, End: 6000}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + out := &CaptureNode{UID: 1} + wi := &WindowInto{UID: 2, Fn: window.NewCustom(&kvSizedWindowFn{}), Out: out} + root := &FixedRoot{UID: 3, Elements: []MainInput{{Key: FullValue{ + Windows: window.SingleGlobalWindow, + Timestamp: test.ts, + Elm: test.key, + Elm2: test.val, + }}}, Out: wi} + + p, err := NewPlan("a", []Unit{root, wi, out}) + if err != nil { + t.Fatalf("failed to construct plan: %v", err) + } + if err := p.Execute(ctx, "1", DataContext{}); err != nil { + t.Fatalf("execute failed: %v", err) + } + if err := p.Down(ctx); err != nil { + t.Fatalf("down failed: %v", err) + } + + if len(out.Elements) != 1 { + t.Fatalf("got %v elements, want 1", len(out.Elements)) + } + if got := out.Elements[0].Windows; !window.IsEqualList(got, []typex.Window{test.want}) { + t.Errorf("WindowInto assigned %v, want %v", got, test.want) + } + }) + } +} + +// multiWindowFn assigns every timestamp to two windows, earliest first. +type multiWindowFn struct{} + +func (f *multiWindowFn) AssignWindows(ts typex.EventTime) []typex.Window { + return []typex.Window{ + window.IntervalWindow{Start: 0, End: 1000}, + window.IntervalWindow{Start: 1000, End: 2000}, + } +} + +// TestMapWindowCustom checks that side input mapping accepts a custom +// WindowFn only when it assigns to exactly one window. Picking among several +// candidates relies on an ordering only the built-in kinds guarantee. +func TestMapWindowCustom(t *testing.T) { + tests := []struct { + name string + wfn *window.Fn + in typex.Window + want typex.Window + wantErr bool + }{ + { + name: "single window", + wfn: window.NewCustom(&fixedCustomWindowFn{SizeMs: 1000}), + in: window.IntervalWindow{Start: 100, End: 200}, + want: window.IntervalWindow{Start: 0, End: 1000}, + }, + { + name: "multiple windows rejected", + wfn: window.NewCustom(&multiWindowFn{}), + in: window.IntervalWindow{Start: 100, End: 200}, + wantErr: true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := mustWindowMapper(t, tc.wfn).MapWindow(tc.in) + if tc.wantErr { + if err == nil { + t.Fatalf("MapWindow(%v) = %v, want error", tc.in, got) + } + return + } + if err != nil { + t.Fatalf("MapWindow(%v) failed: %v", tc.in, err) + } + if !got.Equals(tc.want) { + t.Errorf("MapWindow(%v) = %v, want %v", tc.in, got, tc.want) + } + }) + } +} + +// elemSizedWindowFn derives the window size from the element value. +type elemSizedWindowFn struct{} + +func (f *elemSizedWindowFn) AssignWindows(ts typex.EventTime, elem int64) []typex.Window { + size := typex.EventTime(elem) + start := ts - ((ts%size)+size)%size + return []typex.Window{window.IntervalWindow{Start: start, End: start + size}} +} + +func BenchmarkAssignWindowsCustom(b *testing.B) { + fn := window.NewCustom(&fixedCustomWindowFn{SizeMs: 3000}) + inv, err := invokerFor(fn) + if err != nil { + b.Fatalf("invokerFor(%v) failed: %v", fn, err) + } + b.ReportAllocs() + for i := 0; i < b.N; i++ { + assignWindows(fn, inv, 1500, nil, nil) + } +} + +// TestWindowIntoElementAware checks that an element-aware custom WindowFn +// receives the element value rather than the enclosing FullValue. +func TestWindowIntoElementAware(t *testing.T) { + tests := []struct { + name string + ts typex.EventTime + elm int64 + want typex.Window + }{ + {"element sets 3s size", 1500, 3000, window.IntervalWindow{Start: 0, End: 3000}}, + {"element sets 6s size", 1500, 6000, window.IntervalWindow{Start: 0, End: 6000}}, + {"element selects later window", 4500, 3000, window.IntervalWindow{Start: 3000, End: 6000}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + out := &CaptureNode{UID: 1} + wi := &WindowInto{UID: 2, Fn: window.NewCustom(&elemSizedWindowFn{}), Out: out} + root := &FixedRoot{UID: 3, Elements: []MainInput{{Key: FullValue{ + Windows: window.SingleGlobalWindow, + Timestamp: test.ts, + Elm: test.elm, + }}}, Out: wi} + + p, err := NewPlan("a", []Unit{root, wi, out}) + if err != nil { + t.Fatalf("failed to construct plan: %v", err) + } + if err := p.Execute(ctx, "1", DataContext{}); err != nil { + t.Fatalf("execute failed: %v", err) + } + if err := p.Down(ctx); err != nil { + t.Fatalf("down failed: %v", err) + } + + if len(out.Elements) != 1 { + t.Fatalf("got %v elements, want 1", len(out.Elements)) + } + if got := out.Elements[0].Windows; !window.IsEqualList(got, []typex.Window{test.want}) { + t.Errorf("WindowInto assigned %v, want %v", got, test.want) } }) } } +// fixedCustomWindowFn is a test custom WindowFn that mimics fixed windows. +type fixedCustomWindowFn struct { + SizeMs int64 +} + +func (f *fixedCustomWindowFn) AssignWindows(ts typex.EventTime) []typex.Window { + size := typex.EventTime(f.SizeMs) + // Euclidean remainder; correct floor for negative ts. + start := ts - ((ts%size)+size)%size + return []typex.Window{window.IntervalWindow{Start: start, End: start + size}} +} + func makeNoncedWindowValues(in []typex.Window, expect []typex.Window) ([]MainInput, []FullValue) { if len(in) != len(expect) { panic("provided window slices must be the same length") @@ -233,3 +449,21 @@ func makeNoncedWindowValues(in []typex.Window, expect []typex.Window) ([]MainInp } return inV, expectV } + +func mustInvokerFor(t *testing.T, wfn *window.Fn) *window.WindowFnInvoker { + t.Helper() + inv, err := invokerFor(wfn) + if err != nil { + t.Fatalf("invokerFor(%v) failed: %v", wfn, err) + } + return inv +} + +func mustWindowMapper(t *testing.T, wfn *window.Fn) *windowMapper { + t.Helper() + m, err := newWindowMapper(wfn) + if err != nil { + t.Fatalf("newWindowMapper(%v) failed: %v", wfn, err) + } + return m +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go b/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go index 265d936a71ec..7cf9d8fa1d9c 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go @@ -18,12 +18,16 @@ package graphx import ( + "encoding/json" "reflect" "strings" "testing" + "time" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" v1pb "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/v1" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" ) func TestEncodeType(t *testing.T) { @@ -89,3 +93,96 @@ func TestEncodeType(t *testing.T) { } }) } + +func TestWindowFnRoundTrip(t *testing.T) { + tests := []struct { + name string + fn *window.Fn + }{ + {"global", window.NewGlobalWindows()}, + {"fixed", window.NewFixedWindows(5 * time.Second)}, + {"sliding", window.NewSlidingWindows(1*time.Second, 3*time.Second)}, + {"sessions", window.NewSessions(10 * time.Second)}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + pb := encodeWindowFn(tc.fn) + got := decodeWindowFn(pb) + if !tc.fn.Equals(got) { + t.Errorf("roundtrip mismatch: got %v, want %v", got, tc.fn) + } + }) + } +} + +// TestEncodeWindowFnCustomKeepsKind covers the internal v1 proto path, which +// preserves only the Kind so that EncodeMultiEdge does not fail. Custom +// WindowFns travel via the Beam model proto instead; see TestMakeWindowFnCustom. +func TestEncodeWindowFnCustomKeepsKind(t *testing.T) { + fn := &window.Fn{Kind: window.CustomWindows} + pb := encodeWindowFn(fn) + got := decodeWindowFn(pb) + if got.Kind != window.CustomWindows { + t.Errorf("kind mismatch: got %v, want %v", got.Kind, window.CustomWindows) + } +} + +// serializeTestWindowFn is a custom WindowFn used to check the FunctionSpec +// that makeWindowFn emits. +type serializeTestWindowFn struct { + SizeMs int64 +} + +func (f *serializeTestWindowFn) AssignWindows(ts typex.EventTime) []typex.Window { + return []typex.Window{window.IntervalWindow{Start: ts, End: ts + typex.EventTime(f.SizeMs)}} +} + +func init() { + window.RegisterWindowFn[*serializeTestWindowFn]() +} + +// TestMakeWindowFnCustom checks the URN and JSON envelope that a custom +// WindowFn is marshalled into, which unmarshalWindowFn has to read back. +func TestMakeWindowFnCustom(t *testing.T) { + spec, err := makeWindowFn(window.NewCustom(&serializeTestWindowFn{SizeMs: 3000})) + if err != nil { + t.Fatalf("makeWindowFn failed: %v", err) + } + if got := spec.GetUrn(); got != URNCustomWindowFn { + t.Errorf("urn = %v, want %v", got, URNCustomWindowFn) + } + + var envelope struct { + Type string `json:"type"` + Payload json.RawMessage `json:"payload"` + } + if err := json.Unmarshal(spec.GetPayload(), &envelope); err != nil { + t.Fatalf("payload is not a JSON envelope: %v", err) + } + + wantKey, ok := runtime.TypeKey(reflect.TypeOf(serializeTestWindowFn{})) + if !ok { + t.Fatal("serializeTestWindowFn has no type key") + } + if envelope.Type != wantKey { + t.Errorf("envelope type = %v, want %v", envelope.Type, wantKey) + } + + var back serializeTestWindowFn + if err := json.Unmarshal(envelope.Payload, &back); err != nil { + t.Fatalf("envelope payload does not decode into the WindowFn: %v", err) + } + if back.SizeMs != 3000 { + t.Errorf("envelope payload SizeMs = %v, want 3000", back.SizeMs) + } +} + +// TestMakeWindowFnCustomUnregistered checks that marshalling refuses a +// WindowFn the type registry does not know, since the harness could not look +// it up on the other side. +func TestMakeWindowFnCustomUnregistered(t *testing.T) { + type localWindowFn struct{} + if _, err := makeWindowFn(&window.Fn{Kind: window.CustomWindows, CustomFn: &localWindowFn{}}); err == nil { + t.Error("makeWindowFn succeeded for an unregistered type, want error") + } +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/translate.go b/sdks/go/pkg/beam/core/runtime/graphx/translate.go index 3994397e7ba5..ef626513e3cd 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/translate.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/translate.go @@ -17,7 +17,9 @@ package graphx import ( "context" + "encoding/json" "fmt" + "reflect" "sort" "strings" @@ -26,11 +28,14 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window/trigger" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/contextreg" v1pb "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/v1" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/pipelinex" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/state" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/jsonx" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/protox" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx" "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" "github.com/apache/beam/sdks/v2/go/pkg/beam/options/resource" @@ -58,6 +63,10 @@ const ( URNFixedWindowsWindowFn = "beam:window_fn:fixed_windows:v1" URNSlidingWindowsWindowFn = "beam:window_fn:sliding_windows:v1" URNSessionsWindowFn = "beam:window_fn:session_windows:v1" + // URNCustomWindowFn is Go specific, mirroring + // beam:window_fn:serialized_java:v1 and beam:window_fn:pickled_python:v1. + // Other SDKs and runners cannot rehydrate it. + URNCustomWindowFn = "beam:window_fn:serialized_go:v1" // SDK constants URNDoFn = "beam:go:transform:dofn:v1" @@ -69,6 +78,7 @@ const ( URNWindowMappingGlobal = "beam:go:windowmapping:global:v1" URNWindowMappingFixed = "beam:go:windowmapping:fixed:v1" URNWindowMappingSliding = "beam:go:windowmapping:sliding:v1" + URNWindowMappingCustom = "beam:go:windowmapping:custom:v1" URNProgressReporting = "beam:protocol:progress_reporting:v1" URNMultiCore = "beam:protocol:multi_core_bundle_processing:v1" @@ -358,6 +368,11 @@ func getSideWindowMappingUrn(winFn *window.Fn) string { mappingUrn = URNWindowMappingSliding case window.Sessions: panic("session windowing is not supported for side inputs") + case window.CustomWindows: + if winFn.NeedsElement() { + panic("element-aware custom WindowFn is not supported for side inputs") + } + mappingUrn = URNWindowMappingCustom } return mappingUrn } @@ -1461,6 +1476,34 @@ func makeWindowFn(w *window.Fn) (*pipepb.FunctionSpec, error) { }, ), }, nil + case window.CustomWindows: + t := reflect.TypeOf(w.CustomFn) + structType := reflectx.SkipPtr(t) + if _, ok := window.LookupWindowFn(structType); !ok { + return nil, errors.Errorf("custom WindowFn type %v is not registered; call window.RegisterWindowFn during init()", t) + } + key, ok := runtime.TypeKey(structType) + if !ok { + return nil, errors.Errorf("custom WindowFn type %v has no type key", t) + } + structPayload, err := jsonx.Marshal(w.CustomFn) + if err != nil { + return nil, errors.Wrapf(err, "marshaling custom WindowFn %v", t) + } + envelope, err := json.Marshal(struct { + Type string `json:"type"` + Payload json.RawMessage `json:"payload"` + }{ + Type: key, + Payload: structPayload, + }) + if err != nil { + return nil, errors.Wrapf(err, "marshaling custom WindowFn envelope") + } + return &pipepb.FunctionSpec{ + Urn: URNCustomWindowFn, + Payload: envelope, + }, nil default: return nil, errors.Errorf("unexpected windowing strategy: %v", w) } @@ -1470,7 +1513,7 @@ func makeWindowCoder(w *window.Fn) (*coder.WindowCoder, error) { switch w.Kind { case window.GlobalWindows: return coder.NewGlobalWindow(), nil - case window.FixedWindows, window.SlidingWindows, window.Sessions, URNSlidingWindowsWindowFn: + case window.FixedWindows, window.SlidingWindows, window.Sessions, window.CustomWindows, URNSlidingWindowsWindowFn: return coder.NewIntervalWindow(), nil default: return nil, errors.Errorf("unexpected windowing strategy for coder: %v", w) diff --git a/sdks/go/pkg/beam/windowing.go b/sdks/go/pkg/beam/windowing.go index 0977cf3d094a..b09ab575b8e2 100644 --- a/sdks/go/pkg/beam/windowing.go +++ b/sdks/go/pkg/beam/windowing.go @@ -17,11 +17,14 @@ package beam import ( "fmt" + "reflect" "time" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window/trigger" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx" "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" ) @@ -84,6 +87,9 @@ func TryWindowInto(s Scope, wfn *window.Fn, col PCollection, opts ...WindowIntoO if !col.IsValid() { return PCollection{}, errors.New("invalid input pcollection") } + if err := validateWindowFnElements(wfn, col.Type()); err != nil { + return PCollection{}, err + } ws := window.WindowingStrategy{Fn: wfn, Trigger: trigger.DefaultTrigger{}} for _, opt := range opts { switch opt := opt.(type) { @@ -105,3 +111,40 @@ func TryWindowInto(s Scope, wfn *window.Fn, col PCollection, opts ...WindowIntoO ret := PCollection{edge.Output[0].To} return ret, nil } + +// validateWindowFnElements rejects a custom WindowFn whose AssignWindows +// element parameters cannot receive the PCollection's element. WindowInto binds +// the element the way a DoFn receives its main input, so a KV supplies a key +// and a value while anything else supplies a single element. +func validateWindowFnElements(wfn *window.Fn, t typex.FullType) error { + if wfn.Kind != window.CustomWindows { + return nil + } + elems, ok := window.LookupWindowFn(reflectx.SkipPtr(reflect.TypeOf(wfn.CustomFn))) + if !ok || len(elems) == 0 { + return nil + } + + var supplied []typex.FullType + switch { + case typex.IsKV(t): + supplied = t.Components() + case t.Class() == typex.Composite: + return errors.Errorf("element-aware custom WindowFn %T cannot window a %v PCollection", wfn.CustomFn, t) + default: + supplied = []typex.FullType{t} + } + + if len(elems) != len(supplied) { + return errors.Errorf("custom WindowFn %T takes %v element parameter(s) in AssignWindows, but a %v PCollection supplies %v", wfn.CustomFn, len(elems), t, len(supplied)) + } + for i, elem := range elems { + if supplied[i].Class() == typex.Universal { + continue // Not bound yet, so nothing to check here. + } + if !supplied[i].Type().AssignableTo(elem) { + return errors.Errorf("custom WindowFn %T AssignWindows element parameter %v is %v, which cannot receive %v from a %v PCollection", wfn.CustomFn, i+1, elem, supplied[i].Type(), t) + } + } + return nil +} diff --git a/sdks/go/pkg/beam/windowing_test.go b/sdks/go/pkg/beam/windowing_test.go new file mode 100644 index 000000000000..d934978e2e7c --- /dev/null +++ b/sdks/go/pkg/beam/windowing_test.go @@ -0,0 +1,100 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package beam + +import ( + "reflect" + "testing" + "time" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" +) + +type wfnTsOnly struct{} + +func (f *wfnTsOnly) AssignWindows(typex.EventTime) []typex.Window { return nil } + +type wfnInt64 struct{} + +func (f *wfnInt64) AssignWindows(typex.EventTime, int64) []typex.Window { return nil } + +type wfnAny struct{} + +func (f *wfnAny) AssignWindows(typex.EventTime, any) []typex.Window { return nil } + +type wfnKV struct{} + +func (f *wfnKV) AssignWindows(typex.EventTime, string, int64) []typex.Window { return nil } + +type wfnBytes struct{} + +func (f *wfnBytes) AssignWindows(typex.EventTime, []byte) []typex.Window { return nil } + +func init() { + window.RegisterWindowFn[*wfnTsOnly]() + window.RegisterWindowFn[*wfnInt64]() + window.RegisterWindowFn[*wfnAny]() + window.RegisterWindowFn[*wfnKV]() + window.RegisterWindowFn[*wfnBytes]() +} + +func TestValidateWindowFnElements(t *testing.T) { + int64T := typex.New(reflect.TypeFor[int64]()) + stringT := typex.New(reflect.TypeFor[string]()) + kvT := typex.NewKV(stringT, int64T) + + tests := []struct { + name string + wfn *window.Fn + col typex.FullType + wantErr bool + }{ + {"built-in ignores shape", window.NewFixedWindows(time.Second), kvT, false}, + {"timestamp only ignores shape", window.NewCustom(&wfnTsOnly{}), kvT, false}, + {"single element matches", window.NewCustom(&wfnInt64{}), int64T, false}, + {"any element accepts anything", window.NewCustom(&wfnAny{}), stringT, false}, + {"kv matches", window.NewCustom(&wfnKV{}), kvT, false}, + {"single element type mismatch", window.NewCustom(&wfnInt64{}), stringT, true}, + {"single element against kv", window.NewCustom(&wfnInt64{}), kvT, true}, + {"kv against single element", window.NewCustom(&wfnKV{}), int64T, true}, + {"kv component mismatch", window.NewCustom(&wfnKV{}), typex.NewKV(int64T, int64T), true}, + {"element aware against cogbk", window.NewCustom(&wfnInt64{}), typex.NewCoGBK(stringT, int64T), true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + err := validateWindowFnElements(tc.wfn, tc.col) + if (err != nil) != tc.wantErr { + t.Errorf("validateWindowFnElements(%v, %v) error = %v, want error presence %v", tc.wfn, tc.col, err, tc.wantErr) + } + }) + } +} + +// TestTryWindowIntoValidatesElements checks that TryWindowInto performs the +// element validation rather than leaving the mismatch to fail inside a bundle. +func TestTryWindowIntoValidatesElements(t *testing.T) { + p := NewPipeline() + s := p.Root() + col := Impulse(s) // PCollection<[]byte> + + if _, err := TryWindowInto(s, window.NewCustom(&wfnBytes{}), col); err != nil { + t.Errorf("TryWindowInto with a matching WindowFn failed: %v", err) + } + if _, err := TryWindowInto(s, window.NewCustom(&wfnInt64{}), col); err == nil { + t.Error("TryWindowInto with a mismatched WindowFn succeeded, want error") + } +} diff --git a/sdks/go/test/integration/integration.go b/sdks/go/test/integration/integration.go index 535a7fb8413c..50d3e3332f47 100644 --- a/sdks/go/test/integration/integration.go +++ b/sdks/go/test/integration/integration.go @@ -143,6 +143,12 @@ var portableFilters = []string{ // no support for BundleFinalizer "TestParDoBundleFinalizer.*", + + // Custom WindowFns use a Go-specific WindowFn URN that Java-based + // runners reject when rehydrating the windowing strategy. + "TestWindowSums_Custom", + "TestWindowSums_ElementAware", + "TestValidateCustomWindowedSideInputs", } var prismFilters = []string{ @@ -204,6 +210,12 @@ var flinkFilters = []string{ // no support for BundleFinalizer "TestParDoBundleFinalizer.*", + + // Custom WindowFns use a Go-specific WindowFn URN that Java-based + // runners reject when rehydrating the windowing strategy. + "TestWindowSums_Custom", + "TestWindowSums_ElementAware", + "TestValidateCustomWindowedSideInputs", } var sparkFilters = []string{ @@ -242,6 +254,12 @@ var sparkFilters = []string{ "TestTimers_ProcessingTime_Unbounded", // Side inputs in executable stage not supported. // no support for BundleFinalizer "TestParDoBundleFinalizer.*", + + // Custom WindowFns use a Go-specific WindowFn URN that Java-based + // runners reject when rehydrating the windowing strategy. + "TestWindowSums_Custom", + "TestWindowSums_ElementAware", + "TestValidateCustomWindowedSideInputs", } var dataflowFilters = []string{ @@ -284,6 +302,12 @@ var dataflowFilters = []string{ "TestOomParDo", // Runner V2 doesn't support OrderedListState SDK feature. "TestOrderedListState", + + // Custom WindowFns use a Go-specific WindowFn URN that Java-based + // runners reject when rehydrating the windowing strategy. + "TestWindowSums_Custom", + "TestWindowSums_ElementAware", + "TestValidateCustomWindowedSideInputs", } // CheckFilters checks if an integration test is filtered to be skipped, either diff --git a/sdks/go/test/integration/primitives/windowinto.go b/sdks/go/test/integration/primitives/windowinto.go index f5d01bdfbba5..436e45ce6c0f 100644 --- a/sdks/go/test/integration/primitives/windowinto.go +++ b/sdks/go/test/integration/primitives/windowinto.go @@ -16,12 +16,14 @@ package primitives import ( + "reflect" "time" "github.com/apache/beam/sdks/v2/go/pkg/beam" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window/trigger" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" "github.com/apache/beam/sdks/v2/go/pkg/beam/register" "github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert" "github.com/apache/beam/sdks/v2/go/pkg/beam/testing/teststream" @@ -36,6 +38,15 @@ func init() { register.Emitter3[beam.EventTime, string, int]() register.Emitter1[int]() register.Iter1[int]() + + window.RegisterWindowFn[*customFixedWindowFn]() + window.RegisterWindowFn[*elemAwareWindowFn]() + + register.DoFn2x0[[]byte, func(beam.EventTime, elemWithSize)](&createElemAwareData{}) + register.Emitter2[beam.EventTime, elemWithSize]() + register.Function1x1(extractValue) + + beam.RegisterType(reflect.TypeOf((*elemWithSize)(nil)).Elem()) } // createTimestampedData produces data timestamped with the ordinal. @@ -413,3 +424,126 @@ func TriggerOrFinally(s beam.Scope) { beam.Trigger(trigger), }, 4) } + +// customFixedWindowFn is a user-defined custom WindowFn that replicates +// 3-second fixed windows, used to validate the custom WindowFn path +// end-to-end against the known-correct output from WindowSums_GBK. +type customFixedWindowFn struct { + SizeMs int64 +} + +func (f *customFixedWindowFn) AssignWindows(ts typex.EventTime) []typex.Window { + size := typex.EventTime(f.SizeMs) + // Euclidean remainder, so the floor is correct for negative ts too. + start := ts - ((ts%size)+size)%size + return []typex.Window{window.IntervalWindow{Start: start, End: start + size}} +} + +// ValidateCustomWindowedSideInputs checks that side inputs windowed with +// a custom WindowFn produce the same results as the equivalent fixed windows. +// Uses 1s custom fixed windows for the side input, same as "Fixed-Same" in +// ValidateWindowedSideInputs, expecting output 2, 4, 6. +func ValidateCustomWindowedSideInputs(s beam.Scope) { + timestampedData := beam.ParDo(s, &createTimestampedData{Data: []int{1, 2, 3}}, beam.Impulse(s)) + timestampedData = beam.DropKey(s, timestampedData) + + windowSize := 1 * time.Second + customSideFn := window.NewCustom(&customFixedWindowFn{SizeMs: 1000}) + + // Main in standard 1s fixed windows, side in custom 1s fixed windows. + // Each window has one element; the side input in the same window adds + // the element to itself: 1+1=2, 2+2=4, 3+3=6. + wData := beam.WindowInto(s.Scope("MainWindow"), window.NewFixedWindows(windowSize), timestampedData) + wSide := beam.WindowInto(s.Scope("SideWindow"), customSideFn, timestampedData) + sums := beam.ParDo(s.Scope("Combine"), sumSideInputs, wData, beam.SideInput{Input: wSide}) + sums = beam.WindowInto(s.Scope("Rewindow"), window.NewGlobalWindows(), sums) + passert.Equals(s, sums, 2, 4, 6) +} + +// WindowSums_Custom validates that a custom WindowFn produces the same +// results as a 3-second fixed window. The magic square rows sum to 15 +// in each window, same as WindowSums_GBK. +func WindowSums_Custom(s beam.Scope) { + timestampedData := beam.ParDo(s, &createTimestampedData{Data: []int{4, 9, 2, 3, 5, 7, 8, 1, 6}}, beam.Impulse(s)) + + wfn := window.NewCustom(&customFixedWindowFn{SizeMs: 3000}) + windowed := beam.WindowInto(s.Scope("Custom"), wfn, timestampedData) + sums := gbkSumPerKey(s.Scope("Sum"), windowed) + sums = beam.WindowInto(s.Scope("Rewindow"), window.NewGlobalWindows(), sums) + sums = beam.DropKey(s, sums) + passert.Equals(s, sums, 15, 15, 15) +} + +// elemWithSize carries a value and a window size in milliseconds. +// Each element dictates which window it belongs to. +type elemWithSize struct { + Value int + SizeMs int64 +} + +// elemAwareWindowFn uses the element's SizeMs field to determine the +// window size, demonstrating data-driven window assignment. +type elemAwareWindowFn struct{} + +func (f *elemAwareWindowFn) AssignWindows(ts typex.EventTime, elem elemWithSize) []typex.Window { + size := typex.EventTime(elem.SizeMs) + if size <= 0 { + size = 1000 // fallback: 1s + } + start := ts - ((ts%size)+size)%size + return []typex.Window{window.IntervalWindow{Start: start, End: start + size}} +} + +// createElemAwareData produces elemWithSize values with timestamps. +// The PCollection is single-valued (not KV) so that WindowInto sees +// the elemWithSize directly via elm.Elm. +type createElemAwareData struct { + Data []elemWithSize +} + +func (f *createElemAwareData) ProcessElement(_ []byte, emit func(beam.EventTime, elemWithSize)) { + for i, v := range f.Data { + timestamp := mtime.FromMilliseconds(int64((i + 1) * 1000)).Subtract(10 * time.Millisecond) + emit(timestamp, v) + } +} + +// WindowSums_ElementAware validates that an element-aware custom WindowFn +// correctly routes elements into different windows based on their content. +// +// We emit 6 elements at timestamps 990ms, 1990ms, ..., 5990ms. +// The first 3 elements carry SizeMs=3000 (3s windows) and values 4, 9, 2. +// The last 3 elements carry SizeMs=6000 (6s windows) and values 3, 5, 7. +// +// With 3s windows: ts 990->[0,3000), ts 1990->[0,3000), ts 2990->[0,3000) +// +// -> sum = 4+9+2 = 15 +// +// With 6s windows: ts 3990->[0,6000), ts 4990->[0,6000), ts 5990->[0,6000) +// +// -> sum = 3+5+7 = 15 +// +// After windowing, we extract values, add a fixed key, GBK, and sum. +func WindowSums_ElementAware(s beam.Scope) { + data := []elemWithSize{ + {Value: 4, SizeMs: 3000}, + {Value: 9, SizeMs: 3000}, + {Value: 2, SizeMs: 3000}, + {Value: 3, SizeMs: 6000}, + {Value: 5, SizeMs: 6000}, + {Value: 7, SizeMs: 6000}, + } + timestampedData := beam.ParDo(s, &createElemAwareData{Data: data}, beam.Impulse(s)) + + wfn := window.NewCustom(&elemAwareWindowFn{}) + windowed := beam.WindowInto(s.Scope("ElemAware"), wfn, timestampedData) + // Extract the Value field for summation. + values := beam.ParDo(s.Scope("ExtractValue"), extractValue, windowed) + sums := stats.Sum(s.Scope("Sum"), values) + sums = beam.WindowInto(s.Scope("Rewindow"), window.NewGlobalWindows(), sums) + passert.Equals(s, sums, 15, 15) +} + +func extractValue(e elemWithSize) int { + return e.Value +} diff --git a/sdks/go/test/integration/primitives/windowinto_test.go b/sdks/go/test/integration/primitives/windowinto_test.go index 39a1df6e9e74..02502846c2f5 100644 --- a/sdks/go/test/integration/primitives/windowinto_test.go +++ b/sdks/go/test/integration/primitives/windowinto_test.go @@ -97,3 +97,18 @@ func TestTriggerOrFinally(t *testing.T) { integration.CheckFilters(t) ptest.BuildAndRun(t, TriggerOrFinally) } + +func TestWindowSums_Custom(t *testing.T) { + integration.CheckFilters(t) + ptest.BuildAndRun(t, WindowSums_Custom) +} + +func TestValidateCustomWindowedSideInputs(t *testing.T) { + integration.CheckFilters(t) + ptest.BuildAndRun(t, ValidateCustomWindowedSideInputs) +} + +func TestWindowSums_ElementAware(t *testing.T) { + integration.CheckFilters(t) + ptest.BuildAndRun(t, WindowSums_ElementAware) +}