Add custom window support in Go SDK

CHANGES.md

@@ -79,6 +79,7 @@
 * [Flink Runner] Added opt-in static round-robin split assignment for small bounded sources via the new `sourceStaticSplitThresholdMb` pipeline option. The default of 0 keeps the existing lazy pull-based assignment ([#39873](https://github.com/apache/beam/issues/39873)).
 * Added automatic caching of bounded, single-pane side-input views for classic Java Flink DataStream execution ([#39866](https://github.com/apache/beam/issues/39866)).
 * (Python) Added `Sample.Any`, the Python equivalent of Java's `Sample.any`, which returns up to n arbitrary elements from a PCollection ([#18552](https://github.com/apache/beam/issues/18552)).
+* (Go) Added support for custom WindowFns ([#20627](https://github.com/apache/beam/issues/20627)).
 
 ## Breaking Changes
 

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,13 +55,51 @@ func NewSessions(gap time.Duration) *Fn {
 	return &Fn{Kind: Sessions, Gap: gap}
 }
 
+// NewCustom returns a WindowFn backed by a user-defined custom window
+// function. The fn value must be a pointer-to-struct whose concrete type
+// has been registered with [RegisterWindowFn] during init. Custom window
+// functions are non-merging and must return [IntervalWindow] values from
+// their AssignWindows method.
+//
+// NewCustom panics if fn is nil or its type was not registered.
+func NewCustom(fn any) *Fn {
+	if fn == nil {
+		panic("window.NewCustom: fn must not be nil")
+	}
+	t := reflect.TypeOf(fn)
+	st := t
+	if t.Kind() == reflect.Pointer {
+		st = t.Elem()
+	}
+	if _, ok := LookupWindowFn(st); !ok {
+		panic(fmt.Sprintf("window.NewCustom: type %v is not registered; call window.RegisterWindowFn during init()", t))
+	}
+	return &Fn{Kind: CustomWindows, CustomFn: fn}
+}
+
 // Fn defines the window fn.
 type Fn struct {
 	Kind Kind
 
 	Size   time.Duration // FixedWindows, SlidingWindows
 	Period time.Duration // SlidingWindows
 	Gap    time.Duration // Sessions
+
+	CustomFn any // CustomWindows (nil for built-in kinds)
+}
+
+// NeedsElement reports whether a CustomWindows Fn has an element-aware
+// AssignWindows signature. Returns false for all built-in window kinds.
+func (w *Fn) NeedsElement() bool {
+	if w.Kind != CustomWindows || w.CustomFn == nil {
+		return false
+	}
+	t := reflect.TypeOf(w.CustomFn)
+	if t.Kind() == reflect.Pointer {
+		t = t.Elem()
+	}
+	elems, _ := LookupWindowFn(t)
+	return len(elems) > 0
 }
 
 // TODO(herohde) 4/17/2018: do we need to expose the window type as well?
@@ -82,6 +122,8 @@ func (w *Fn) String() string {
 		return fmt.Sprintf("%v[%v@%v]", w.Kind, w.Size, w.Period)
 	case Sessions:
 		return fmt.Sprintf("%v[%v]", w.Kind, w.Gap)
+	case CustomWindows:
+		return fmt.Sprintf("%v[%T %+v]", w.Kind, w.CustomFn, w.CustomFn)
 	default:
 		return string(w.Kind)
 	}
@@ -105,6 +147,8 @@ func (w *Fn) Equals(o *Fn) bool {
 		return w.Period == o.Period && w.Size == o.Size
 	case Sessions:
 		return w.Gap == o.Gap
+	case CustomWindows:
+		return reflect.DeepEqual(w.CustomFn, o.CustomFn)
 	default:
 		panic(fmt.Sprintf("unknown window type: %v", w))
 	}

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)
+	}
+}

sdks/go/pkg/beam/core/graph/window/invoke.go

@@ -0,0 +1,128 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package window
+
+import (
+	"reflect"
+
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
+)
+
+// tsOnlyAssigner is the fast-path interface for timestamp-only custom WindowFns.
+type tsOnlyAssigner interface {
+	AssignWindows(typex.EventTime) []typex.Window
+}
+
+// anyElemAssigner is the fast-path interface for element-aware WindowFns whose
+// element parameter is typed as any.
+type anyElemAssigner interface {
+	AssignWindows(typex.EventTime, any) []typex.Window
+}
+
+// anyKVAssigner is the fast-path interface for KV-aware WindowFns whose key and
+// value parameters are both typed as any.
+type anyKVAssigner interface {
+	AssignWindows(typex.EventTime, any, any) []typex.Window
+}
+
+// WindowFnInvoker wraps a custom WindowFn instance and dispatches AssignWindows
+// through a path chosen once at construction: a typed interface assertion where
+// the signature permits one, otherwise reflect.Value.Call, which costs a small
+// per-call allocation.
+//
+// Element parameters follow the convention the DoFn invoker uses for main
+// input: a KV element arrives as a key and a value, anything else as a single
+// element.
+type WindowFnInvoker struct {
+	call func(ts typex.EventTime, elm, elm2 any) []typex.Window
+}
+
+// NewWindowFnInvoker builds an invoker for fn. The concrete type of fn must
+// have been previously registered via RegisterWindowFn.
+func NewWindowFnInvoker(fn any) (*WindowFnInvoker, error) {
+	t := reflect.TypeOf(fn)
+	if t == nil {
+		return nil, errors.New("window.NewWindowFnInvoker: fn must not be nil")
+	}
+	structType := t
+	if t.Kind() == reflect.Pointer {
+		structType = t.Elem()
+	}
+
+	elems, ok := LookupWindowFn(structType)
+	if !ok {
+		return nil, errors.Errorf("window.NewWindowFnInvoker: type %v is not registered; call window.RegisterWindowFn during init()", t)
+	}
+
+	inv := &WindowFnInvoker{}
+
+	switch len(elems) {
+	case 0:
+		if a, ok := fn.(tsOnlyAssigner); ok {
+			inv.call = func(ts typex.EventTime, _, _ any) []typex.Window {
+				return a.AssignWindows(ts)
+			}
+			return inv, nil
+		}
+	case 1:
+		if a, ok := fn.(anyElemAssigner); ok {
+			inv.call = func(ts typex.EventTime, elm, _ any) []typex.Window {
+				return a.AssignWindows(ts, elm)
+			}
+			return inv, nil
+		}
+	default:
+		if a, ok := fn.(anyKVAssigner); ok {
+			inv.call = func(ts typex.EventTime, elm, elm2 any) []typex.Window {
+				return a.AssignWindows(ts, elm, elm2)
+			}
+			return inv, nil
+		}
+	}
+
+	// Concrete element types cannot be reached through an interface assertion.
+	m := reflect.ValueOf(fn).MethodByName("AssignWindows")
+	if !m.IsValid() {
+		return nil, errors.Errorf("window.NewWindowFnInvoker: %v has no AssignWindows method", t)
+	}
+
+	switch len(elems) {
+	case 0:
+		inv.call = func(ts typex.EventTime, _, _ any) []typex.Window {
+			out := m.Call([]reflect.Value{reflect.ValueOf(ts)})
+			return out[0].Interface().([]typex.Window)
+		}
+	case 1:
+		inv.call = func(ts typex.EventTime, elm, _ any) []typex.Window {
+			out := m.Call([]reflect.Value{reflect.ValueOf(ts), reflect.ValueOf(elm)})
+			return out[0].Interface().([]typex.Window)
+		}
+	default:
+		inv.call = func(ts typex.EventTime, elm, elm2 any) []typex.Window {
+			out := m.Call([]reflect.Value{reflect.ValueOf(ts), reflect.ValueOf(elm), reflect.ValueOf(elm2)})
+			return out[0].Interface().([]typex.Window)
+		}
+	}
+	return inv, nil
+}
+
+// Invoke calls AssignWindows on the underlying WindowFn. Signatures that do not
+// take element parameters ignore elm and elm2; elm2 carries the value of a KV
+// element.
+func (inv *WindowFnInvoker) Invoke(ts typex.EventTime, elm, elm2 any) []typex.Window {
+	return inv.call(ts, elm, elm2)
+}