From cb7fa60a869238c91f0c4be5f218b45901c65810 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:18:49 +0100 Subject: [PATCH 01/11] Add custom window support in Go SDK Custom WindowFns are plain structs, registered with RegisterWindowFn and applied with window.NewCustom. Registration checks the AssignWindows method by reflection, so a bad signature fails when the pipeline is built rather than inside a bundle. AssignWindows takes an event time and, optionally, the element: AssignWindows(ts typex.EventTime) []typex.Window AssignWindows(ts typex.EventTime, elem T) []typex.Window The second form lets the element decide which window it lands in. Using structs rather than an interface keeps this consistent with how DoFns are written, and lets one registry hold the element type that dispatch and translation both need. WindowInto passes the element to AssignWindows rather than the FullValue around it, and keeps the invoker that Up builds instead of making a new one per element. Building one costs a reflect.TypeOf, a read lock on the package-level registry and a closure, so on a per-element path it would put every bundle goroutine on the same lock. BenchmarkAssignWindowsCustom, Apple M3 Max, go1.27: new invoker per element 43.67 ns/op 72 B/op 4 allocs/op invoker built once 16.40 ns/op 32 B/op 2 allocs/op Serialization reuses the Beam model proto: a FunctionSpec with a URN and a JSON payload. The internal v1 proto needs no new fields. Custom windows are non-merging and must return IntervalWindow values. Only non-KV PCollections work: on a KV input the element is the key alone. Addresses #20627 --- sdks/go/pkg/beam/core/graph/window/fn.go | 46 ++++++ sdks/go/pkg/beam/core/graph/window/fn_test.go | 75 +++++++++ sdks/go/pkg/beam/core/graph/window/invoke.go | 117 ++++++++++++++ .../pkg/beam/core/graph/window/invoke_test.go | 147 ++++++++++++++++++ .../go/pkg/beam/core/graph/window/register.go | 147 ++++++++++++++++++ .../pkg/beam/core/runtime/exec/translate.go | 55 ++++++- sdks/go/pkg/beam/core/runtime/exec/window.go | 29 +++- .../pkg/beam/core/runtime/exec/window_test.go | 111 ++++++++++++- .../core/runtime/graphx/serialize_test.go | 35 +++++ .../pkg/beam/core/runtime/graphx/translate.go | 38 ++++- .../test/integration/primitives/windowinto.go | 136 ++++++++++++++++ .../integration/primitives/windowinto_test.go | 15 ++ 12 files changed, 943 insertions(+), 8 deletions(-) create mode 100644 sdks/go/pkg/beam/core/graph/window/invoke.go create mode 100644 sdks/go/pkg/beam/core/graph/window/invoke_test.go create mode 100644 sdks/go/pkg/beam/core/graph/window/register.go diff --git a/sdks/go/pkg/beam/core/graph/window/fn.go b/sdks/go/pkg/beam/core/graph/window/fn.go index df32a97b89c2..a1b890e214de 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 LookupWindowFnMeta(st) == nil { + 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,24 @@ 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() + } + if meta := LookupWindowFnMeta(t); meta != nil { + return meta.NeedsElement() + } + return false } // TODO(herohde) 4/17/2018: do we need to expose the window type as well? @@ -82,6 +124,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[%v]", w.Kind, reflect.TypeOf(w.CustomFn)) default: return string(w.Kind) } @@ -105,6 +149,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..95bacc2c288a --- /dev/null +++ b/sdks/go/pkg/beam/core/graph/window/invoke.go @@ -0,0 +1,117 @@ +// 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" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" +) + +// 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 (interface{}). +type anyElemAssigner interface { + AssignWindows(typex.EventTime, any) []typex.Window +} + +// WindowFnInvoker wraps a custom WindowFn instance and dispatches +// AssignWindows calls through one of three paths, chosen once at +// construction: +// +// 1. Type-assert to tsOnlyAssigner — zero allocation. +// 2. Type-assert to anyElemAssigner — zero allocation. +// 3. reflect.Value.Method.Call — small per-call allocation. +type WindowFnInvoker struct { + call func(typex.EventTime, any) []typex.Window + needsElement bool +} + +// NewWindowFnInvoker builds an invoker for fn. The concrete type of fn +// must have been previously registered via RegisterWindowFn. +// Panics if fn's type is not registered. +func NewWindowFnInvoker(fn any) *WindowFnInvoker { + t := reflect.TypeOf(fn) + if t == nil { + panic("window.NewWindowFnInvoker: fn must not be nil") + } + structType := t + if t.Kind() == reflect.Pointer { + structType = t.Elem() + } + + meta := LookupWindowFnMeta(structType) + if meta == nil { + panic(fmt.Sprintf("window.NewWindowFnInvoker: type %v is not registered; call window.RegisterWindowFn during init()", t)) + } + + inv := &WindowFnInvoker{needsElement: meta.NeedsElement()} + + if !meta.NeedsElement() { + // Fast path 1: timestamp-only. + if a, ok := fn.(tsOnlyAssigner); ok { + inv.call = func(ts typex.EventTime, _ any) []typex.Window { + return a.AssignWindows(ts) + } + return inv + } + // Unreachable for well-typed registrations, but fall through to reflect. + } else { + // Fast path 2: element typed as any. + if a, ok := fn.(anyElemAssigner); ok { + inv.call = func(ts typex.EventTime, elem any) []typex.Window { + return a.AssignWindows(ts, elem) + } + return inv + } + } + + // Slow path 3: concrete element type — use reflect. + rv := reflect.ValueOf(fn) + m := rv.MethodByName("AssignWindows") + if !m.IsValid() { + panic(fmt.Sprintf("window.NewWindowFnInvoker: %v has no AssignWindows method", t)) + } + + if !meta.NeedsElement() { + inv.call = func(ts typex.EventTime, _ any) []typex.Window { + out := m.Call([]reflect.Value{reflect.ValueOf(ts)}) + return out[0].Interface().([]typex.Window) + } + } else { + inv.call = func(ts typex.EventTime, elem any) []typex.Window { + out := m.Call([]reflect.Value{reflect.ValueOf(ts), reflect.ValueOf(elem)}) + return out[0].Interface().([]typex.Window) + } + } + return inv +} + +// Invoke calls AssignWindows on the underlying WindowFn. +// If the WindowFn is timestamp-only, elem is ignored. +func (inv *WindowFnInvoker) Invoke(ts typex.EventTime, elem any) []typex.Window { + return inv.call(ts, elem) +} + +// NeedsElement reports whether the underlying WindowFn accepts an element. +func (inv *WindowFnInvoker) NeedsElement() bool { + return inv.needsElement +} 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..f70f4750ca27 --- /dev/null +++ b/sdks/go/pkg/beam/core/graph/window/invoke_test.go @@ -0,0 +1,147 @@ +// 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 — fast path 2. +type elemAwareAnyWindowFn struct { + SizeMs int64 +} + +func (f *elemAwareAnyWindowFn) AssignWindows(ts typex.EventTime, _ any) []typex.Window { + size := typex.EventTime(f.SizeMs) + start := ts - (ts % size) + return []typex.Window{IntervalWindow{Start: start, End: start + size}} +} + +// elemAwareConcreteWindowFn accepts a concrete element type — reflect path. +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) + return []typex.Window{IntervalWindow{Start: start, End: start + size}} +} + +func init() { + RegisterWindowFn[*elemAwareAnyWindowFn]() + RegisterWindowFn[*elemAwareConcreteWindowFn]() +} + +func TestWindowFnInvoker_TimestampOnly(t *testing.T) { + fn := &testWindowFn{BucketSize: 3000} + inv := NewWindowFnInvoker(fn) + + if inv.NeedsElement() { + t.Fatal("NeedsElement() = true, want false") + } + + windows := inv.Invoke(1500, 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 := NewWindowFnInvoker(fn) + + if !inv.NeedsElement() { + t.Fatal("NeedsElement() = false, want true") + } + + windows := inv.Invoke(7500, "ignored") + 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 := NewWindowFnInvoker(fn) + + if !inv.NeedsElement() { + t.Fatal("NeedsElement() = false, want true") + } + + // Element provides window size of 5000ms. + windows := inv.Invoke(7500, 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, 5000) = %v, want %v", windows[0], want) + } + + // Element <= 0: falls back to default. + windows = inv.Invoke(1500, int64(0)) + 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_NeedsElementCorrectness(t *testing.T) { + tests := []struct { + name string + fn any + want bool + }{ + {"timestamp-only", &testWindowFn{BucketSize: 1000}, false}, + {"any-elem", &elemAwareAnyWindowFn{SizeMs: 1000}, true}, + {"concrete-elem", &elemAwareConcreteWindowFn{DefaultSizeMs: 1000}, true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + inv := NewWindowFnInvoker(tc.fn) + if got := inv.NeedsElement(); got != tc.want { + t.Errorf("NeedsElement() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestWindowFnInvoker_PanicOnUnregistered(t *testing.T) { + type unregisteredFn struct{} + defer func() { + if r := recover(); r == nil { + t.Error("NewWindowFnInvoker did not panic on unregistered type") + } + }() + NewWindowFnInvoker(&unregisteredFn{}) +} 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..840c842636a4 --- /dev/null +++ b/sdks/go/pkg/beam/core/graph/window/register.go @@ -0,0 +1,147 @@ +// 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" +) + +// windowFnMeta holds validated metadata about a registered custom WindowFn type. +type windowFnMeta struct { + // Type is the struct type (with pointer stripped), e.g. myWindowFn. + Type reflect.Type + // ElemType is the concrete element parameter type, or nil for + // timestamp-only signatures. + ElemType reflect.Type +} + +// NeedsElement reports whether this WindowFn signature accepts an element. +func (m *windowFnMeta) NeedsElement() bool { + return m.ElemType != nil +} + +var ( + windowFnRegistryMu sync.RWMutex + windowFnRegistry = map[reflect.Type]*windowFnMeta{} +) + +// LookupWindowFnMeta returns the validated metadata for a registered custom +// WindowFn type. The key is the struct type (pointer stripped). +// Returns nil if not registered. +func LookupWindowFnMeta(t reflect.Type) *windowFnMeta { + windowFnRegistryMu.RLock() + defer windowFnRegistryMu.RUnlock() + return windowFnRegistry[t] +} + +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 +// +// RegisterWindowFn panics if the type is invalid or already registered. +// +// 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)) + } + + meta := validateAssignWindows(t, m) + + windowFnRegistryMu.Lock() + defer windowFnRegistryMu.Unlock() + + if _, dup := windowFnRegistry[structType]; dup { + panic(fmt.Sprintf("window.RegisterWindowFn: %v is already registered", t)) + } + windowFnRegistry[structType] = meta + + runtime.RegisterType(reflect.TypeOf(v)) +} + +// validateAssignWindows checks that the method has a valid signature and +// returns the corresponding metadata. +func validateAssignWindows(ptrType reflect.Type, m reflect.Method) *windowFnMeta { + mt := m.Type + // Method type includes the receiver as first param. + // Valid shapes: + // (receiver, typex.EventTime) -> []typex.Window numIn=2 + // (receiver, typex.EventTime, elemType) -> []typex.Window numIn=3 + + if mt.NumOut() != 1 || mt.Out(0) != windowSliceType { + panic(fmt.Sprintf( + "window.RegisterWindowFn: %v.AssignWindows must return []typex.Window, got %v", + ptrType, mt)) + } + + switch mt.NumIn() { + case 2: + // (receiver, typex.EventTime) + if mt.In(1) != eventTimeType { + panic(fmt.Sprintf( + "window.RegisterWindowFn: %v.AssignWindows first param must be typex.EventTime, got %v", + ptrType, mt.In(1))) + } + return &windowFnMeta{Type: ptrType.Elem(), ElemType: nil} + + case 3: + // (receiver, typex.EventTime, elemType) + if mt.In(1) != eventTimeType { + panic(fmt.Sprintf( + "window.RegisterWindowFn: %v.AssignWindows first param must be typex.EventTime, got %v", + ptrType, mt.In(1))) + } + elemType := mt.In(2) + return &windowFnMeta{Type: ptrType.Elem(), ElemType: elemType} + + default: + panic(fmt.Sprintf( + "window.RegisterWindowFn: %v.AssignWindows must take (typex.EventTime) or (typex.EventTime, T), got %d params (excluding receiver)", + ptrType, mt.NumIn()-1)) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/exec/translate.go b/sdks/go/pkg/beam/core/runtime/exec/translate.go index 13b40ea0d1c6..09d6d29fbb8d 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 window.LookupWindowFnMeta(t) == nil { + 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()), nil 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)), nil 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)), nil + 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) + } + meta := window.LookupWindowFnMeta(t) + if meta == nil { + return nil, errors.Errorf("type %v is not registered via window.RegisterWindowFn", t) + } + if meta.NeedsElement() { + 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())), nil default: return nil, fmt.Errorf("unsupported window mapping fn URN %v", urn) } diff --git a/sdks/go/pkg/beam/core/runtime/exec/window.go b/sdks/go/pkg/beam/core/runtime/exec/window.go index fabd6af933d8..24d167e5cffd 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,16 +41,27 @@ func (w *WindowInto) ID() UnitID { } func (w *WindowInto) Up(ctx context.Context) error { + w.invoker = invokerFor(w.Fn) return nil } +// 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 { + if wfn.Kind != window.CustomWindows { + return nil + } + return window.NewWindowFnInvoker(wfn.CustomFn) +} + func (w *WindowInto) StartBundle(ctx context.Context, id string, data DataContext) error { return w.Out.StartBundle(ctx, id, data) } 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), Timestamp: elm.Timestamp, Elm: elm.Elm, Elm2: elm.Elm2, @@ -57,7 +70,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 any) []typex.Window { switch wfn.Kind { case window.GlobalWindows: return window.SingleGlobalWindow @@ -82,6 +97,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) + default: panic(fmt.Sprintf("Unexpected window fn: %v", wfn)) } @@ -170,10 +188,15 @@ type WindowMapper interface { type windowMapper struct { wfn *window.Fn + inv *window.WindowFnInvoker // non-nil for CustomWindows +} + +func newWindowMapper(wfn *window.Fn) *windowMapper { + return &windowMapper{wfn: wfn, inv: invokerFor(wfn)} } 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) if len(candidates) == 0 { return nil, fmt.Errorf("failed to map main input window to side input window with WindowFn %v", f.wfn.String()) } 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..33fabc431a62 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window_test.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window_test.go @@ -113,10 +113,32 @@ 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) + out := assignWindows(test.fn, invokerFor(test.fn), test.in, nil) if !window.IsEqualList(out, test.out) { t.Errorf("assignWindows(%v, %v) = %v, want %v", test.fn, test.in, out, test.out) } @@ -219,6 +241,93 @@ func TestMapWindows(t *testing.T) { } } +func init() { + window.RegisterWindowFn[*fixedCustomWindowFn]() + window.RegisterWindowFn[*elemSizedWindowFn]() +} + +// 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 := invokerFor(fn) + b.ReportAllocs() + for b.Loop() { + assignWindows(fn, inv, 1500, 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) + start := ts - (ts % size) + if ts < 0 { + // Go's % truncates toward zero, so for negative dividends + // ts%size is non-positive and ts-(ts%size) rounds toward + // zero instead of toward -inf. The double-mod expression + // computes the Euclidean (non-negative) remainder, giving + // a correct floor to the window boundary. + 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") 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..3adf628671c4 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go @@ -21,7 +21,9 @@ import ( "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" ) @@ -89,3 +91,36 @@ 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) + } + }) + } +} + +func TestWindowFnRoundTrip_CustomKind(t *testing.T) { + // Custom WindowFns are serialized via the Beam model proto + // (FunctionSpec), not the internal v1 proto. The v1 path only + // preserves the Kind so that EncodeMultiEdge does not fail. + 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) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/translate.go b/sdks/go/pkg/beam/core/runtime/graphx/translate.go index 3994397e7ba5..5e57690eca39 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" @@ -69,6 +74,8 @@ 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" + URNCustomWindowFn = "beam:go:windowfn:custom:v1" URNProgressReporting = "beam:protocol:progress_reporting:v1" URNMultiCore = "beam:protocol:multi_core_bundle_processing:v1" @@ -358,6 +365,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 +1473,30 @@ func makeWindowFn(w *window.Fn) (*pipepb.FunctionSpec, error) { }, ), }, nil + case window.CustomWindows: + t := reflect.TypeOf(w.CustomFn) + key, ok := runtime.TypeKey(reflectx.SkipPtr(t)) + if !ok { + return nil, errors.Errorf("custom WindowFn type %v is not registered", 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 +1506,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/test/integration/primitives/windowinto.go b/sdks/go/test/integration/primitives/windowinto.go index f5d01bdfbba5..29f6ec79c1be 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,13 @@ 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) } // createTimestampedData produces data timestamped with the ordinal. @@ -413,3 +422,130 @@ 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) + start := ts - (ts.Add(time.Duration(f.SizeMs)*time.Millisecond) % mtime.FromDuration(time.Duration(f.SizeMs)*time.Millisecond)) + end := start + size + return []typex.Window{window.IntervalWindow{Start: start, End: end}} +} + +// 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 +} + +func init() { + beam.RegisterType(reflect.TypeOf((*elemWithSize)(nil)).Elem()) +} + +// 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) + 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) +} From 30a31b9474518614285d6d832e957d6e3b03bac0 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:03 +0100 Subject: [PATCH 02/11] Skip custom window tests on Java-based runners A custom WindowFn is serialized with a Go-specific URN in the windowing strategy. WindowingStrategyTranslation.fromProto throws IllegalArgumentException for any URN it does not know, so every Java-based runner rejects the pipeline outright. This is not hypothetical: goPortablePreCommit runs ulrValidatesRunner, the Java universal local runner, so GoPortable fails on these tests however correct the Go side is. Skip the three tests on portable, flink, samza, spark and dataflow. Direct and prism are untouched and both pass, which states the scope of the feature plainly: Go-side runners only, until custom windows have a form other runners can read. --- sdks/go/test/integration/integration.go | 30 +++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/sdks/go/test/integration/integration.go b/sdks/go/test/integration/integration.go index 9d04f08de4fb..fba7e6dc137c 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 samzaFilters = []string{ @@ -251,6 +263,12 @@ var samzaFilters = []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{ @@ -289,6 +307,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{ @@ -331,6 +355,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 From 7ce2155d0f1c40c70e262a6b4ca57b872c9cf108 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:03 +0100 Subject: [PATCH 03/11] Require one window for custom side input mapping Side input mapping runs the side WindowFn at the main window's maximum timestamp and takes the last window it returns. That picks the earliest window, but only because the built-in sliding windows return theirs latest first. A custom WindowFn makes no such promise, so one returning its windows earliest first got the latest window instead, and the side input read the wrong data with no error. Java only offers this maximum-timestamp mapping from PartitioningWindowFn, which always assigns to exactly one window. WindowFn.getDefaultWindowMappingFn is otherwise abstract, and Sessions throws rather than inventing a mapping. Do the same here: if a custom WindowFn returns more than one window, report an error instead of guessing which one was meant. Rejecting element-aware WindowFns does not cover this. That rule is about having no element to pass to AssignWindows. TestMapWindowCustom covers both cases. Without the check the multi-window case returns [1000:2000) where the code means the earliest window. The test mappers go through newWindowMapper so a custom WindowFn reaches MapWindow with its invoker set. --- sdks/go/pkg/beam/core/runtime/exec/window.go | 8 +++ .../pkg/beam/core/runtime/exec/window_test.go | 58 ++++++++++++++++++- 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/sdks/go/pkg/beam/core/runtime/exec/window.go b/sdks/go/pkg/beam/core/runtime/exec/window.go index 24d167e5cffd..c74baa4a23d6 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window.go @@ -200,6 +200,14 @@ func (f *windowMapper) MapWindow(w typex.Window) (typex.Window, error) { 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, fmt.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 33fabc431a62..e48f9803874a 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window_test.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window_test.go @@ -184,7 +184,7 @@ func TestMapWindow(t *testing.T) { }, } for _, test := range tests { - mapper := &windowMapper{wfn: test.wfn} + mapper := newWindowMapper(test.wfn) outputWin, err := mapper.MapWindow(test.in) if err != nil { t.Fatalf("MapWindow for test %v failed, got %v", test.name, err) @@ -220,7 +220,7 @@ func TestMapWindows(t *testing.T) { inV, expected := makeNoncedWindowValues(tc.in, tc.expect) out := &CaptureNode{UID: 1} - unit := &MapWindows{UID: 2, Fn: &windowMapper{wfn: tc.wFn}, Out: out} + unit := &MapWindows{UID: 2, Fn: newWindowMapper(tc.wFn), Out: out} a := &FixedRoot{UID: 3, Elements: inV, Out: unit} p, err := NewPlan(tc.name, []Unit{a, unit, out}) @@ -244,6 +244,60 @@ func TestMapWindows(t *testing.T) { func init() { window.RegisterWindowFn[*fixedCustomWindowFn]() window.RegisterWindowFn[*elemSizedWindowFn]() + window.RegisterWindowFn[*multiWindowFn]() +} + +// 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 := newWindowMapper(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. From 97548d67019ae9d77fd9132bd215b0b9f27d555b Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:16 +0100 Subject: [PATCH 04/11] Accept a KV element in custom AssignWindows An element-aware AssignWindows received only FullValue.Elm, so windowing a KV PCollection quietly passed the key and dropped the value. There is no single-value KV type to hand over instead: typex.KV is an empty marker used to build type trees, not a carrier. The DoFn invoker already answers this. exec/fn.go binds Elm to the next parameter and, when Elm2 is set, binds it to the one after, so every DoFn takes a KV as two parameters. AssignWindows now works the same way, with a third accepted shape: AssignWindows(ts typex.EventTime, k K, v V) []typex.Window The registry holds the list of element parameter types, which carries the types and the count together. That also replaces LookupWindowFnMeta, which was exported but returned an unexported struct, so callers in other packages could call its methods without being able to name the result. Its Type field was set on both construction paths and never read, the struct type already being the map key. Registration had no tests, despite being the contract that is supposed to reject misuse early. TestRegisterWindowFn covers the three accepted shapes and the five rejected ones, and TestWindowIntoKV drives a KV element through WindowInto to check the value arrives. Whether a PCollection's shape matches the registered arity is still unchecked. --- sdks/go/pkg/beam/core/graph/window/fn.go | 8 +- sdks/go/pkg/beam/core/graph/window/invoke.go | 93 +++++++----- .../pkg/beam/core/graph/window/invoke_test.go | 79 ++++++++++- .../go/pkg/beam/core/graph/window/register.go | 86 +++++------ .../beam/core/graph/window/register_test.go | 133 ++++++++++++++++++ .../pkg/beam/core/runtime/exec/translate.go | 8 +- sdks/go/pkg/beam/core/runtime/exec/window.go | 8 +- .../pkg/beam/core/runtime/exec/window_test.go | 61 +++++++- 8 files changed, 372 insertions(+), 104 deletions(-) create mode 100644 sdks/go/pkg/beam/core/graph/window/register_test.go diff --git a/sdks/go/pkg/beam/core/graph/window/fn.go b/sdks/go/pkg/beam/core/graph/window/fn.go index a1b890e214de..ad8a75b38775 100644 --- a/sdks/go/pkg/beam/core/graph/window/fn.go +++ b/sdks/go/pkg/beam/core/graph/window/fn.go @@ -71,7 +71,7 @@ func NewCustom(fn any) *Fn { if t.Kind() == reflect.Pointer { st = t.Elem() } - if LookupWindowFnMeta(st) == nil { + 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} @@ -98,10 +98,8 @@ func (w *Fn) NeedsElement() bool { if t.Kind() == reflect.Pointer { t = t.Elem() } - if meta := LookupWindowFnMeta(t); meta != nil { - return meta.NeedsElement() - } - return false + elems, _ := LookupWindowFn(t) + return len(elems) > 0 } // TODO(herohde) 4/17/2018: do we need to expose the window type as well? diff --git a/sdks/go/pkg/beam/core/graph/window/invoke.go b/sdks/go/pkg/beam/core/graph/window/invoke.go index 95bacc2c288a..d8dcd4d3eceb 100644 --- a/sdks/go/pkg/beam/core/graph/window/invoke.go +++ b/sdks/go/pkg/beam/core/graph/window/invoke.go @@ -27,26 +27,34 @@ 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 (interface{}). +// 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 } -// WindowFnInvoker wraps a custom WindowFn instance and dispatches -// AssignWindows calls through one of three paths, chosen once at -// construction: +// 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. // -// 1. Type-assert to tsOnlyAssigner — zero allocation. -// 2. Type-assert to anyElemAssigner — zero allocation. -// 3. reflect.Value.Method.Call — 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(typex.EventTime, any) []typex.Window + call func(ts typex.EventTime, elm, elm2 any) []typex.Window needsElement bool + isKV bool } -// NewWindowFnInvoker builds an invoker for fn. The concrete type of fn -// must have been previously registered via RegisterWindowFn. +// NewWindowFnInvoker builds an invoker for fn. The concrete type of fn must +// have been previously registered via RegisterWindowFn. // Panics if fn's type is not registered. func NewWindowFnInvoker(fn any) *WindowFnInvoker { t := reflect.TypeOf(fn) @@ -58,60 +66,77 @@ func NewWindowFnInvoker(fn any) *WindowFnInvoker { structType = t.Elem() } - meta := LookupWindowFnMeta(structType) - if meta == nil { + elems, ok := LookupWindowFn(structType) + if !ok { panic(fmt.Sprintf("window.NewWindowFnInvoker: type %v is not registered; call window.RegisterWindowFn during init()", t)) } - inv := &WindowFnInvoker{needsElement: meta.NeedsElement()} + inv := &WindowFnInvoker{needsElement: len(elems) > 0, isKV: len(elems) > 1} - if !meta.NeedsElement() { - // Fast path 1: timestamp-only. + switch len(elems) { + case 0: if a, ok := fn.(tsOnlyAssigner); ok { - inv.call = func(ts typex.EventTime, _ any) []typex.Window { + inv.call = func(ts typex.EventTime, _, _ any) []typex.Window { return a.AssignWindows(ts) } return inv } - // Unreachable for well-typed registrations, but fall through to reflect. - } else { - // Fast path 2: element typed as any. + case 1: if a, ok := fn.(anyElemAssigner); ok { - inv.call = func(ts typex.EventTime, elem any) []typex.Window { - return a.AssignWindows(ts, elem) + inv.call = func(ts typex.EventTime, elm, _ any) []typex.Window { + return a.AssignWindows(ts, elm) + } + return inv + } + 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 } } - // Slow path 3: concrete element type — use reflect. - rv := reflect.ValueOf(fn) - m := rv.MethodByName("AssignWindows") + // Concrete element types cannot be reached through an interface assertion. + m := reflect.ValueOf(fn).MethodByName("AssignWindows") if !m.IsValid() { panic(fmt.Sprintf("window.NewWindowFnInvoker: %v has no AssignWindows method", t)) } - if !meta.NeedsElement() { - inv.call = func(ts typex.EventTime, _ any) []typex.Window { + 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) } - } else { - inv.call = func(ts typex.EventTime, elem any) []typex.Window { - out := m.Call([]reflect.Value{reflect.ValueOf(ts), reflect.ValueOf(elem)}) + 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 } -// Invoke calls AssignWindows on the underlying WindowFn. -// If the WindowFn is timestamp-only, elem is ignored. -func (inv *WindowFnInvoker) Invoke(ts typex.EventTime, elem any) []typex.Window { - return inv.call(ts, elem) +// 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) } // NeedsElement reports whether the underlying WindowFn accepts an element. func (inv *WindowFnInvoker) NeedsElement() bool { return inv.needsElement } + +// IsKV reports whether the underlying WindowFn takes a KV element as a +// separate key and value. +func (inv *WindowFnInvoker) IsKV() bool { + return inv.isKV +} diff --git a/sdks/go/pkg/beam/core/graph/window/invoke_test.go b/sdks/go/pkg/beam/core/graph/window/invoke_test.go index f70f4750ca27..3376c65dec2b 100644 --- a/sdks/go/pkg/beam/core/graph/window/invoke_test.go +++ b/sdks/go/pkg/beam/core/graph/window/invoke_test.go @@ -46,9 +46,80 @@ func (f *elemAwareConcreteWindowFn) AssignWindows(ts typex.EventTime, elem int64 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 := NewWindowFnInvoker(tc.fn) + if !inv.NeedsElement() { + t.Error("NeedsElement() = false, want true") + } + if !inv.IsKV() { + t.Error("IsKV() = false, want true") + } + + 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_IsKV(t *testing.T) { + tests := []struct { + name string + fn any + want bool + }{ + {"timestamp-only", &testWindowFn{BucketSize: 1000}, false}, + {"single element", &elemAwareConcreteWindowFn{DefaultSizeMs: 1000}, false}, + {"kv element", &kvConcreteWindowFn{}, true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := NewWindowFnInvoker(tc.fn).IsKV(); got != tc.want { + t.Errorf("IsKV() = %v, want %v", got, tc.want) + } + }) + } } func TestWindowFnInvoker_TimestampOnly(t *testing.T) { @@ -59,7 +130,7 @@ func TestWindowFnInvoker_TimestampOnly(t *testing.T) { t.Fatal("NeedsElement() = true, want false") } - windows := inv.Invoke(1500, nil) + windows := inv.Invoke(1500, nil, nil) if len(windows) != 1 { t.Fatalf("got %d windows, want 1", len(windows)) } @@ -77,7 +148,7 @@ func TestWindowFnInvoker_AnyElem(t *testing.T) { t.Fatal("NeedsElement() = false, want true") } - windows := inv.Invoke(7500, "ignored") + windows := inv.Invoke(7500, "ignored", nil) if len(windows) != 1 { t.Fatalf("got %d windows, want 1", len(windows)) } @@ -96,7 +167,7 @@ func TestWindowFnInvoker_ConcreteElem(t *testing.T) { } // Element provides window size of 5000ms. - windows := inv.Invoke(7500, int64(5000)) + windows := inv.Invoke(7500, int64(5000), nil) if len(windows) != 1 { t.Fatalf("got %d windows, want 1", len(windows)) } @@ -106,7 +177,7 @@ func TestWindowFnInvoker_ConcreteElem(t *testing.T) { } // Element <= 0: falls back to default. - windows = inv.Invoke(1500, int64(0)) + windows = inv.Invoke(1500, int64(0), nil) if len(windows) != 1 { t.Fatalf("got %d windows, want 1", len(windows)) } diff --git a/sdks/go/pkg/beam/core/graph/window/register.go b/sdks/go/pkg/beam/core/graph/window/register.go index 840c842636a4..43ecc60189d2 100644 --- a/sdks/go/pkg/beam/core/graph/window/register.go +++ b/sdks/go/pkg/beam/core/graph/window/register.go @@ -24,32 +24,24 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" ) -// windowFnMeta holds validated metadata about a registered custom WindowFn type. -type windowFnMeta struct { - // Type is the struct type (with pointer stripped), e.g. myWindowFn. - Type reflect.Type - // ElemType is the concrete element parameter type, or nil for - // timestamp-only signatures. - ElemType reflect.Type -} - -// NeedsElement reports whether this WindowFn signature accepts an element. -func (m *windowFnMeta) NeedsElement() bool { - return m.ElemType != nil -} - +// 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]*windowFnMeta{} + windowFnRegistry = map[reflect.Type][]reflect.Type{} ) -// LookupWindowFnMeta returns the validated metadata for a registered custom -// WindowFn type. The key is the struct type (pointer stripped). -// Returns nil if not registered. -func LookupWindowFnMeta(t reflect.Type) *windowFnMeta { +// 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() - return windowFnRegistry[t] + elems, ok = windowFnRegistry[t] + return elems, ok } var ( @@ -66,6 +58,10 @@ var ( // // 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 or already registered. // @@ -91,7 +87,7 @@ func RegisterWindowFn[T any]() { panic(fmt.Sprintf("window.RegisterWindowFn: %v has no AssignWindows method", t)) } - meta := validateAssignWindows(t, m) + elems := validateAssignWindows(t, m) windowFnRegistryMu.Lock() defer windowFnRegistryMu.Unlock() @@ -99,49 +95,37 @@ func RegisterWindowFn[T any]() { if _, dup := windowFnRegistry[structType]; dup { panic(fmt.Sprintf("window.RegisterWindowFn: %v is already registered", t)) } - windowFnRegistry[structType] = meta + windowFnRegistry[structType] = elems runtime.RegisterType(reflect.TypeOf(v)) } // validateAssignWindows checks that the method has a valid signature and -// returns the corresponding metadata. -func validateAssignWindows(ptrType reflect.Type, m reflect.Method) *windowFnMeta { +// 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 - // Method type includes the receiver as first param. - // Valid shapes: - // (receiver, typex.EventTime) -> []typex.Window numIn=2 - // (receiver, typex.EventTime, elemType) -> []typex.Window numIn=3 if mt.NumOut() != 1 || mt.Out(0) != windowSliceType { panic(fmt.Sprintf( "window.RegisterWindowFn: %v.AssignWindows must return []typex.Window, got %v", ptrType, mt)) } - - switch mt.NumIn() { - case 2: - // (receiver, typex.EventTime) - if mt.In(1) != eventTimeType { - panic(fmt.Sprintf( - "window.RegisterWindowFn: %v.AssignWindows first param must be typex.EventTime, got %v", - ptrType, mt.In(1))) - } - return &windowFnMeta{Type: ptrType.Elem(), ElemType: nil} - - case 3: - // (receiver, typex.EventTime, elemType) - if mt.In(1) != eventTimeType { - panic(fmt.Sprintf( - "window.RegisterWindowFn: %v.AssignWindows first param must be typex.EventTime, got %v", - ptrType, mt.In(1))) - } - elemType := mt.In(2) - return &windowFnMeta{Type: ptrType.Elem(), ElemType: elemType} - - default: + if mt.NumIn() < 2 || mt.NumIn() > 4 { panic(fmt.Sprintf( - "window.RegisterWindowFn: %v.AssignWindows must take (typex.EventTime) or (typex.EventTime, T), got %d params (excluding receiver)", + "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..ea275255c3a6 --- /dev/null +++ b/sdks/go/pkg/beam/core/graph/window/register_test.go @@ -0,0 +1,133 @@ +// 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, + }, + { + name: "already registered", + register: RegisterWindowFn[*regTsOnly], + wantPanic: true, + }, + } + 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/translate.go b/sdks/go/pkg/beam/core/runtime/exec/translate.go index 09d6d29fbb8d..87ce15432776 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/translate.go +++ b/sdks/go/pkg/beam/core/runtime/exec/translate.go @@ -291,7 +291,7 @@ func unmarshalWindowFn(wfn *pipepb.FunctionSpec) (*window.Fn, error) { if !ok { return nil, errors.Errorf("custom WindowFn type key %q not found in registry", envelope.Type) } - if window.LookupWindowFnMeta(t) == nil { + if _, ok := window.LookupWindowFn(t); !ok { return nil, errors.Errorf("type %v is not registered via window.RegisterWindowFn", t) } val := reflect.New(t) @@ -349,11 +349,11 @@ func unmarshalAndMakeWindowMapping(wmfn *pipepb.FunctionSpec) (WindowMapper, err if !ok { return nil, errors.Errorf("custom WindowFn type key %q not found in registry", envelope.Type) } - meta := window.LookupWindowFnMeta(t) - if meta == nil { + elems, ok := window.LookupWindowFn(t) + if !ok { return nil, errors.Errorf("type %v is not registered via window.RegisterWindowFn", t) } - if meta.NeedsElement() { + 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) diff --git a/sdks/go/pkg/beam/core/runtime/exec/window.go b/sdks/go/pkg/beam/core/runtime/exec/window.go index c74baa4a23d6..042f8226b8df 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window.go @@ -61,7 +61,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, w.invoker, elm.Timestamp, elm.Elm), + Windows: assignWindows(w.Fn, w.invoker, elm.Timestamp, elm.Elm, elm.Elm2), Timestamp: elm.Timestamp, Elm: elm.Elm, Elm2: elm.Elm2, @@ -72,7 +72,7 @@ func (w *WindowInto) ProcessElement(ctx context.Context, elm *FullValue, values // 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 any) []typex.Window { +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 @@ -98,7 +98,7 @@ func assignWindows(wfn *window.Fn, inv *window.WindowFnInvoker, ts typex.EventTi return []typex.Window{window.IntervalWindow{Start: ts, End: ts.Add(wfn.Gap)}} case window.CustomWindows: - return inv.Invoke(ts, elm) + return inv.Invoke(ts, elm, elm2) default: panic(fmt.Sprintf("Unexpected window fn: %v", wfn)) @@ -196,7 +196,7 @@ func newWindowMapper(wfn *window.Fn) *windowMapper { } func (f *windowMapper) MapWindow(w typex.Window) (typex.Window, error) { - candidates := assignWindows(f.wfn, f.inv, w.MaxTimestamp(), nil) + 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()) } 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 e48f9803874a..94915b692307 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window_test.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window_test.go @@ -138,7 +138,7 @@ func TestAssignWindow(t *testing.T) { } for _, test := range tests { - out := assignWindows(test.fn, invokerFor(test.fn), test.in, nil) + out := assignWindows(test.fn, invokerFor(test.fn), test.in, nil, nil) if !window.IsEqualList(out, test.out) { t.Errorf("assignWindows(%v, %v) = %v, want %v", test.fn, test.in, out, test.out) } @@ -245,6 +245,63 @@ 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. @@ -314,7 +371,7 @@ func BenchmarkAssignWindowsCustom(b *testing.B) { inv := invokerFor(fn) b.ReportAllocs() for b.Loop() { - assignWindows(fn, inv, 1500, nil) + assignWindows(fn, inv, 1500, nil, nil) } } From 84580459fbd0ccf5377ab47a7f45109d05869579 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:16 +0100 Subject: [PATCH 05/11] Check the windowed element against AssignWindows Nothing compared a custom WindowFn's element parameters against the PCollection being windowed, so a mismatch showed up as a reflect panic from inside a bundle, wrapped by callNoPanic, with nothing pointing back to the WindowInto call. TryWindowInto now rejects, when the graph is built, an element-aware WindowFn whose parameters cannot take the PCollection's element: a wrong element type, a parameter count that disagrees with the PCollection shape either way, and a composite other than KV such as CoGBK. Timestamp-only WindowFns and the built-in kinds accept any shape. Universal components are skipped rather than rejected, since they are not bound yet at this point. The check uses assignability rather than the wider conversion the DoFn invoker does through ConvertFn, which lives in exec and is not reachable here. The difference only shows up on concrete mismatches, which are the ones worth refusing. --- sdks/go/pkg/beam/windowing.go | 43 +++++++++++++ sdks/go/pkg/beam/windowing_test.go | 100 +++++++++++++++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 sdks/go/pkg/beam/windowing_test.go 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") + } +} From 592f6d48bdaf43ffaecb817425218dd6134c1921 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:30 +0100 Subject: [PATCH 06/11] Firm up the custom WindowFn invoker and registry Two rough edges in the new API. NewWindowFnInvoker panicked on a nil or unregistered WindowFn, and WindowInto.Up called it while returning error. Combine.Up and ParDo.Up both report construction failures through that return. The panic is not reachable today, since Up gets its Fn from unmarshalWindowFn which checks the registry and builds through NewCustom, and both DecodeMultiEdge callers throw away the v1-decoded WindowFn whose CustomFn is nil. It is still the wrong contract. The nearest precedent, newInvoker for DoFns, is infallible because *funcx.Fn can only be obtained by passing validation; NewWindowFnInvoker takes any and has to look in the registry, and window.Fn has exported fields so one can be built with no check at all. RegisterWindowFn panicked on any duplicate, stricter than either sibling registry. runtime.RegisterFunction overwrites and says so in a comment, and runtime.RegisterType panics only when a different type collides on the key. RegisterWindowFn already calls RegisterType, so the first call filled both registries and a second panicked in the window one before reaching the tolerant check. That broke real cases: the same type registered from a library's init and a test's init aborts the process at startup. The element types are derived from the type's AssignWindows method, so a duplicate key cannot carry different data and an overwrite is enough. --- sdks/go/pkg/beam/core/graph/window/invoke.go | 19 +++++----- .../pkg/beam/core/graph/window/invoke_test.go | 38 ++++++++++++------- .../go/pkg/beam/core/graph/window/register.go | 8 ++-- .../beam/core/graph/window/register_test.go | 9 +++-- .../pkg/beam/core/runtime/exec/translate.go | 8 ++-- sdks/go/pkg/beam/core/runtime/exec/window.go | 17 ++++++--- .../pkg/beam/core/runtime/exec/window_test.go | 31 ++++++++++++--- 7 files changed, 85 insertions(+), 45 deletions(-) diff --git a/sdks/go/pkg/beam/core/graph/window/invoke.go b/sdks/go/pkg/beam/core/graph/window/invoke.go index d8dcd4d3eceb..663fff9a2aba 100644 --- a/sdks/go/pkg/beam/core/graph/window/invoke.go +++ b/sdks/go/pkg/beam/core/graph/window/invoke.go @@ -16,10 +16,10 @@ package window import ( - "fmt" "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. @@ -55,11 +55,10 @@ type WindowFnInvoker struct { // NewWindowFnInvoker builds an invoker for fn. The concrete type of fn must // have been previously registered via RegisterWindowFn. -// Panics if fn's type is not registered. -func NewWindowFnInvoker(fn any) *WindowFnInvoker { +func NewWindowFnInvoker(fn any) (*WindowFnInvoker, error) { t := reflect.TypeOf(fn) if t == nil { - panic("window.NewWindowFnInvoker: fn must not be nil") + return nil, errors.New("window.NewWindowFnInvoker: fn must not be nil") } structType := t if t.Kind() == reflect.Pointer { @@ -68,7 +67,7 @@ func NewWindowFnInvoker(fn any) *WindowFnInvoker { elems, ok := LookupWindowFn(structType) if !ok { - panic(fmt.Sprintf("window.NewWindowFnInvoker: type %v is not registered; call window.RegisterWindowFn during init()", t)) + return nil, errors.Errorf("window.NewWindowFnInvoker: type %v is not registered; call window.RegisterWindowFn during init()", t) } inv := &WindowFnInvoker{needsElement: len(elems) > 0, isKV: len(elems) > 1} @@ -79,28 +78,28 @@ func NewWindowFnInvoker(fn any) *WindowFnInvoker { inv.call = func(ts typex.EventTime, _, _ any) []typex.Window { return a.AssignWindows(ts) } - return inv + 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 + 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 + return inv, nil } } // Concrete element types cannot be reached through an interface assertion. m := reflect.ValueOf(fn).MethodByName("AssignWindows") if !m.IsValid() { - panic(fmt.Sprintf("window.NewWindowFnInvoker: %v has no AssignWindows method", t)) + return nil, errors.Errorf("window.NewWindowFnInvoker: %v has no AssignWindows method", t) } switch len(elems) { @@ -120,7 +119,7 @@ func NewWindowFnInvoker(fn any) *WindowFnInvoker { return out[0].Interface().([]typex.Window) } } - return inv + return inv, nil } // Invoke calls AssignWindows on the underlying WindowFn. Signatures that do not diff --git a/sdks/go/pkg/beam/core/graph/window/invoke_test.go b/sdks/go/pkg/beam/core/graph/window/invoke_test.go index 3376c65dec2b..6dddfef93773 100644 --- a/sdks/go/pkg/beam/core/graph/window/invoke_test.go +++ b/sdks/go/pkg/beam/core/graph/window/invoke_test.go @@ -83,7 +83,7 @@ func TestWindowFnInvoker_KV(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - inv := NewWindowFnInvoker(tc.fn) + inv := mustInvoker(t, tc.fn) if !inv.NeedsElement() { t.Error("NeedsElement() = false, want true") } @@ -115,7 +115,7 @@ func TestWindowFnInvoker_IsKV(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - if got := NewWindowFnInvoker(tc.fn).IsKV(); got != tc.want { + if got := mustInvoker(t, tc.fn).IsKV(); got != tc.want { t.Errorf("IsKV() = %v, want %v", got, tc.want) } }) @@ -124,7 +124,7 @@ func TestWindowFnInvoker_IsKV(t *testing.T) { func TestWindowFnInvoker_TimestampOnly(t *testing.T) { fn := &testWindowFn{BucketSize: 3000} - inv := NewWindowFnInvoker(fn) + inv := mustInvoker(t, fn) if inv.NeedsElement() { t.Fatal("NeedsElement() = true, want false") @@ -142,7 +142,7 @@ func TestWindowFnInvoker_TimestampOnly(t *testing.T) { func TestWindowFnInvoker_AnyElem(t *testing.T) { fn := &elemAwareAnyWindowFn{SizeMs: 5000} - inv := NewWindowFnInvoker(fn) + inv := mustInvoker(t, fn) if !inv.NeedsElement() { t.Fatal("NeedsElement() = false, want true") @@ -160,7 +160,7 @@ func TestWindowFnInvoker_AnyElem(t *testing.T) { func TestWindowFnInvoker_ConcreteElem(t *testing.T) { fn := &elemAwareConcreteWindowFn{DefaultSizeMs: 1000} - inv := NewWindowFnInvoker(fn) + inv := mustInvoker(t, fn) if !inv.NeedsElement() { t.Fatal("NeedsElement() = false, want true") @@ -199,7 +199,7 @@ func TestWindowFnInvoker_NeedsElementCorrectness(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - inv := NewWindowFnInvoker(tc.fn) + inv := mustInvoker(t, tc.fn) if got := inv.NeedsElement(); got != tc.want { t.Errorf("NeedsElement() = %v, want %v", got, tc.want) } @@ -207,12 +207,24 @@ func TestWindowFnInvoker_NeedsElementCorrectness(t *testing.T) { } } -func TestWindowFnInvoker_PanicOnUnregistered(t *testing.T) { +func TestWindowFnInvoker_ErrorOnUnregistered(t *testing.T) { type unregisteredFn struct{} - defer func() { - if r := recover(); r == nil { - t.Error("NewWindowFnInvoker did not panic on unregistered type") - } - }() - NewWindowFnInvoker(&unregisteredFn{}) + 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 index 43ecc60189d2..f599b5b83c2a 100644 --- a/sdks/go/pkg/beam/core/graph/window/register.go +++ b/sdks/go/pkg/beam/core/graph/window/register.go @@ -63,7 +63,8 @@ var ( // 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 or already registered. +// RegisterWindowFn panics if the type is invalid. Registering the same type +// more than once is allowed. // // Example: // @@ -92,9 +93,8 @@ func RegisterWindowFn[T any]() { windowFnRegistryMu.Lock() defer windowFnRegistryMu.Unlock() - if _, dup := windowFnRegistry[structType]; dup { - panic(fmt.Sprintf("window.RegisterWindowFn: %v is already registered", t)) - } + // 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)) diff --git a/sdks/go/pkg/beam/core/graph/window/register_test.go b/sdks/go/pkg/beam/core/graph/window/register_test.go index ea275255c3a6..67645e6f1178 100644 --- a/sdks/go/pkg/beam/core/graph/window/register_test.go +++ b/sdks/go/pkg/beam/core/graph/window/register_test.go @@ -102,9 +102,12 @@ func TestRegisterWindowFn(t *testing.T) { wantPanic: true, }, { - name: "already registered", - 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 { diff --git a/sdks/go/pkg/beam/core/runtime/exec/translate.go b/sdks/go/pkg/beam/core/runtime/exec/translate.go index 87ce15432776..22f0eaecaa92 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/translate.go +++ b/sdks/go/pkg/beam/core/runtime/exec/translate.go @@ -308,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 newWindowMapper(window.NewGlobalWindows()), nil + return newWindowMapper(window.NewGlobalWindows()) case graphx.URNWindowMappingFixed: var payload pipepb.FixedWindowsPayload if err := proto.Unmarshal(wmfn.GetPayload(), &payload); err != nil { @@ -319,7 +319,7 @@ func unmarshalAndMakeWindowMapping(wmfn *pipepb.FunctionSpec) (WindowMapper, err return nil, err } size := sizePB.AsDuration() - return newWindowMapper(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 { @@ -336,7 +336,7 @@ func unmarshalAndMakeWindowMapping(wmfn *pipepb.FunctionSpec) (WindowMapper, err return nil, err } size := sizePB.AsDuration() - return newWindowMapper(window.NewSlidingWindows(period, size)), nil + return newWindowMapper(window.NewSlidingWindows(period, size)) case graphx.URNWindowMappingCustom: var envelope struct { Type string `json:"type"` @@ -360,7 +360,7 @@ func unmarshalAndMakeWindowMapping(wmfn *pipepb.FunctionSpec) (WindowMapper, err 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())), nil + 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/window.go b/sdks/go/pkg/beam/core/runtime/exec/window.go index 042f8226b8df..bfc9b2ce1f6f 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window.go @@ -41,16 +41,17 @@ func (w *WindowInto) ID() UnitID { } func (w *WindowInto) Up(ctx context.Context) error { - w.invoker = invokerFor(w.Fn) - 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 { +func invokerFor(wfn *window.Fn) (*window.WindowFnInvoker, error) { if wfn.Kind != window.CustomWindows { - return nil + return nil, nil } return window.NewWindowFnInvoker(wfn.CustomFn) } @@ -191,8 +192,12 @@ type windowMapper struct { inv *window.WindowFnInvoker // non-nil for CustomWindows } -func newWindowMapper(wfn *window.Fn) *windowMapper { - return &windowMapper{wfn: wfn, inv: invokerFor(wfn)} +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) { 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 94915b692307..9211e4505f24 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window_test.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window_test.go @@ -138,7 +138,7 @@ func TestAssignWindow(t *testing.T) { } for _, test := range tests { - out := assignWindows(test.fn, invokerFor(test.fn), test.in, nil, nil) + out := assignWindows(test.fn, mustInvokerFor(t, test.fn), test.in, nil, nil) if !window.IsEqualList(out, test.out) { t.Errorf("assignWindows(%v, %v) = %v, want %v", test.fn, test.in, out, test.out) } @@ -184,7 +184,7 @@ func TestMapWindow(t *testing.T) { }, } for _, test := range tests { - mapper := newWindowMapper(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) @@ -220,7 +220,7 @@ func TestMapWindows(t *testing.T) { inV, expected := makeNoncedWindowValues(tc.in, tc.expect) out := &CaptureNode{UID: 1} - unit := &MapWindows{UID: 2, Fn: newWindowMapper(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}) @@ -340,7 +340,7 @@ func TestMapWindowCustom(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - got, err := newWindowMapper(tc.wfn).MapWindow(tc.in) + 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) @@ -368,7 +368,10 @@ func (f *elemSizedWindowFn) AssignWindows(ts typex.EventTime, elem int64) []type func BenchmarkAssignWindowsCustom(b *testing.B) { fn := window.NewCustom(&fixedCustomWindowFn{SizeMs: 3000}) - inv := invokerFor(fn) + inv, err := invokerFor(fn) + if err != nil { + b.Fatalf("invokerFor(%v) failed: %v", fn, err) + } b.ReportAllocs() for b.Loop() { assignWindows(fn, inv, 1500, nil, nil) @@ -453,3 +456,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 +} From d4911896992f2e55c6bb0f9c86c118134550ed45 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:30 +0100 Subject: [PATCH 07/11] Move the custom WindowFn URN into beam:window_fn The URN goes into WindowingStrategy.window_fn but was named beam:go:windowfn:custom:v1, inventing a namespace. Every other SDK puts its private custom WindowFn in the shared beam:window_fn namespace for that field: Java uses beam:window_fn:serialized_java:v1 and Python uses beam:window_fn:pickled_python:v1. The Go built-ins right above it already use beam:window_fn too. Rename to beam:window_fn:serialized_go:v1 and group the constant with the other window_fn URNs instead of leaving it beside the window mapping URNs, which fill a different field. A comment records that other runners cannot read it, next to the name rather than buried in the marshalling code. Nothing has shipped with the old URN, so no compatibility shim is needed. URNWindowMappingCustom stays as it is. It fills SideInput.window_mapping_fn next to the three Go-private mapping URNs and already matches their naming. Both constants are needed: they describe different things in different proto fields, the same way Java and Python each carry a pair. --- sdks/go/pkg/beam/core/runtime/graphx/translate.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/sdks/go/pkg/beam/core/runtime/graphx/translate.go b/sdks/go/pkg/beam/core/runtime/graphx/translate.go index 5e57690eca39..70bc9bd40083 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/translate.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/translate.go @@ -63,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" @@ -75,7 +79,6 @@ const ( URNWindowMappingFixed = "beam:go:windowmapping:fixed:v1" URNWindowMappingSliding = "beam:go:windowmapping:sliding:v1" URNWindowMappingCustom = "beam:go:windowmapping:custom:v1" - URNCustomWindowFn = "beam:go:windowfn:custom:v1" URNProgressReporting = "beam:protocol:progress_reporting:v1" URNMultiCore = "beam:protocol:multi_core_bundle_processing:v1" From 1f7cd589b1701041172d0790774867c8ce68ff6e Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:41 +0100 Subject: [PATCH 08/11] Test the custom WindowFn wire format The FunctionSpec a custom WindowFn is written to, and read back from, had no unit tests. Integration tests cover the happy path through the direct runner, but the four error branches in unmarshalWindowFn had nothing, and TestWindowFnRoundTrip_CustomKind tested the internal v1 path that custom WindowFns deliberately do not use, as its own comment explained. Rename that test to say what it checks. Pin both halves instead: what makeWindowFn writes, and what unmarshalWindowFn accepts, including a malformed envelope, an unknown type key, a type in the type registry that was never registered as a WindowFn, and a payload that does not decode. Writing the test turned up a check that did not do what its error said. makeWindowFn tested runtime.TypeKey, which only reports whether a type has a package path and a name, and then said "is not registered" when it did not. Any locally declared type passes that, so an unregistered WindowFn was written out with a type key the harness cannot resolve, failing at run time instead of at construction. It now asks the WindowFn registry, and reports a missing type key separately. Reaching it takes building window.Fn by hand, since NewCustom checks registration, so this was latent rather than live. Fn.String printed only the type of the custom WindowFn, so two WindowFns of the same type with different settings looked identical. A failing assertion read "got CUS[*pkg.myFn], want CUS[*pkg.myFn]". Print the value too, which also helps the side input mapping error and any log line. --- sdks/go/pkg/beam/core/graph/window/fn.go | 2 +- .../beam/core/runtime/exec/translate_test.go | 87 +++++++++++++++++++ .../core/runtime/graphx/serialize_test.go | 70 ++++++++++++++- .../pkg/beam/core/runtime/graphx/translate.go | 8 +- 4 files changed, 160 insertions(+), 7 deletions(-) diff --git a/sdks/go/pkg/beam/core/graph/window/fn.go b/sdks/go/pkg/beam/core/graph/window/fn.go index ad8a75b38775..a212637919ba 100644 --- a/sdks/go/pkg/beam/core/graph/window/fn.go +++ b/sdks/go/pkg/beam/core/graph/window/fn.go @@ -123,7 +123,7 @@ func (w *Fn) String() string { case Sessions: return fmt.Sprintf("%v[%v]", w.Kind, w.Gap) case CustomWindows: - return fmt.Sprintf("%v[%v]", w.Kind, reflect.TypeOf(w.CustomFn)) + return fmt.Sprintf("%v[%T %+v]", w.Kind, w.CustomFn, w.CustomFn) default: return string(w.Kind) } 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/graphx/serialize_test.go b/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go index 3adf628671c4..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,6 +18,7 @@ package graphx import ( + "encoding/json" "reflect" "strings" "testing" @@ -26,6 +27,7 @@ import ( "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) { @@ -113,10 +115,10 @@ func TestWindowFnRoundTrip(t *testing.T) { } } -func TestWindowFnRoundTrip_CustomKind(t *testing.T) { - // Custom WindowFns are serialized via the Beam model proto - // (FunctionSpec), not the internal v1 proto. The v1 path only - // preserves the Kind so that EncodeMultiEdge does not fail. +// 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) @@ -124,3 +126,63 @@ func TestWindowFnRoundTrip_CustomKind(t *testing.T) { 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 70bc9bd40083..ef626513e3cd 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/translate.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/translate.go @@ -1478,9 +1478,13 @@ func makeWindowFn(w *window.Fn) (*pipepb.FunctionSpec, error) { }, nil case window.CustomWindows: t := reflect.TypeOf(w.CustomFn) - key, ok := runtime.TypeKey(reflectx.SkipPtr(t)) + 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 is not registered", t) + return nil, errors.Errorf("custom WindowFn type %v has no type key", t) } structPayload, err := jsonx.Marshal(w.CustomFn) if err != nil { From 46731e9356bb23384c297371824185f1b0622db1 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:58 +0100 Subject: [PATCH 09/11] Tidy the custom window code and note it in CHANGES Cleanups with no behaviour to argue about. The example and test WindowFns floored timestamps three different ways. Some used the Euclidean remainder, some a bare ts%size that rounds toward zero and so picks the wrong window for negative timestamps, and customFixedWindowFn wrote ts - (ts.Add(size) % mtime.FromDuration(size)) which reduces to ts-(ts%size) for non-negative ts and works by accident. Use one correct form everywhere, since these are the snippets someone writing a custom WindowFn will copy. fixedCustomWindowFn also drops a branch and a paragraph of comment for the same one-line expression. NeedsElement and IsKV on WindowFnInvoker had no caller outside tests, and the fields behind them were set in the constructor and never read there either, because dispatch switches on the element count directly. They were exported API on a new type that only assertions used. The two tests that checked the flags go with them; the tests that invoke a WindowFn and check which window the element lands in already cover the behaviour. The identically named (*window.Fn).NeedsElement is unrelated and stays: getSideWindowMappingUrn uses it. window_test.go had four names for the same idea: out, expected, expect and want. Across the exec test package want and wantErr are the convention by 37 declarations to 7, so settle on want. The benchmark uses b.N to match the other eleven benchmark files, and the side input mapping error uses errors.Errorf, which exec production code prefers 82 to 17. Fold the second init in windowinto.go into the first, drop two em dashes from comments, and record the feature in CHANGES. --- CHANGES.md | 1 + sdks/go/pkg/beam/core/graph/window/invoke.go | 17 +---- .../pkg/beam/core/graph/window/invoke_test.go | 67 ++----------------- sdks/go/pkg/beam/core/runtime/exec/window.go | 2 +- .../pkg/beam/core/runtime/exec/window_test.go | 49 ++++++-------- .../test/integration/primitives/windowinto.go | 14 ++-- 6 files changed, 37 insertions(+), 113 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index cc1ec48ba188..d138a24399cd 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -75,6 +75,7 @@ ([#38139](https://github.com/apache/beam/issues/38139)). * (Python) Added type alias for with_exception_handling to be used for typehints. ([#38173](https://github.com/apache/beam/issues/38173)). * Added plugin mechanism to support different Lineage implementations (Java) ([#36790](https://github.com/apache/beam/issues/36790)). +* (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/invoke.go b/sdks/go/pkg/beam/core/graph/window/invoke.go index 663fff9a2aba..715906488da9 100644 --- a/sdks/go/pkg/beam/core/graph/window/invoke.go +++ b/sdks/go/pkg/beam/core/graph/window/invoke.go @@ -48,9 +48,7 @@ type anyKVAssigner interface { // 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 - needsElement bool - isKV bool + call func(ts typex.EventTime, elm, elm2 any) []typex.Window } // NewWindowFnInvoker builds an invoker for fn. The concrete type of fn must @@ -70,7 +68,7 @@ func NewWindowFnInvoker(fn any) (*WindowFnInvoker, error) { return nil, errors.Errorf("window.NewWindowFnInvoker: type %v is not registered; call window.RegisterWindowFn during init()", t) } - inv := &WindowFnInvoker{needsElement: len(elems) > 0, isKV: len(elems) > 1} + inv := &WindowFnInvoker{} switch len(elems) { case 0: @@ -128,14 +126,3 @@ func NewWindowFnInvoker(fn any) (*WindowFnInvoker, error) { func (inv *WindowFnInvoker) Invoke(ts typex.EventTime, elm, elm2 any) []typex.Window { return inv.call(ts, elm, elm2) } - -// NeedsElement reports whether the underlying WindowFn accepts an element. -func (inv *WindowFnInvoker) NeedsElement() bool { - return inv.needsElement -} - -// IsKV reports whether the underlying WindowFn takes a KV element as a -// separate key and value. -func (inv *WindowFnInvoker) IsKV() bool { - return inv.isKV -} diff --git a/sdks/go/pkg/beam/core/graph/window/invoke_test.go b/sdks/go/pkg/beam/core/graph/window/invoke_test.go index 6dddfef93773..3f54187e69da 100644 --- a/sdks/go/pkg/beam/core/graph/window/invoke_test.go +++ b/sdks/go/pkg/beam/core/graph/window/invoke_test.go @@ -21,18 +21,20 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" ) -// elemAwareAnyWindowFn accepts an element typed as any — fast path 2. +// 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) + start := ts - ((ts%size)+size)%size return []typex.Window{IntervalWindow{Start: start, End: start + size}} } -// elemAwareConcreteWindowFn accepts a concrete element type — reflect path. +// elemAwareConcreteWindowFn accepts a concrete element type, so the invoker +// has to dispatch through reflect. type elemAwareConcreteWindowFn struct { DefaultSizeMs int64 } @@ -42,7 +44,7 @@ func (f *elemAwareConcreteWindowFn) AssignWindows(ts typex.EventTime, elem int64 if size <= 0 { size = typex.EventTime(f.DefaultSizeMs) } - start := ts - (ts % size) + start := ts - ((ts%size)+size)%size return []typex.Window{IntervalWindow{Start: start, End: start + size}} } @@ -84,12 +86,6 @@ func TestWindowFnInvoker_KV(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { inv := mustInvoker(t, tc.fn) - if !inv.NeedsElement() { - t.Error("NeedsElement() = false, want true") - } - if !inv.IsKV() { - t.Error("IsKV() = false, want true") - } windows := inv.Invoke(7500, "key", int64(5000)) if len(windows) != 1 { @@ -103,33 +99,10 @@ func TestWindowFnInvoker_KV(t *testing.T) { } } -func TestWindowFnInvoker_IsKV(t *testing.T) { - tests := []struct { - name string - fn any - want bool - }{ - {"timestamp-only", &testWindowFn{BucketSize: 1000}, false}, - {"single element", &elemAwareConcreteWindowFn{DefaultSizeMs: 1000}, false}, - {"kv element", &kvConcreteWindowFn{}, true}, - } - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - if got := mustInvoker(t, tc.fn).IsKV(); got != tc.want { - t.Errorf("IsKV() = %v, want %v", got, tc.want) - } - }) - } -} - func TestWindowFnInvoker_TimestampOnly(t *testing.T) { fn := &testWindowFn{BucketSize: 3000} inv := mustInvoker(t, fn) - if inv.NeedsElement() { - t.Fatal("NeedsElement() = true, want false") - } - windows := inv.Invoke(1500, nil, nil) if len(windows) != 1 { t.Fatalf("got %d windows, want 1", len(windows)) @@ -144,10 +117,6 @@ func TestWindowFnInvoker_AnyElem(t *testing.T) { fn := &elemAwareAnyWindowFn{SizeMs: 5000} inv := mustInvoker(t, fn) - if !inv.NeedsElement() { - t.Fatal("NeedsElement() = false, want true") - } - windows := inv.Invoke(7500, "ignored", nil) if len(windows) != 1 { t.Fatalf("got %d windows, want 1", len(windows)) @@ -162,10 +131,6 @@ func TestWindowFnInvoker_ConcreteElem(t *testing.T) { fn := &elemAwareConcreteWindowFn{DefaultSizeMs: 1000} inv := mustInvoker(t, fn) - if !inv.NeedsElement() { - t.Fatal("NeedsElement() = false, want true") - } - // Element provides window size of 5000ms. windows := inv.Invoke(7500, int64(5000), nil) if len(windows) != 1 { @@ -187,26 +152,6 @@ func TestWindowFnInvoker_ConcreteElem(t *testing.T) { } } -func TestWindowFnInvoker_NeedsElementCorrectness(t *testing.T) { - tests := []struct { - name string - fn any - want bool - }{ - {"timestamp-only", &testWindowFn{BucketSize: 1000}, false}, - {"any-elem", &elemAwareAnyWindowFn{SizeMs: 1000}, true}, - {"concrete-elem", &elemAwareConcreteWindowFn{DefaultSizeMs: 1000}, true}, - } - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - inv := mustInvoker(t, tc.fn) - if got := inv.NeedsElement(); got != tc.want { - t.Errorf("NeedsElement() = %v, want %v", got, tc.want) - } - }) - } -} - func TestWindowFnInvoker_ErrorOnUnregistered(t *testing.T) { type unregisteredFn struct{} if _, err := NewWindowFnInvoker(&unregisteredFn{}); err == nil { diff --git a/sdks/go/pkg/beam/core/runtime/exec/window.go b/sdks/go/pkg/beam/core/runtime/exec/window.go index bfc9b2ce1f6f..3857e5675372 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/window.go +++ b/sdks/go/pkg/beam/core/runtime/exec/window.go @@ -211,7 +211,7 @@ func (f *windowMapper) MapWindow(w typex.Window) (typex.Window, error) { // 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, fmt.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 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. 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 9211e4505f24..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(), @@ -139,18 +139,18 @@ func TestAssignWindow(t *testing.T) { for _, test := range tests { out := assignWindows(test.fn, mustInvokerFor(t, test.fn), test.in, nil, nil) - if !window.IsEqualList(out, test.out) { - t.Errorf("assignWindows(%v, %v) = %v, want %v", test.fn, test.in, out, test.out) + 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", @@ -189,18 +189,18 @@ func TestMapWindow(t *testing.T) { 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", @@ -217,7 +217,7 @@ 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: mustWindowMapper(t, tc.wFn), Out: out} @@ -234,8 +234,8 @@ 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...)) } }) } @@ -373,7 +373,7 @@ func BenchmarkAssignWindowsCustom(b *testing.B) { b.Fatalf("invokerFor(%v) failed: %v", fn, err) } b.ReportAllocs() - for b.Loop() { + for i := 0; i < b.N; i++ { assignWindows(fn, inv, 1500, nil, nil) } } @@ -430,15 +430,8 @@ type fixedCustomWindowFn struct { func (f *fixedCustomWindowFn) AssignWindows(ts typex.EventTime) []typex.Window { size := typex.EventTime(f.SizeMs) - start := ts - (ts % size) - if ts < 0 { - // Go's % truncates toward zero, so for negative dividends - // ts%size is non-positive and ts-(ts%size) rounds toward - // zero instead of toward -inf. The double-mod expression - // computes the Euclidean (non-negative) remainder, giving - // a correct floor to the window boundary. - start = ts - (ts%size+size)%size - } + // Euclidean remainder; correct floor for negative ts. + start := ts - ((ts%size)+size)%size return []typex.Window{window.IntervalWindow{Start: start, End: start + size}} } diff --git a/sdks/go/test/integration/primitives/windowinto.go b/sdks/go/test/integration/primitives/windowinto.go index 29f6ec79c1be..436e45ce6c0f 100644 --- a/sdks/go/test/integration/primitives/windowinto.go +++ b/sdks/go/test/integration/primitives/windowinto.go @@ -45,6 +45,8 @@ func init() { 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. @@ -432,9 +434,9 @@ type customFixedWindowFn struct { func (f *customFixedWindowFn) AssignWindows(ts typex.EventTime) []typex.Window { size := typex.EventTime(f.SizeMs) - start := ts - (ts.Add(time.Duration(f.SizeMs)*time.Millisecond) % mtime.FromDuration(time.Duration(f.SizeMs)*time.Millisecond)) - end := start + size - return []typex.Window{window.IntervalWindow{Start: start, End: end}} + // 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 @@ -479,10 +481,6 @@ type elemWithSize struct { SizeMs int64 } -func init() { - beam.RegisterType(reflect.TypeOf((*elemWithSize)(nil)).Elem()) -} - // elemAwareWindowFn uses the element's SizeMs field to determine the // window size, demonstrating data-driven window assignment. type elemAwareWindowFn struct{} @@ -492,7 +490,7 @@ func (f *elemAwareWindowFn) AssignWindows(ts typex.EventTime, elem elemWithSize) if size <= 0 { size = 1000 // fallback: 1s } - start := ts - (ts % size) + start := ts - ((ts%size)+size)%size return []typex.Window{window.IntervalWindow{Start: start, End: start + size}} } From 35c8868c6220e4f4303fcf78e744c212dd9c2d35 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:19:58 +0100 Subject: [PATCH 10/11] Fix interval window encoding Two defects in the interval window encoder, neither reachable before custom WindowFns, both reachable now that a user function can return any typex.Window. Encode returned nil when EncodeSingle failed, throwing the error away after the window count had already been written. The result is a truncated windowed value header reported as success, which corrupts the stream instead of failing the bundle. EncodeSingle asserted straight to IntervalWindow. NewCustom documents that a custom WindowFn must return IntervalWindow values and makeWindowCoder hands custom windowing the interval coder regardless, but nothing checked. Returning any other window panicked with "interface conversion: typex.Window is window.GlobalWindow, not window.IntervalWindow", naming neither the WindowFn nor the rule it broke. Both encoders now return an error saying what to fix, and the value encoder no longer asserts to IntervalWindow either, so a non-window element is reported rather than crashing. This is the only part of the branch outside the files custom windows touch, and it stands on its own if it is better reviewed separately. --- sdks/go/pkg/beam/core/runtime/exec/coder.go | 13 ++++++-- .../pkg/beam/core/runtime/exec/coder_test.go | 30 +++++++++++++++++++ 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/sdks/go/pkg/beam/core/runtime/exec/coder.go b/sdks/go/pkg/beam/core/runtime/exec/coder.go index 2c21ebea56b5..cc0accefdbbf 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/coder.go +++ b/sdks/go/pkg/beam/core/runtime/exec/coder.go @@ -1139,7 +1139,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 @@ -1147,7 +1147,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 } @@ -1193,7 +1196,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 75d18e533cf1..d6e17bf7808d 100644 --- a/sdks/go/pkg/beam/core/runtime/exec/coder_test.go +++ b/sdks/go/pkg/beam/core/runtime/exec/coder_test.go @@ -261,3 +261,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) + } + }) + } +} From a59e1a3260dd7ff620d44852049afea55d2ae587 Mon Sep 17 00:00:00 2001 From: Hannes Gustafsson Date: Sat, 12 Sep 2026 08:52:45 +0100 Subject: [PATCH 11/11] Fix conflicts with master CHANGES.md moves the entry to the 2.77.0 section, and the test filters follow master's lists, which no longer include samza. --- CHANGES.md | 207 ++++++++++++++++++++++-- sdks/go/test/integration/integration.go | 61 +------ 2 files changed, 200 insertions(+), 68 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index d138a24399cd..8ea41fe6e13e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -55,7 +55,7 @@ * ([#X](https://github.com/apache/beam/issues/X)). --> -# [2.74.0] - Unreleased +# [2.77.0] - Unreleased ## Highlights @@ -65,21 +65,27 @@ ## I/Os * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* Added `schema_update_options` to `WriteToBigQuery` file loads, allowing BigQuery load jobs to add nullable fields or relax required fields when appending data (Python) ([#21141](https://github.com/apache/beam/issues/21141)). +* BigQueryIO now supports reading BigQuery Lakehouse runtime catalog (BigLake metastore) Iceberg tables with the Storage Read API, using 4-part `project.catalog.namespace.table` identifiers (or a `TableReference` with a composite `catalog.namespace` dataset id). Previously such references were silently mis-parsed (Java) ([#39597](https://github.com/apache/beam/issues/39597)) . +* SolaceIO now supports reading and writing binary and text content data payload (Java) ([#39875](https://github.com/apache/beam/issues/39875)). +* ClickHouseIO: support writing `Decimal(P, S)` / `Decimal32/64/128/256` columns (Java) ([#39840](https://github.com/apache/beam/issues/39840)). ## New Features / Improvements * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). -* TriggerStateMachineRunner changes from BitSetCoder to SentinelBitSetCoder to - encode finished bitset. SentinelBitSetCoder and BitSetCoder are state - compatible. Both coders can decode encoded bytes from the other coder - ([#38139](https://github.com/apache/beam/issues/38139)). -* (Python) Added type alias for with_exception_handling to be used for typehints. ([#38173](https://github.com/apache/beam/issues/38173)). -* Added plugin mechanism to support different Lineage implementations (Java) ([#36790](https://github.com/apache/beam/issues/36790)). +* (Java/Python) `Watch` can bound its deduplication state by event time, retiring an output key once the greatest emitted timestamp has moved more than the allowed lateness past it. Java adds `Watch.growthOf(...).withTimestampCursor()`. Python adds `allowed_lateness` for the existing `timestamp_cursor` option ([#18459](https://github.com/apache/beam/issues/18459)). +* (Java) Spark Structured Streaming runner: stateful ParDo with state, timers, `@RequiresTimeSortedInput` and tagged outputs is now supported in batch mode ([#39779](https://github.com/apache/beam/issues/39779)). +* (Python) Added support for Vertex AI Model Monitoring V2 in RunInference ([#39738](https://github.com/apache/beam/issues/39738)). +* [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 -* X behavior was changed ([#X](https://github.com/apache/beam/issues/X)). +* Portable Java SDK now encodes SchemaCoders in a portable way ([#34672](https://github.com/apache/beam/issues/34672)). + - Original custom Java coder encoding can still be obtained using [StreamingOptions.setUpdateCompatibilityVersion("2.76")](https://github.com/apache/beam/blob/2cf0930e7ae1aa389c26ce6639b584877a3e31d9/sdks/java/core/src/main/java/org/apache/beam/sdk/options/StreamingOptions.java#L47) ([#34672](https://github.com/apache/beam/issues/34672)). + - Fixes ([#36496](https://github.com/apache/beam/issues/36496)), ([#30276](https://github.com/apache/beam/issues/30276)), ([#29245](https://github.com/apache/beam/issues/29245)). ## Deprecations @@ -87,7 +93,12 @@ ## Bugfixes -* Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when multiple requests share the same enrichment key (Python) ([#38035](https://github.com/apache/beam/issues/38035)). +* (Java) Fixed the Spark runner firing processing-time timers in reverse timestamp order ([#39824](https://github.com/apache/beam/issues/39824)). +* (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)). +* (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)). +* (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)). +* (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)). +* (Go) Fixed GCS glob matching silently dropping objects when the glob pattern contains multi-byte characters ([#39969](https://github.com/apache/beam/issues/39969)). ## Security Fixes @@ -98,16 +109,172 @@ [comment]: # ( When updating known issues after release, make sure also update website blog in website/www/site/content/blog.) * ([#X](https://github.com/apache/beam/issues/X)). -# [2.73.0] - 2026-04-?? +# [2.76.0] - 2026-08-31 + +## Highlights + +* Added a full Iceberg batch and streaming changelog source (CDC) ([#38831](https://github.com/apache/beam/issues/38831)) +* (Java) Added per-element OpenTelemetry trace propagation across stages in the Dataflow Streaming Runner. Enable it with `--experiments=enable_otel_defaults,element_metadata_supported,disable_portable_worker`. Cloud Trace incurs additional cost. ([#33176](https://github.com/apache/beam/issues/33176)) +* (Java) Added OpenTelemetry header propagation support for both reads and writes in KafkaIO and PubSubIO. ([#33176](https://github.com/apache/beam/issues/33176)) +* (Java) Added OpenTelemetry tracing support for SpannerIO change streams ([#33176](https://github.com/apache/beam/issues/33176)) +* (Python) JmsIO (IBM MQ, ActiveMQ, and other providers) is now supported in Python via cross-language ([#30716](https://github.com/apache/beam/issues/30716)). + +## I/Os + +* Upgraded Iceberg dependency to 1.11.0 (Java) ([#38925](https://github.com/apache/beam/issues/38925)). +* Add ArrowFlight IO (Java) ([#20116](https://github.com/apache/beam/issues/20116)). +* Added a Delta Lake batch changelog source (CDC) ([#39492](https://github.com/apache/beam/issues/39492)) + +## New Features / Improvements + +* Added `GroupIntoBatches` transform and the standard + `beam:coder:sharded_key:v1` coder to the Go SDK, along with + `beam.Coder.IsDeterministic`, `beam.PCollection.WindowingStrategy`, + and `coder.RegisterDeterministicCoder` for opt-in deterministic + custom coders (Go) ([#19868](https://github.com/apache/beam/issues/19868)). +* TriggerStateMachineRunner changes from BitSetCoder to SentinelBitSetCoder to + encode finished bitset. SentinelBitSetCoder and BitSetCoder are state + compatible. Both coders can decode encoded bytes from the other coder + ([#38139](https://github.com/apache/beam/issues/38139)). +* (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). +* (Java) Supported acknowledge mode for JmsIO ([#39253](https://github.com/apache/beam/issues/39253)). +* (Python) Staged files directory is now automatically added to `sys.path` on the Python SDK worker at startup. This makes Python files provided via the '--files_to_stage' pipeline option importable in the pipeline code and makes it easier to initialize Python SDK harness at startup via the `--beam_plugins` pipeline option. For more information, see the [Staging Individual Files](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/#staging-files) section of the dependency management docs. This behavior can be disabled by passing the '--experiments=no_staged_dir_in_sys_path' pipeline option ([#39431](https://github.com/apache/beam/issues/39431)). +* (Python) Added `equal_to_approx`, an `assert_that` matcher that compares numeric pipeline outputs with a configurable tolerance ([#18028](https://github.com/apache/beam/issues/18028)). +* (Python) `Timestamp` now supports variable subsecond precision, up to nanoseconds. The portable + `beam:logical_type:timestamp:v1` logical type now maps to Python's `Timestamp` ([#39344](https://github.com/apache/beam/issues/39344)). +* (Python) Added `UnboundedSource`, an interface for reading an infinite stream of records with checkpointing, watermark reporting, and bundle finalization. Read one with `beam.io.Read` + ([#19137](https://github.com/apache/beam/issues/19137)). +* (Python) Added `Watch`, a transform that polls a growing set of outputs for each input element, deduplicates outputs across poll rounds, and stops per a user-supplied termination condition + ([#21521](https://github.com/apache/beam/issues/21521)). +* (Python) Added support to analyze core dumps created after python worker segmentation faults with `pystack` (or `gdb` if installed) using the `--profiler_agent=coredump` pipeline option. ([#39484](https://github.com/apache/beam/issues/39484)). + +## Breaking Changes + +* (Python) Removed `google-perftools` from the SDK container images. Users who wish to use `--profiler_agent=tcmalloc` should install google-perftools APT package in their custom container images separately ([#39323](https://github.com/apache/beam/issues/39323)). +* [IcebergIO] Reading a `timestamptz` column will now return a `Timestamp.MICROS` Beam logical type to preserve + microseconds (the old Beam `Schema.FieldType#DATETIME` primitive type truncates past milliseconds). This may break + the following use cases when a `timestamptz` column is present: + * Existing streaming read pipelines. + * Managed Iceberg batch reads when upgraded from an older SDK. + * Python reads. + + Use pipeline option `--updateCompatibilityVersion=2.75.0` (or any older version) to keep the old behavior ([#39344](https://github.com/apache/beam/issues/39344)). +* `DoFn.process` returning a `str`, `bytes`, or `dict` (instead of an iterable wrapping one) now raises a `TypeError` rather than silently iterating per-character/byte/key (Python) ([#18712](https://github.com/apache/beam/issues/18712)). +* (Java) Added `DRAINING` and `DRAINED` states to `PipelineResult`, including runner state mappings and Dataflow update handling ([#39020](https://github.com/apache/beam/issues/39020)). +* (Java) IcebergIO and projects that use it must now be built with Java 17 or later as a result of Iceberg 1.11.0 upgrade ([#38925](https://github.com/apache/beam/issues/38925)). + +## Bugfixes + +* Fixed unresolved runtime `ValueProvider` options being stringified in Python Dataflow Flex Templates ([#39499](https://github.com/apache/beam/issues/39499)). +* Fixed unbounded checkpoint state growth for splittable DoFns that self-checkpoint on the portable Flink runner (Java) ([#27648](https://github.com/apache/beam/issues/27648)). +* Improved Java pipeline performance by avoiding repeated `DoFn` type descriptor resolution when creating cached invokers ([#39309](https://github.com/apache/beam/issues/39309)). +* (Python) Fixed a memory leak in Python SDK caused by storing exceptions with potentially large stack frames in a cache ([#39406](https://github.com/apache/beam/issues/39406)). + +# [2.75.0] - 2026-07-08 ## Highlights +* Python SDK now supports memory profiling with Memray ([#38853](https://github.com/apache/beam/issues/38853)). +* (Python) Added [Qdrant](https://qdrant.tech/) VectorDatabaseWriteConfig implementation ([#38141](https://github.com/apache/beam/issues/38141)). + +## I/Os + +* Support for reading from Delta Lake added (Java) ([#38551](https://github.com/apache/beam/issues/38551)). +* ClickHouseIO: support writing `DateTime64(precision[, 'timezone'])` columns with sub-second precision (Java) ([#38466](https://github.com/apache/beam/issues/38466)). +* Upgraded IO Expansion Service to Java 17 ([#38974](https://github.com/apache/beam/issues/38974)). +* SpannerIO: Added support for Cloud Spanner Directed Reads (Java) ([#X](https://github.com/apache/beam/issues/X)). + +## New Features / Improvements + +* Dataflow Runner v2 has been renamed to Dataflow Portable Runner. Please refer to Dataflow [public documentation](https://docs.cloud.google.com/dataflow/docs/runner-v2) on when to enable Portable Runner.([#39000](https://github.com/apache/beam/issues/39000)). +* (Java) Enabled state tag encoding v2 by default for new Dataflow Streaming Engine jobs. It can be disabled by passing `--experiments=disable_streaming_engine_state_tag_encoding_v2` or `--updateCompatibilityVersion=2.74.0` pipeline option. Note that the tag encoding version cannot change during a job update. Jobs using tag encoding v2 (enabled by default for new jobs on 2.75.0+) cannot be downgraded to Beam versions prior to 2.73.0, as only versions 2.73.0 and later support tag encoding v2. ([#38705](https://github.com/apache/beam/issues/38705)). +* (Python) Added instrumentation to support off-the-shelf profiling agents when launching Python SDK Harness ([#38853](https://github.com/apache/beam/issues/38853)). +* (Java) Added support to the FnApi Data stream protocol allowing runners to isolate bundles slowly processing input from other bundles. ([#39001](https://github.com/apache/beam/issues/39001)). +* (YAML) Switched js2py library to Quickjs ([#38473](https://github.com/apache/beam/issues/38473)). +* (YAML) Added HuggingFaceModelHandler for YAML usage ([#38696](https://github.com/apache/beam/issues/38696)). +* (YAML) Added WriteToMongoDB transform ([#38376](https://github.com/apache/beam/issues/38376)). +* (YAML) Added WriteToDatadog transform ([#38362](https://github.com/apache/beam/issues/38362)). +* (Java) Flink 2.1 and 2.2 support is added ([#38947](https://github.com/apache/beam/issues/38947)) ([#38978](https://github.com/apache/beam/issues/38978)); Flink 1.17 and 1.18 support is dropped. +* (Python) MqttIO is now supported in Python via cross-language ([#21060](https://github.com/apache/beam/issues/21060)). + +## Breaking Changes + +* (Python) Typehints of dataclass fields are honored during type inferences. To restore the behavior of fallback-to-any, + use pipeline option `--exclude_infer_dataclass_field_type` ([#38797](https://github.com/apache/beam/issues/38797)). + However fixing forward is recommended. + +## Bugfixes + +* Fixed GCS filesystem glob matching to correctly handle `/` in object names and support `**` for recursive matching (Go) ([#38059](https://github.com/apache/beam/issues/38059)). +* Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when multiple requests share the same enrichment key (Python) ([#38035](https://github.com/apache/beam/issues/38035)). +* Fixed IcebergIO writing manifest column bounds padded with trailing `0x00` bytes, which broke equality predicate pushdown in some query engines (Java) ([#38580](https://github.com/apache/beam/issues/38580)). + +## Known Issues + +* (Java) Projects using the Flink runner with Flink 2.1 or later alongside libraries requiring `org.lz4:lz4-java` (e.g., Kafka clients) may encounter a Gradle capability conflict, because Flink 2.1+ ships `at.yawk.lz4:lz4-java` which declares the same capability. To resolve, add a `capabilitiesResolution` rule to your `build.gradle` that selects `at.yawk.lz4:lz4-java` ([#38947](https://github.com/apache/beam/issues/38947)). +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). +* (Java) Pipelines with a moderate to heavy Cloud Storage read workload might experience a performance regression ([#39548](https://github.com/apache/beam/issues/39548)). +* (Java) Pipelines using the Dataflow Runner and Java versions 17+ may experience spiky memory caused by a JVM upgrade in the runner image ([#39897](https://github.com/apache/beam/issues/39897). + +# [2.74.0] - 2026-06-02 + +## Highlights + +* Spark 4 runner support for Java SDK ([#38255](https://github.com/apache/beam/issues/38255)). + +## I/Os + +* IcebergIO: support declaring a table's sort order on dynamic table creation via the new `sort_fields` config ([#38269](https://github.com/apache/beam/issues/38269)). +* IcebergIO: support writing with hash distribution mode, and with autosharding ([#38061](https://github.com/apache/beam/issues/38061)). + +## New Features / Improvements + +* (Java) Added an experimental Kafka Streams runner, which executes a Beam pipeline as an ordinary Kafka Streams application with no cluster to operate. It supports a subset of the model and is not built by default; pass `-Pwith-kafka-streams-runner` to include it ([#18479](https://github.com/apache/beam/issues/18479)). +* Capability introduces an indicator for aggregations and timers firing during a pipeline drain, allowing users and sinks to recognize and appropriately handle potentially incomplete or partial data ([#36884](https://github.com/apache/beam/issues/36884)). +* Added support for setting disk provisioned IOPS and throughput in Dataflow runner via `--diskProvisionedIops` and `--diskProvisionedThroughputMibps` pipeline options (Java/Go/Python) ([#38349](https://github.com/apache/beam/issues/38349)). +* TriggerStateMachineRunner changes from BitSetCoder to SentinelBitSetCoder to + encode finished bitset. SentinelBitSetCoder and BitSetCoder are state + compatible. Both coders can decode encoded bytes from the other coder + ([#38139](https://github.com/apache/beam/issues/38139)). +* (Python) Added type alias for with_exception_handling to be used for typehints. ([#38173](https://github.com/apache/beam/issues/38173)). +* (Java) BatchElements transform for Java SDK ([#38369](https://github.com/apache/beam/issues/38369)) +* Added plugin mechanism to support different Lineage implementations (Java) ([#36790](https://github.com/apache/beam/issues/36790)). +* (Python) Supported Python user type in Beam SQL. For example, SQL statements like `SELECT some_field from PCOLLECTION` can now operate a PCollection of Beam Row containing pickable Python user type ([#20738](https://github.com/apache/beam/issues/20738)). +* (Python) Introduced `beam.coders.registry.register_row` as preferred API to register a named tuple or dataclass with a Beam Row. At pipelne runtime, the original type associated with the registered row are preserved across the serialization boundary ([#38108](https://github.com/apache/beam/issues/38108)). +* (Python) Added `type_overrides` parameter to `WriteToBigQuery` allowing users to specify custom BigQuery to Python type mappings when using Storage Write API. This enables support for types like DATE, DATETIME, and JSON (Python) ([#25946](https://github.com/apache/beam/issues/25946)). + +## Breaking Changes + +* (Python) Made Beartype the default fallback type checking tool. This can be disabled with the `--disable_beartype` pipeline option. ([#38275](https://github.com/apache/beam/issues/38275)) + +## Deprecations + +* Dropped Java 8 support ([#31678](https://github.com/apache/beam/issues/31678)). +* Removed Samza Runner support ([#35448](https://github.com/apache/beam/issues/35448)). + +## Bugfixes + +* Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when multiple requests share the same enrichment key (Python) ([#38035](https://github.com/apache/beam/issues/38035)). +* Added `max_batch_duration_secs` passthrough support in Python Enrichment BigQuery and CloudSQL handlers so batching duration can be forwarded to `BatchElements` ([#38243](https://github.com/apache/beam/issues/38243)). + +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). +* (Java) Pipelines with a moderate to heavy Cloud Storage read workload might experience a performance regression ([#39548](https://github.com/apache/beam/issues/39548)). +* (Java) Pipelines using the Dataflow Runner and Java versions 17+ may experience spiky memory caused by a JVM upgrade in the runner image ([#39897](https://github.com/apache/beam/issues/39897). + +# [2.73.0] - 2026-04-29 + +## Highlights + + ## I/Os * DebeziumIO (Java): added `OffsetRetainer` interface and `FileSystemOffsetRetainer` implementation to persist and restore CDC offsets across pipeline restarts, and exposed `withStartOffset` / `withOffsetRetainer` on `DebeziumIO.Read` and the cross-language `ReadBuilder` ([#28248](https://github.com/apache/beam/issues/28248)). ## New Features / Improvements +* (Python) Added BigQuery CDC streaming source ([#37724](https://github.com/apache/beam/issues/37724)) * Added `ADKAgentModelHandler` for running Google Agent Development Kit (ADK) agents (Python) ([#37917](https://github.com/apache/beam/issues/37917)). * (Python) Added exception chaining to preserve error context in CloudSQLEnrichmentHandler, processes utilities, and core transforms ([#37422](https://github.com/apache/beam/issues/37422)). * (Python) Added a pipeline option `--experiments=pip_no_build_isolation` to disable build isolation when installing dependencies in the runtime environment ([#37331](https://github.com/apache/beam/issues/37331)). @@ -121,6 +288,7 @@ ## Breaking Changes * The Python SDK container's `boot.go` now passes pipeline options through a file instead of the `PIPELINE_OPTIONS` environment variable. If a user pairs a new Python SDK container with an older SDK version (which does not support the file-based approach), the pipeline options will not be recognized and the pipeline will fail. Users must ensure their SDK and container versions are synchronized ([#37370](https://github.com/apache/beam/issues/37370)). +* Python DoFn.with_exception_handling now respects user DoFn typehints. This can break update compatibility if coders change. It can also break pipeline compilation if existing typehints are incorrect. To update safely sepcify the pipeline option `--update_compatibility_version=2.72.0`. To fix typehints replace any incorrect typehints that were previously ignored ([#37590](https://github.com/apache/beam/issues/37590)) ## Bugfixes @@ -131,6 +299,10 @@ * Fixed [CVE-2023-46604](https://www.cve.org/CVERecord?id=CVE-2023-46604) (CVSS 10.0) and [CVE-2022-41678](https://www.cve.org/CVERecord?id=CVE-2022-41678) by upgrading ActiveMQ from 5.14.5 to 5.19.2 (Java) ([#37943](https://github.com/apache/beam/issues/37943)). * Fixed [CVE-2024-1597](https://www.cve.org/CVERecord?id=CVE-2024-1597), [CVE-2022-31197](https://www.cve.org/CVERecord?id=CVE-2022-31197), and [CVE-2022-21724](https://www.cve.org/CVERecord?id=CVE-2022-21724) by upgrading PostgreSQL JDBC Driver from 42.2.16 to 42.6.2 (Java) ([#37942](https://github.com/apache/beam/issues/37942)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.72.0] - 2026-03-30 ## Highlights @@ -164,6 +336,10 @@ * Fixed [CVE-2024-28397](https://www.cve.org/CVERecord?id=CVE-2024-28397) by switching from js2py to pythonmonkey (Yaml) ([#37560](https://github.com/apache/beam/issues/37560)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.71.0] - 2026-01-22 ## I/Os @@ -185,6 +361,7 @@ ## Known Issues +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.70.0] - 2025-12-16 @@ -207,6 +384,10 @@ Now Beam has full support for Milvus integration including Milvus enrichment and * (Python) Python 3.9 reached EOL in October 2025 and support for the language version has been removed. ([#36665](https://github.com/apache/beam/issues/36665)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.69.0] - 2025-10-28 ## Highlights @@ -266,6 +447,10 @@ Now Beam has full support for Milvus integration including Milvus enrichment and ([#36141](https://github.com/apache/beam/issues/36141)). * Fixed Spanner Change Stream reading stuck issue due to watermark of partition moving backwards ([#36470](https://github.com/apache/beam/issues/36470)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.68.0] - 2025-09-22 ## Highlights @@ -320,6 +505,7 @@ Now Beam has full support for Milvus integration including Milvus enrichment and ## Known Issues * ([#36470](https://github.com/apache/beam/issues/36470)). Spanner Change Stream reading stuck issue due to watermark of partition moving backwards. This issue exists in 2.67.0 and 2.68.0. To mitigate the issue, either use old version 2.66.0 or go to 2.69.0. +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.67.0] - 2025-08-12 @@ -368,6 +554,7 @@ Now Beam has full support for Milvus integration including Milvus enrichment and * ([#35666](https://github.com/apache/beam/issues/35666)). YAML Flatten incorrectly drops fields when input PCollections' schema are different. This issue exists for all versions since 2.52.0. * ([#36470](https://github.com/apache/beam/issues/36470)). Spanner Change Stream reading stuck issue due to watermark of partition moving backwards. This issue exists in 2.67.0 and 2.68.0. To mitigate the issue, either use old version 2.66.0 or go to 2.69.0. +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.66.0] - 2025-07-01 diff --git a/sdks/go/test/integration/integration.go b/sdks/go/test/integration/integration.go index fba7e6dc137c..50d3e3332f47 100644 --- a/sdks/go/test/integration/integration.go +++ b/sdks/go/test/integration/integration.go @@ -67,7 +67,7 @@ var directFilters = []string{ // The direct runner does not yet support cross-language. "TestXLang.*", "TestKafkaIO.*", - "TestBigQueryIO.*", + "TestBigQueryIO_[^WQ].*", "TestBigtableIO.*", "TestSpannerIO.*", "TestDebeziumIO_BasicRead", @@ -205,61 +205,8 @@ var flinkFilters = []string{ "TestTestStreamToGBK", "TestTestStreamTimersEventTime", - "TestTimers_EventTime_Unbounded", // (failure when comparing on side inputs (NPE on window lookup)) - "TestTimers_ProcessingTime.*", // Flink doesn't support processing time timers. - - // 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 samzaFilters = []string{ - // TODO(https://github.com/apache/beam/issues/20987): Samza tests invalid encoding. - "TestReshuffle", - "TestReshuffleKV", - // The Samza runner does not support the TestStream primitive - "TestTestStream.*", - // The trigger and pane tests uses TestStream - "TestTrigger.*", - "TestPanes", - // TODO(https://github.com/apache/beam/issues/21244): Samza doesn't yet support post job metrics, used by WordCount - "TestWordCount.*", - // TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners. - "TestBigQueryIO.*", - "TestBigtableIO.*", - "TestSpannerIO.*", - // The Samza runner does not support self-checkpointing - "TestCheckpointing", - // The samza runner does not support pipeline drain for SDF. - "TestDrain", - // FhirIO currently only supports Dataflow runner - "TestFhirIO.*", - // OOMs currently only lead to heap dumps on Dataflow runner - "TestOomParDo", - // The samza runner does not support user state. - "TestValueState", - "TestValueStateWindowed", - "TestValueStateClear", - "TestBagState", - "TestBagStateClear", - "TestCombiningState", - "TestMapState", - "TestMapStateClear", - "TestSetState", - "TestSetStateClear", - // TODO(https://github.com/apache/beam/issues/26126): Java runner issue (AcitveBundle has no regsitered handler) - "TestDebeziumIO_BasicRead", - - "TestOrderedListState", - - // Samza does not support state. - "TestTimers.*", - "TestBagStateBlindWrite", + "TestTimers_EventTime_WithNoOutputTimestamp", // Encounter error: TimestampCombiner moved element from TIMESTAMP_MAX_VALUE to earlier time (end of global window) for window GlobalWindow + "TestTimers_ProcessingTime.*", // Flink doesn't support processing time timers. // no support for BundleFinalizer "TestParDoBundleFinalizer.*", @@ -405,8 +352,6 @@ func CheckFilters(t *testing.T) { filters = portableFilters case "flink", "FlinkRunner": filters = flinkFilters - case "samza", "SamzaRunner": - filters = samzaFilters case "spark", "SparkRunner": filters = sparkFilters case "dataflow", "DataflowRunner":