This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 3a0985a07da [Go SDK] Add GroupIntoBatches transform (#19868) (#38220)
3a0985a07da is described below
commit 3a0985a07daa426198ece0811146efc60292f4b5
Author: Florian TREHAUT <[email protected]>
AuthorDate: Tue Aug 4 21:04:11 2026 +0700
[Go SDK] Add GroupIntoBatches transform (#19868) (#38220)
* [Go SDK] Add Coder.IsDeterministic and ShardedKey standard coder
This introduces the supporting infrastructure required by the upcoming
GroupIntoBatches transform (#19868):
- (*coder.Coder).IsDeterministic() reports whether a coder produces
byte-stable output. Primitives (bytes, bool, varint, double, string)
are deterministic; composite coders (KV, CoGBK, Nullable, Iterable,
LP, ShardedKey) are deterministic iff every component is. Custom
user-registered coders are non-deterministic by default and opt in
via the new RegisterDeterministicCoder registration helper.
- beam.Coder.IsDeterministic() forwards to the inner coder's method so
transform authors can gate on determinism without reaching into
internals.
- beam.PCollection.WindowingStrategy() exposes the input's windowing
strategy publicly so transforms honoring allowed lateness (e.g.
GroupIntoBatches) can read it without package-private access.
- typex.ShardedKey[K] is a concrete Go generic struct representing a
sharded user key. The accompanying Kind (coder.ShardedKey) and
beam:coder:sharded_key:v1 URN wiring (graphx marshal/unmarshal, exec
encode/decode) produce the exact wire format documented in
standard_coders.yaml:501-521 — verified byte-identical against the
four published fixtures.
Cross-SDK byte compatibility is required for Dataflow/Flink
interoperability; a single divergent byte would silently corrupt
pipelines. Roundtrip tests cover all four yaml fixtures.
* [Go SDK] Add GroupIntoBatches transform and ShardedKey composite (#19868)
Builds on top of the Coder.IsDeterministic / DeterministicCoder
foundation and introduces the full user-facing surface for batching
PCollection<KV<K,V>> elements by key.
* typex.ShardedKey is added as a new Composite marker type (alongside
KV, CoGBK, WindowedValue, Timers). Its runtime representation is a
two-part FullValue (Elm=key, Elm2=[]byte shardID).
* coder.NewSK builds the associated coder; graphx/coder and exec/coder
wire the beam:coder:sharded_key:v1 URN in both directions. The wire
format is byte-identical to the Java util.ShardedKey.Coder and the
Python sharded_key coder — verified against all four
standard_coders.yaml fixtures (lines 501-521).
* beam.PCollection.WindowingStrategy and beam.Coder.IsDeterministic are
exposed publicly, matching the access pattern already used inside the
beam package (pardo.go, gbk.go).
* transforms/batch introduces GroupIntoBatches, a stateful DoFn that
buffers per-key values in a state.Bag and flushes when BatchSize /
BatchSizeBytes / MaxBufferingDuration / end-of-window + allowed
lateness triggers fire. The transform honors the input's allowed
lateness (Java parity; Python currently ignores it) and panics at
pipeline-build time on invalid params, non-KV inputs, or
non-deterministic key coders.
* CHANGES.md is updated under [2.74.0] - Unreleased.
Scope note: this release ships GroupIntoBatches with string keys and
string values. The underlying ShardedKey infrastructure is fully in
place (type, coder, URN, tests); GroupIntoBatchesWithShardedKey and
arbitrary K/V generics are follow-up work once the Go SDK binds
universal types through state.Bag element coders.
End-to-end Prism integration testing of the stateful DoFn path remains
a follow-up — the pipeline hangs on job completion in the bounded
case, pending investigation of Prism's watermark signalling for
event-time timers set on the GlobalWindow maxTimestamp. All unit tests
(coder roundtrip, Params validation, primitive sizer) pass.
* [Go SDK] Support generic K,V and WithShardedKey in GroupIntoBatches
(#19868)
Extends GroupIntoBatches to arbitrary key/value types and adds
GroupIntoBatchesWithShardedKey, completing the Apache Beam
GroupIntoBatches feature parity with Java/Python (#19868).
Generic K, V support:
- Replaces the string-only DoFn with a typex.T / typex.V universal
pair, resolved by beam.ParDo's type-binding engine at graph
construction. Values flow through a state.Bag[[]byte] encoded via
a cached beam.ElementEncoder/Decoder lazily initialised from
beam.EncodedType{T: valueType} — a single reflect.Type captured at
graph time and serialised across the SDK-worker boundary.
- Separates into two concrete DoFn shapes: the plain
groupIntoBatchesFn (event-time timer only) and
groupIntoBatchesBufferedFn (event-time + processing-time). A single
DoFn with an unused processing-time timer family stalls Prism
waiting for the family's completion signal — splitting the shape
by params.MaxBufferingDuration avoids the stall.
- ProcessingTime timer is only wired when the user requests
buffering, eliminating the Prism stall we hit on the initial
implementation.
WithShardedKey:
- Adds GroupIntoBatchesWithShardedKey(s, params, col) that
round-trips KV<K, V> → KV<[]byte-shardKey, V> → batched →
KV<K, []V>. ShardIDs are 24-byte worker-UUID + atomic-counter
tuples matching Java/Python layouts; downstream workers see
independent state per shard, so a single hot logical key's
processing spreads across workers on distributed runners.
- Output shape: PCollection<KV<K, []V>>, identical to
GroupIntoBatches. The Go SDK's type-binding engine does not
accept custom generic structs as DoFn output types, so we do not
surface ShardedKey<K> to the user. Cross-SDK bytes-compat
ShardedKey coder infrastructure is still wired at the core/typex
+ core/graph/coder level for future bidirectional pipelines.
Testing:
- End-to-end Prism tests for GroupIntoBatches across count, byte and
per-key-isolation triggers, including a non-string value type
(int).
- GroupIntoBatchesWithShardedKey pipeline construction test (Prism
panics on the 3-stage round-trip pipeline with "assignment to nil
map" in aggregateStageKind.buildEventTimeBundle — a runner-side
regression we verify does NOT reproduce on non-Prism runners).
Follow-up items documented in the package godoc.
* [Go SDK] Fix ShardedKey coder serialization for generic closures (#19868)
Go generic functions produce closures with identical compiler-assigned
symbol names across type instantiations — all RegisterShardedKeyType[K]
instantiations generated closures named
"RegisterShardedKeyType[...].func1", causing cross-worker
deserialization to resolve the wrong enc/dec function (last-registered
wins).
Root cause: reflectx.FunctionName calls runtime.FuncForPC which returns
the compiler name; Go does not qualify closure names by type parameter.
Fix: three surgical additions to core SDK infrastructure:
1. reflectx.MakeFuncWithName wraps a Func with a caller-supplied Name()
so the serializer (encodeUserFn → u.Fn.Name()) emits a
type-qualified name like "batch.encShardedKey[string]".
2. runtime.RegisterFunctionWithName registers a function under a custom
name in the resolution cache so the deserializer (decodeUserFn →
ResolveFunction) finds it.
3. coder.RegisterDeterministicCoderWithFuncs accepts pre-wrapped
funcx.Fn values carrying the qualified names, bypassing the
automatic name derivation in NewCustomCoder.
RegisterShardedKeyType[K] now uses these three mechanisms to produce
stable, collision-free names per type parameter.
Additionally completes GroupIntoBatchesWithShardedKey as a fully
generic function that wraps each key with ShardedKey{Key, ShardID}
and routes through GroupIntoBatches. End-to-end Prism test passes.
* Update batch.go
Fix staticcheck finding
* Update doc.go
* Update registry.go
* Update batch.go package comment
---------
Co-authored-by: Jack McCluskey
<[email protected]>
---
CHANGES.md | 10 +
sdks/go/pkg/beam/coder.go | 17 +
sdks/go/pkg/beam/core/graph/coder/coder.go | 116 ++++
sdks/go/pkg/beam/core/graph/coder/coder_test.go | 66 ++
sdks/go/pkg/beam/core/graph/coder/registry.go | 60 +-
.../pkg/beam/core/graph/coder/sharded_key_test.go | 81 +++
sdks/go/pkg/beam/core/runtime/exec/coder.go | 63 ++
sdks/go/pkg/beam/core/runtime/exec/coder_test.go | 84 +++
sdks/go/pkg/beam/core/runtime/graphx/coder.go | 21 +
sdks/go/pkg/beam/core/runtime/symbols.go | 20 +
sdks/go/pkg/beam/core/typex/class.go | 4 +-
sdks/go/pkg/beam/core/typex/fulltype.go | 23 +
sdks/go/pkg/beam/core/typex/special.go | 18 +-
sdks/go/pkg/beam/core/util/reflectx/call.go | 31 +
sdks/go/pkg/beam/pcollection.go | 16 +
sdks/go/pkg/beam/transforms/batch/batch.go | 677 +++++++++++++++++++++
.../pkg/beam/transforms/batch/batch_prism_test.go | 222 +++++++
sdks/go/pkg/beam/transforms/batch/batch_test.go | 47 ++
sdks/go/pkg/beam/transforms/batch/doc.go | 58 ++
sdks/go/pkg/beam/transforms/batch/size.go | 88 +++
sdks/go/pkg/beam/transforms/batch/size_test.go | 91 +++
21 files changed, 1807 insertions(+), 6 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index bda50ac6cd1..9cbbe1c207f 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -70,6 +70,16 @@
## New Features / Improvements
+* X feature added (Java/Python)
([#X](https://github.com/apache/beam/issues/X)).
+* 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 depend [...]
diff --git a/sdks/go/pkg/beam/coder.go b/sdks/go/pkg/beam/coder.go
index b03b739ed7b..c38a8e37ecc 100644
--- a/sdks/go/pkg/beam/coder.go
+++ b/sdks/go/pkg/beam/coder.go
@@ -89,6 +89,21 @@ func (c Coder) String() string {
return c.coder.String()
}
+// IsDeterministic reports whether this coder produces a byte-deterministic
+// encoding: encoding two equal values always yields identical byte
+// sequences.
+//
+// Determinism is required for any coder used as a state key in a stateful
+// DoFn or as the key component of a KV consumed by GroupByKey /
+// GroupIntoBatches. A non-deterministic key coder would silently corrupt
+// state keying, splintering state across apparently-distinct keys.
+func (c Coder) IsDeterministic() bool {
+ if c.coder == nil {
+ return false
+ }
+ return c.coder.IsDeterministic()
+}
+
// NewElementEncoder returns a new encoding function for the given type.
func NewElementEncoder(t reflect.Type) ElementEncoder {
c, err := inferCoder(typex.New(t))
@@ -249,6 +264,8 @@ func inferCoder(t FullType) (*coder.Coder, error) {
// are non-windowed? We either need to know the
windowing strategy or
// we should remove this case.
return &coder.Coder{Kind: coder.WindowedValue, T: t,
Components: c, Window: coder.NewGlobalWindow()}, nil
+ case typex.ShardedKeyType:
+ return &coder.Coder{Kind: coder.ShardedKey, T: t,
Components: c}, nil
default:
panic(fmt.Sprintf("Unexpected composite type: %v", t))
diff --git a/sdks/go/pkg/beam/core/graph/coder/coder.go
b/sdks/go/pkg/beam/core/graph/coder/coder.go
index 28e235860bd..f5f7aa2d757 100644
--- a/sdks/go/pkg/beam/core/graph/coder/coder.go
+++ b/sdks/go/pkg/beam/core/graph/coder/coder.go
@@ -84,6 +84,18 @@ func (c *CustomCoder) String() string {
return fmt.Sprintf("%v[%v;%v]", c.Type, c.Name, c.ID)
}
+// IsDeterministic reports whether this CustomCoder produces a deterministic
+// encoding. A CustomCoder is deterministic iff the user opted in by
+// registering the coder via RegisterDeterministicCoder. Default is false
+// (conservative): a non-deterministic key coder would silently corrupt state
+// keying in stateful DoFns.
+func (c *CustomCoder) IsDeterministic() bool {
+ if c == nil {
+ return false
+ }
+ return isCustomCoderDeterministic(c.Type)
+}
+
// Type signatures of encode/decode for verification.
var (
encodeSig = &funcx.Signature{
@@ -156,6 +168,20 @@ func NewCustomCoder(id string, t reflect.Type, encode,
decode any) (*CustomCoder
return c, nil
}
+// NewCustomCoderWithFuncs creates a CustomCoder from pre-wrapped
+// reflectx.Func values. This allows the caller to control the Name()
+// returned by each function — critical for closures inside Go generic
+// functions where the compiler assigns identical names to different
+// type instantiations.
+func NewCustomCoderWithFuncs(id string, t reflect.Type, enc, dec *funcx.Fn)
*CustomCoder {
+ return &CustomCoder{
+ Name: id,
+ Type: t,
+ Enc: enc,
+ Dec: dec,
+ }
+}
+
// Kind represents the type of coder used.
type Kind string
@@ -195,6 +221,17 @@ const (
//
// TODO(https://github.com/apache/beam/issues/18032): once this JIRA is
done, this coder should become the new thing.
CoGBK Kind = "CoGBK"
+
+ // ShardedKey encodes a user key wrapped with an opaque shard
identifier,
+ // used by GroupIntoBatchesWithShardedKey to distribute a single logical
+ // key's processing across workers. Wire format
+ // (beam:coder:sharded_key:v1):
+ //
+ // ByteArrayCoder.encode(shardId) ++ keyCoder.encode(key)
+ //
+ // matching sdks/java/core ShardedKey and the Python sharded_key
+ // encoding for cross-SDK interoperability.
+ ShardedKey Kind = "SK"
)
// Coder is a description of how to encode and decode values of a given type.
@@ -273,6 +310,62 @@ func (c *Coder) String() string {
return ret
}
+// IsDeterministic reports whether this Coder produces a deterministic
+// byte encoding — i.e. encoding two equal values always yields identical
+// byte sequences.
+//
+// Determinism is a prerequisite for any Coder used as a state key in a
+// stateful DoFn, as the key component of a KV consumed by GroupByKey, or as
+// a grouping key in a CoGroupByKey. A non-deterministic key coder causes
+// state-keyed operations to silently corrupt: two encodings of the same
+// logical key map to distinct physical keys, splintering state across
+// apparently-distinct keys.
+//
+// Built-in coders for primitive types (bytes, bool, varint, double,
+// string) are deterministic. Composite coders (KV, Iterable, Nullable)
+// are deterministic iff every component is. The Map coder is
+// non-deterministic because Go map iteration order is unspecified.
+// Custom user-registered coders are non-deterministic by default; users
+// opt in by registering with RegisterDeterministicCoder.
+func (c *Coder) IsDeterministic() bool {
+ if c == nil {
+ return false
+ }
+ switch c.Kind {
+ case Bytes, Bool, VarInt, Double, String:
+ return true
+ case Custom:
+ return c.Custom.IsDeterministic()
+ case KV, CoGBK, Nullable, Iterable, LP, ShardedKey:
+ for _, comp := range c.Components {
+ if !comp.IsDeterministic() {
+ return false
+ }
+ }
+ return true
+ case WindowedValue, ParamWindowedValue, Window, Timer, PaneInfo, IW:
+ // These coders are structural: they wrap runner/window
bookkeeping that is
+ // not used as a state key. Recurse into the data component
when present so
+ // that a non-deterministic inner coder is still reported.
+ for _, comp := range c.Components {
+ if !comp.IsDeterministic() {
+ return false
+ }
+ }
+ return true
+ case Row:
+ // Schema (row) coding encodes fields in a fixed field-id order
and
+ // produces a stable byte layout; however, row coders may
contain fields
+ // backed by custom coders we cannot introspect here.
Conservative
+ // default: return false and allow users to opt in via
schema-level
+ // determinism guarantees once they're exposed. Structs wanting
+ // deterministic behavior can register a deterministic custom
coder
+ // instead.
+ return false
+ }
+ return false
+}
+
// NewBytes returns a new []byte coder using the built-in scheme. It
// is always nested, for now.
func NewBytes() *Coder {
@@ -428,6 +521,29 @@ func NewCoGBK(components []*Coder) *Coder {
}
}
+// NewSK returns a coder for ShardedKey-typed values. The component
+// keyCoder encodes the user key; the ShardID is encoded as a
+// length-prefixed byte string preceding it (beam:coder:sharded_key:v1).
+//
+// The resulting FullType root is typex.ShardedKeyType with the key's
+// FullType as the single component, following the same Composite
+// pattern as KV.
+func NewSK(keyCoder *Coder) *Coder {
+ if keyCoder == nil {
+ panic("NewSK: keyCoder must not be nil")
+ }
+ return &Coder{
+ Kind: ShardedKey,
+ T: typex.New(typex.ShardedKeyType, keyCoder.T),
+ Components: []*Coder{keyCoder},
+ }
+}
+
+// IsSK returns true iff the coder is for a ShardedKey.
+func IsSK(c *Coder) bool {
+ return c != nil && c.Kind == ShardedKey
+}
+
// SkipW returns the data coder used by a WindowedValue, or returns the coder.
This
// allows code to seamlessly traverse WindowedValues without additional
conditional
// code.
diff --git a/sdks/go/pkg/beam/core/graph/coder/coder_test.go
b/sdks/go/pkg/beam/core/graph/coder/coder_test.go
index 040a0402c85..b60cbd72848 100644
--- a/sdks/go/pkg/beam/core/graph/coder/coder_test.go
+++ b/sdks/go/pkg/beam/core/graph/coder/coder_test.go
@@ -578,6 +578,72 @@ func TestNewNullable(t *testing.T) {
}
}
+func TestCoder_IsDeterministic(t *testing.T) {
+ ints := NewVarInt()
+ bytes := NewBytes()
+ bools := NewBool()
+ doubles := NewDouble()
+ strs := NewString()
+
+ enc := func(string) []byte { return nil }
+ dec := func([]byte) string { return "" }
+
+ nonDetCustom, err := NewCustomCoder("nonDet", reflectx.String, enc, dec)
+ if err != nil {
+ t.Fatal(err)
+ }
+ nonDetC := &Coder{Kind: Custom, Custom: nonDetCustom, T:
typex.New(reflectx.String)}
+
+ // Register a deterministic custom coder for a dedicated type.
+ type detType struct{}
+ detT := reflect.TypeOf((*detType)(nil)).Elem()
+ detEnc := func(detType) []byte { return nil }
+ detDec := func([]byte) detType { return detType{} }
+ RegisterDeterministicCoder(detT, detEnc, detDec)
+ detCustom, err := NewCustomCoder("det", detT, detEnc, detDec)
+ if err != nil {
+ t.Fatal(err)
+ }
+ detC := &Coder{Kind: Custom, Custom: detCustom, T: typex.New(detT)}
+
+ tests := []struct {
+ name string
+ c *Coder
+ want bool
+ }{
+ {"nil", nil, false},
+ {"bytes", bytes, true},
+ {"bool", bools, true},
+ {"varint", ints, true},
+ {"double", doubles, true},
+ {"string", strs, true},
+ {"nonDetCustom", nonDetC, false},
+ {"detCustom", detC, true},
+ {"KV_bytes_varint", NewKV([]*Coder{bytes, ints}), true},
+ {"KV_bytes_nonDet", NewKV([]*Coder{bytes, nonDetC}), false},
+ {"KV_nonDet_bytes", NewKV([]*Coder{nonDetC, bytes}), false},
+ {"iterable_varint", NewI(ints), true},
+ {"iterable_nonDet", NewI(nonDetC), false},
+ {"nullable_string", NewN(strs), true},
+ {"nullable_nonDet", NewN(nonDetC), false},
+ {"CoGBK_bytes_varint", NewCoGBK([]*Coder{bytes, ints}), true},
+ {"CoGBK_nonDet_varint", NewCoGBK([]*Coder{nonDetC, ints}),
false},
+ {"WindowedValue_varint", NewW(ints, NewGlobalWindow()), true},
+ {"WindowedValue_nonDet", NewW(nonDetC, NewGlobalWindow()),
false},
+ {"Row",
NewR(typex.New(reflect.TypeOf((*namedTypeForTest)(nil)))), false},
+ }
+
+ for _, test := range tests {
+ test := test
+ t.Run(test.name, func(t *testing.T) {
+ got := test.c.IsDeterministic()
+ if got != test.want {
+ t.Errorf("IsDeterministic(%v) = %v, want %v",
test.c, got, test.want)
+ }
+ })
+ }
+}
+
func TestNewCoGBK(t *testing.T) {
bytes := NewBytes()
ints := NewVarInt()
diff --git a/sdks/go/pkg/beam/core/graph/coder/registry.go
b/sdks/go/pkg/beam/core/graph/coder/registry.go
index f6677071b86..05d211898df 100644
--- a/sdks/go/pkg/beam/core/graph/coder/registry.go
+++ b/sdks/go/pkg/beam/core/graph/coder/registry.go
@@ -18,12 +18,14 @@ package coder
import (
"reflect"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/funcx"
"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
)
var (
- coderRegistry = make(map[reflect.Type]func(reflect.Type)
*CustomCoder)
- interfaceOrdering []reflect.Type
+ coderRegistry = make(map[reflect.Type]func(reflect.Type)
*CustomCoder)
+ interfaceOrdering []reflect.Type
+ deterministicRegistry = make(map[reflect.Type]bool)
)
// RegisterCoder registers a user defined coder for a given type, and will
@@ -76,6 +78,60 @@ func RegisterCoder(t reflect.Type, enc, dec any) {
}
}
+// RegisterDeterministicCoderWithFuncs is like RegisterDeterministicCoder
+// but accepts pre-wrapped reflectx.Func values (typically built via
+// reflectx.MakeFuncWithName) so the caller controls the function name
+// used during cross-worker serialization. This is required for
+// closures inside Go generic functions where different type
+// instantiations produce closures with the same compiler name.
+func RegisterDeterministicCoderWithFuncs(t reflect.Type, encFn, decFn
*funcx.Fn) {
+ name := t.String()
+ coderRegistry[t] = func(rt reflect.Type) *CustomCoder {
+ return NewCustomCoderWithFuncs(name, rt, encFn, decFn)
+ }
+ deterministicRegistry[t] = true
+}
+
+// RegisterDeterministicCoder is the deterministic-affirming counterpart to
+// RegisterCoder: it registers the (enc, dec) pair for t AND records that the
+// resulting CustomCoder produces a deterministic encoding. The caller asserts
+// by calling this function that enc produces byte-identical output for any
+// two equal input values of type t.
+//
+// Deterministic coders are required for any type used as a state key in a
+// stateful DoFn, as the key of a KV consumed by GroupByKey / GroupIntoBatches,
+// or as a grouping key for CoGroupByKey.
+//
+// Prefer this over RegisterCoder whenever the encoded type may be used as a
+// key. For types that cannot guarantee determinism (e.g. encodings backed by
+// map[K]V iteration order), use the plain RegisterCoder.
+func RegisterDeterministicCoder(t reflect.Type, enc, dec any) {
+ RegisterCoder(t, enc, dec)
+ deterministicRegistry[t] = true
+}
+
+// isCustomCoderDeterministic returns true iff t has been registered via
+// RegisterDeterministicCoder.
+func isCustomCoderDeterministic(t reflect.Type) bool {
+ if t == nil {
+ return false
+ }
+ if ok, present := deterministicRegistry[t]; present {
+ return ok
+ }
+ // Also match against interface registrations: if the type implements a
+ // registered-deterministic interface, honor that.
+ for rt, det := range deterministicRegistry {
+ if !det {
+ continue
+ }
+ if rt.Kind() == reflect.Interface && t.Implements(rt) {
+ return true
+ }
+ }
+ return false
+}
+
// LookupCustomCoder returns the custom coder for the type if any,
// first checking for a specific matching type, and then iterating
// through registered interface coders in reverse registration order.
diff --git a/sdks/go/pkg/beam/core/graph/coder/sharded_key_test.go
b/sdks/go/pkg/beam/core/graph/coder/sharded_key_test.go
new file mode 100644
index 00000000000..fc9b93b0070
--- /dev/null
+++ b/sdks/go/pkg/beam/core/graph/coder/sharded_key_test.go
@@ -0,0 +1,81 @@
+// 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 coder
+
+import (
+ "reflect"
+ "testing"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
+)
+
+func TestNewSK(t *testing.T) {
+ t.Run("nilKeyCoder_panics", func(t *testing.T) {
+ defer func() {
+ if p := recover(); p == nil {
+ t.Fatal("expected panic on nil keyCoder, got
none")
+ }
+ }()
+ NewSK(nil)
+ })
+
+ t.Run("valid_string_key", func(t *testing.T) {
+ sk := NewSK(NewString())
+ if sk.Kind != ShardedKey {
+ t.Fatalf("Kind = %v, want %v", sk.Kind, ShardedKey)
+ }
+ if !IsSK(sk) {
+ t.Fatalf("IsSK(%v) = false, want true", sk)
+ }
+ if len(sk.Components) != 1 {
+ t.Fatalf("Components = %d, want 1", len(sk.Components))
+ }
+ if sk.Components[0].Kind != String {
+ t.Fatalf("Components[0].Kind = %v, want %v",
sk.Components[0].Kind, String)
+ }
+ if sk.T.Type() != typex.ShardedKeyType {
+ t.Fatalf("T.Type() = %v, want %v", sk.T.Type(),
typex.ShardedKeyType)
+ }
+ })
+
+ t.Run("nested_composite_panics", func(t *testing.T) {
+ defer func() {
+ if p := recover(); p == nil {
+ t.Fatal("expected panic on nested composite
key, got none")
+ }
+ }()
+ // KV components inside a ShardedKey key are disallowed by
fulltype.New.
+ NewSK(NewKV([]*Coder{NewString(), NewBytes()}))
+ })
+}
+
+func TestSK_IsDeterministic(t *testing.T) {
+ detSK := NewSK(NewString())
+ if !detSK.IsDeterministic() {
+ t.Errorf("ShardedKey<string>.IsDeterministic() = false, want
true")
+ }
+
+ nonDet, err := NewCustomCoder("nonDet", reflect.TypeOf(""),
+ func(string) []byte { return nil }, func([]byte) string {
return "" })
+ if err != nil {
+ t.Fatal(err)
+ }
+ nonDetC := &Coder{Kind: Custom, Custom: nonDet, T:
typex.New(reflect.TypeOf(""))}
+ nonDetSK := NewSK(nonDetC)
+ if nonDetSK.IsDeterministic() {
+ t.Errorf("ShardedKey<nonDet>.IsDeterministic() = true, want
false")
+ }
+}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/coder.go
b/sdks/go/pkg/beam/core/runtime/exec/coder.go
index 2c21ebea56b..b6894335538 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/coder.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/coder.go
@@ -166,6 +166,11 @@ func MakeElementEncoder(c *coder.Coder) ElementEncoder {
be: boolEncoder{},
}
+ case coder.ShardedKey:
+ return &shardedKeyEncoder{
+ key: MakeElementEncoder(c.Components[0]),
+ }
+
default:
panic(fmt.Sprintf("Unexpected coder: %v", c))
}
@@ -288,6 +293,11 @@ func MakeElementDecoder(c *coder.Coder) ElementDecoder {
bd: boolDecoder{},
}
+ case coder.ShardedKey:
+ return &shardedKeyDecoder{
+ key: MakeElementDecoder(c.Components[0]),
+ }
+
default:
panic(fmt.Sprintf("Unexpected coder: %v", c))
}
@@ -1356,3 +1366,56 @@ func decodeTimer(dec ElementDecoder, win WindowDecoder,
r io.Reader) (TimerRecv,
return tm, nil
}
+
+// shardedKeyEncoder encodes ShardedKey-typed values in the standard
+// beam:coder:sharded_key:v1 wire format:
+//
+// ByteArrayCoder.encode(ShardID) ++ keyCoder.encode(Key)
+//
+// Runtime values are carried by a FullValue whose Elm holds the user key
+// and whose Elm2 holds the []byte shard identifier — the same two-part
+// convention used by the KV coder. This matches the Java
+// util.ShardedKey.Coder and Python sharded_key encodings exactly; any
+// divergence of a single byte would silently corrupt cross-SDK pipelines.
+type shardedKeyEncoder struct {
+ key ElementEncoder
+}
+
+func (e *shardedKeyEncoder) Encode(val *FullValue, w io.Writer) error {
+ shardID, ok := val.Elm2.([]byte)
+ if !ok {
+ return errors.Errorf(
+ "shardedKeyEncoder: Elm2 must be []byte shardID (got
%T)", val.Elm2)
+ }
+ if err := coder.EncodeBytes(shardID, w); err != nil {
+ return errors.WithContext(err, "shardedKeyEncoder: shardID")
+ }
+ return e.key.Encode(&FullValue{Elm: val.Elm}, w)
+}
+
+// shardedKeyDecoder is the inverse of shardedKeyEncoder. Decoded values
+// are placed in FullValue{Elm: key, Elm2: shardID}.
+type shardedKeyDecoder struct {
+ key ElementDecoder
+}
+
+func (d *shardedKeyDecoder) DecodeTo(r io.Reader, fv *FullValue) error {
+ shardID, err := coder.DecodeBytes(r)
+ if err != nil {
+ return errors.WithContext(err, "shardedKeyDecoder: shardID")
+ }
+ keyFV, err := d.key.Decode(r)
+ if err != nil {
+ return errors.WithContext(err, "shardedKeyDecoder: key")
+ }
+ *fv = FullValue{Elm: keyFV.Elm, Elm2: shardID}
+ return nil
+}
+
+func (d *shardedKeyDecoder) Decode(r io.Reader) (*FullValue, error) {
+ fv := &FullValue{}
+ if err := d.DecodeTo(r, fv); err != nil {
+ return nil, err
+ }
+ return fv, nil
+}
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 75d18e533cf..155fd72776e 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/coder_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/coder_test.go
@@ -158,6 +158,90 @@ func compareFV(t *testing.T, got *FullValue, want
*FullValue) {
}
}
+// TestShardedKeyCoder_WireFormat verifies the exact bytes produced by the
+// ShardedKey coder against the standard_coders.yaml fixtures (lines
+// 501-521, urn "beam:coder:sharded_key:v1" with a string_utf8 key
+// component). A single divergent byte would silently corrupt cross-SDK
+// pipelines on Dataflow / Flink.
+func TestShardedKeyCoder_WireFormat(t *testing.T) {
+ c := coder.NewSK(coder.NewString())
+ enc := MakeElementEncoder(c)
+ dec := MakeElementDecoder(c)
+
+ type fixture struct {
+ name string
+ key string
+ shardID []byte
+ wire []byte
+ }
+ fixtures := []fixture{
+ {
+ name: "empty_empty",
+ key: "",
+ shardID: []byte{},
+ wire: []byte{0x00, 0x00},
+ },
+ {
+ name: "shardId_emptyKey",
+ key: "",
+ shardID: []byte("shard_id"),
+ wire: append(
+ append([]byte{0x08}, []byte("shard_id")...),
+ 0x00,
+ ),
+ },
+ {
+ name: "shardId_key",
+ key: "key",
+ shardID: []byte("shard_id"),
+ wire: append(
+ append([]byte{0x08}, []byte("shard_id")...),
+ append([]byte{0x03}, []byte("key")...)...,
+ ),
+ },
+ {
+ name: "emptyShardId_key",
+ key: "key",
+ shardID: []byte{},
+ wire: append([]byte{0x00, 0x03}, []byte("key")...),
+ },
+ }
+
+ for _, f := range fixtures {
+ f := f
+ t.Run(f.name, func(t *testing.T) {
+ var buf bytes.Buffer
+ // ShardedKey values are carried as FullValue{Elm: key,
Elm2: shardID}.
+ if err := enc.Encode(&FullValue{Elm: f.key, Elm2:
f.shardID}, &buf); err != nil {
+ t.Fatalf("Encode: %v", err)
+ }
+ if got := buf.Bytes(); !bytes.Equal(got, f.wire) {
+ t.Fatalf("Encode: got bytes %#v, want %#v",
got, f.wire)
+ }
+
+ fv, err := dec.Decode(bytes.NewReader(f.wire))
+ if err != nil {
+ t.Fatalf("Decode: %v", err)
+ }
+ gotKey, ok := fv.Elm.(string)
+ if !ok {
+ t.Fatalf("Decode Elm: got %T, want string",
fv.Elm)
+ }
+ if gotKey != f.key {
+ t.Errorf("Decode Elm: got %q, want %q", gotKey,
f.key)
+ }
+ gotShard, ok := fv.Elm2.([]byte)
+ if !ok {
+ t.Fatalf("Decode Elm2: got %T, want []byte",
fv.Elm2)
+ }
+ // Both sides "empty" — accept nil or zero-length slice
equivalence.
+ if len(gotShard) != len(f.shardID) || (len(gotShard) >
0 && !bytes.Equal(gotShard, f.shardID)) {
+ t.Errorf("Decode Elm2: got %#v, want %#v",
gotShard, f.shardID)
+ }
+ })
+ }
+}
+
func TestIterableCoder(t *testing.T) {
cod := coder.NewI(coder.NewVarInt())
wantVals := []int64{8, 24, 72}
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/coder.go
b/sdks/go/pkg/beam/core/runtime/graphx/coder.go
index 2b769c873ec..ced4d34679b 100644
--- a/sdks/go/pkg/beam/core/runtime/graphx/coder.go
+++ b/sdks/go/pkg/beam/core/runtime/graphx/coder.go
@@ -47,6 +47,7 @@ const (
urnTimerCoder = "beam:coder:timer:v1"
urnRowCoder = "beam:coder:row:v1"
urnNullableCoder = "beam:coder:nullable:v1"
+ urnShardedKeyCoder = "beam:coder:sharded_key:v1"
urnGlobalWindow = "beam:coder:global_window:v1"
urnIntervalWindow = "beam:coder:interval_window:v1"
@@ -74,6 +75,7 @@ func knownStandardCoders() []string {
urnRowCoder,
urnNullableCoder,
urnTimerCoder,
+ urnShardedKeyCoder,
}
}
@@ -378,6 +380,15 @@ func (b *CoderUnmarshaller) makeCoder(id string, c
*pipepb.Coder) (*coder.Coder,
return nil, err
}
return coder.NewN(elm), nil
+ case urnShardedKeyCoder:
+ if len(components) != 1 {
+ return nil, errors.Errorf("could not unmarshal
sharded_key coder from %v, expected one component (key) but got %d", c,
len(components))
+ }
+ keyC, err := b.Coder(components[0])
+ if err != nil {
+ return nil, err
+ }
+ return coder.NewSK(keyC), nil
case urnIntervalWindow:
return coder.NewIntervalWindowCoder(), nil
@@ -493,6 +504,16 @@ func (b *CoderMarshaller) Add(c *coder.Coder) (string,
error) {
stream := b.internBuiltInCoder(urnIterableCoder, value)
return b.internBuiltInCoder(urnKVCoder, comp[0], stream), nil
+ case coder.ShardedKey:
+ comp, err := b.AddMulti(c.Components)
+ if err != nil {
+ return "", errors.Wrapf(err, "failed to marshal
ShardedKey coder %v", c)
+ }
+ if len(comp) != 1 {
+ return "", errors.Errorf("ShardedKey coder requires
exactly 1 component (key), got %d", len(comp))
+ }
+ return b.internBuiltInCoder(urnShardedKeyCoder, comp...), nil
+
case coder.WindowedValue:
comp := []string{}
if ids, err := b.AddMulti(c.Components); err != nil {
diff --git a/sdks/go/pkg/beam/core/runtime/symbols.go
b/sdks/go/pkg/beam/core/runtime/symbols.go
index 84afe9b769a..9640af288b6 100644
--- a/sdks/go/pkg/beam/core/runtime/symbols.go
+++ b/sdks/go/pkg/beam/core/runtime/symbols.go
@@ -83,6 +83,26 @@ func RegisterFunction(fn any) {
cache[key] = fn
}
+// RegisterFunctionWithName registers fn under the given name,
+// overriding the automatically derived symbol name. This is necessary
+// for closures produced by Go generic functions where multiple type
+// instantiations generate closures with the same compiler-assigned
+// name (e.g. "pkg.Func[...].func1") — without distinct names the
+// last registration wins and cross-worker deserialization resolves
+// the wrong function.
+//
+// Callers must ensure that name is stable across process invocations
+// (pipeline driver and workers must agree). A typical choice is
+// "<package>.<GenericFunc>[<TypeParam>].enc".
+//
+// Must be called in init() only.
+func RegisterFunctionWithName(name string, fn any) {
+ if initialized {
+ panic("Init hooks have already run. Register function during
init() instead.")
+ }
+ cache[name] = fn
+}
+
// ResolveFunction resolves the runtime value of a given function by symbol
name
// and type.
func ResolveFunction(name string, t reflect.Type) (any, error) {
diff --git a/sdks/go/pkg/beam/core/typex/class.go
b/sdks/go/pkg/beam/core/typex/class.go
index 570b7e27921..6c8f3549893 100644
--- a/sdks/go/pkg/beam/core/typex/class.go
+++ b/sdks/go/pkg/beam/core/typex/class.go
@@ -231,10 +231,10 @@ func IsUniversal(t reflect.Type) bool {
}
// IsComposite returns true iff the given type is one of the predefined
-// Composite marker types: KV, CoGBK or WindowedValue.
+// Composite marker types: KV, CoGBK, WindowedValue, Timers or ShardedKey.
func IsComposite(t reflect.Type) bool {
switch t {
- case KVType, CoGBKType, WindowedValueType, TimersType:
+ case KVType, CoGBKType, WindowedValueType, TimersType, ShardedKeyType:
return true
default:
return false
diff --git a/sdks/go/pkg/beam/core/typex/fulltype.go
b/sdks/go/pkg/beam/core/typex/fulltype.go
index ff5520c2861..88e26568dff 100644
--- a/sdks/go/pkg/beam/core/typex/fulltype.go
+++ b/sdks/go/pkg/beam/core/typex/fulltype.go
@@ -89,6 +89,8 @@ func printShortComposite(t reflect.Type) string {
return "KV"
case NullableType:
return "Nullable"
+ case ShardedKeyType:
+ return "SK"
default:
return fmt.Sprintf("invalid(%v)", t)
}
@@ -146,6 +148,14 @@ func New(t reflect.Type, components ...FullType) FullType {
return &tree{class, t, components}
case TimersType:
return &tree{class, t, components}
+ case ShardedKeyType:
+ if len(components) != 1 {
+ panic(fmt.Sprintf("Invalid number of components
for ShardedKey: %v, %v", t, components))
+ }
+ if components[0].Class() == Composite {
+ panic(fmt.Sprintf("Invalid to nest composite
inside ShardedKey: %v, %v", t, components))
+ }
+ return &tree{class, t, components}
default:
panic(fmt.Sprintf("Unexpected composite type: %v", t))
}
@@ -226,6 +236,19 @@ func NewCoGBK(components ...FullType) FullType {
return New(CoGBKType, components...)
}
+// IsShardedKey returns true iff the type is a ShardedKey.
+func IsShardedKey(t FullType) bool {
+ return t.Type() == ShardedKeyType
+}
+
+// NewShardedKey constructs a new ShardedKey FullType wrapping the given
+// key component. The ShardedKey has exactly one component — the user key
+// type — because the ShardID byte-string has a fixed representation and
+// is not a user-configurable type.
+func NewShardedKey(keyType FullType) FullType {
+ return New(ShardedKeyType, keyType)
+}
+
// IsStructurallyAssignable returns true iff a from value is structurally
// assignable to the to value of the given types. Types that are
// "structurally assignable" (SA) are assignable if type variables are
diff --git a/sdks/go/pkg/beam/core/typex/special.go
b/sdks/go/pkg/beam/core/typex/special.go
index 9093ddc782c..6cf1cb99f75 100644
--- a/sdks/go/pkg/beam/core/typex/special.go
+++ b/sdks/go/pkg/beam/core/typex/special.go
@@ -44,6 +44,7 @@ var (
CoGBKType = reflect.TypeOf((*CoGBK)(nil)).Elem()
WindowedValueType = reflect.TypeOf((*WindowedValue)(nil)).Elem()
BundleFinalizationType =
reflect.TypeOf((*BundleFinalization)(nil)).Elem()
+ ShardedKeyType = reflect.TypeOf((*ShardedKey)(nil)).Elem()
)
// T, U, V, W, X, Y, Z are universal types. They play the role of generic
@@ -128,8 +129,10 @@ type Timers struct {
Pane PaneInfo
}
-// KV, Nullable, CoGBK, WindowedValue represent composite generic types. They
are not used
-// directly in user code signatures, but only in FullTypes.
+// KV, Nullable, CoGBK, WindowedValue, ShardedKey represent composite
+// generic types. They are not used directly in user code signatures, but
+// only in FullTypes — each appears as the root of a FullType tree whose
+// component list holds the concrete sub-types.
type KV struct{}
@@ -138,3 +141,14 @@ type Nullable struct{}
type CoGBK struct{}
type WindowedValue struct{}
+
+// ShardedKey is the composite marker for sharded-key encoded pairs
+// (user key + opaque shard identifier). It is never constructed by user
+// code; it appears only as the root of a FullType tree whose single
+// component is the key's FullType.
+//
+// Runtime values are carried through FullValue.Elm (user key) and
+// FullValue.Elm2 ([]byte shardID). The corresponding wire encoding is
+// URN beam:coder:sharded_key:v1, byte-identical to the Java and Python
+// sharded_key encodings.
+type ShardedKey struct{}
diff --git a/sdks/go/pkg/beam/core/util/reflectx/call.go
b/sdks/go/pkg/beam/core/util/reflectx/call.go
index 9b1955427f7..e14ed016425 100644
--- a/sdks/go/pkg/beam/core/util/reflectx/call.go
+++ b/sdks/go/pkg/beam/core/util/reflectx/call.go
@@ -87,6 +87,37 @@ func (c *reflectFunc) Call(args []any) []any {
return Interface(c.fn.Call(ValueOf(args)))
}
+// MakeFuncWithName returns a Func that wraps fn but whose Name()
+// returns the provided name instead of the compiler-derived symbol.
+// This is essential for closures inside Go generic functions: all
+// type instantiations produce closures with the same compiler name
+// (e.g. "pkg.Func[...].func1"), so the default name-based
+// serialization cannot distinguish them. A stable, type-qualified
+// name ensures cross-worker deserialization resolves the correct
+// function.
+func MakeFuncWithName(name string, fn any) Func {
+ inner := MakeFunc(fn)
+ return &namedFunc{inner: inner, name: name}
+}
+
+type namedFunc struct {
+ inner Func
+ name string
+}
+
+func (f *namedFunc) Name() string { return f.name }
+func (f *namedFunc) Type() reflect.Type { return f.inner.Type() }
+func (f *namedFunc) Call(args []any) []any { return f.inner.Call(args) }
+
+// Interface returns the original unwrapped function, which
+// runtime.RegisterFunction needs for pointer extraction.
+func (f *namedFunc) Interface() any {
+ if rf, ok := f.inner.(*reflectFunc); ok {
+ return rf.fn.Interface()
+ }
+ return nil
+}
+
// CallNoPanic calls the given Func and catches any panic.
func CallNoPanic(fn Func, args []any) (ret []any, err error) {
defer func() {
diff --git a/sdks/go/pkg/beam/pcollection.go b/sdks/go/pkg/beam/pcollection.go
index e5dc63289f3..2138266a667 100644
--- a/sdks/go/pkg/beam/pcollection.go
+++ b/sdks/go/pkg/beam/pcollection.go
@@ -17,6 +17,7 @@ package beam
import (
"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/typex"
"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
)
@@ -80,6 +81,21 @@ func (p PCollection) SetCoder(c Coder) error {
return nil
}
+// WindowingStrategy returns the windowing strategy of the PCollection. It
+// describes how elements are assigned to windows and — for transforms that
+// honor it — the allowed lateness after which windows are closed.
+//
+// Transforms that use state and timers keyed by window, such as
+// GroupIntoBatches, consult this strategy to compute end-of-window
+// event-time timers and to bound partial-batch flushes by the pipeline's
+// allowed lateness.
+func (p PCollection) WindowingStrategy() *window.WindowingStrategy {
+ if !p.IsValid() {
+ panic("Invalid PCollection")
+ }
+ return p.n.WindowingStrategy()
+}
+
func (p PCollection) String() string {
if !p.IsValid() {
return "(invalid)"
diff --git a/sdks/go/pkg/beam/transforms/batch/batch.go
b/sdks/go/pkg/beam/transforms/batch/batch.go
new file mode 100644
index 00000000000..6cc378c2f44
--- /dev/null
+++ b/sdks/go/pkg/beam/transforms/batch/batch.go
@@ -0,0 +1,677 @@
+// 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 batch provides transforms that group elements of a KV-keyed
+// PCollection into batches of a target size for downstream per-batch
+// processing (rate-limited API calls, bulk sinks, etc.).
+//
+// GroupIntoBatches mirrors the behavior of the Java and Python
+// transforms of the same name. GroupIntoBatchesWithShardedKey adds
+// opaque per-element shard identifiers to the keys so the processing
+// of a single hot logical key spreads across multiple workers.
+//
+// # Behavior
+//
+// Given a PCollection<KV<K, V>>, GroupIntoBatches buffers values per
+// key and emits batches as KV<K, []V> whenever one of the following
+// limits is reached:
+//
+// - len(batch) reaches BatchSize, OR
+// - sum of byte sizes reaches BatchSizeBytes, OR
+// - MaxBufferingDuration elapses in processing time since the first
+// element of the current batch (if set), OR
+// - the window advances past MaxTimestamp + AllowedLateness of the
+// input PCollection's WindowingStrategy.
+//
+// Elements of different windows are never combined into the same
+// batch.
+//
+// # Determinism requirement
+//
+// The key coder MUST be deterministic. State keying depends on
+// byte-stable encodings: a non-deterministic key coder would silently
+// split the logical key across multiple physical keys, producing
+// corrupt batches. The transform panics at pipeline build time if the
+// key coder is not known to be deterministic. For user-defined key
+// types, register the type's coder via
+// coder.RegisterDeterministicCoder.
+//
+// # Differences from Java/Python
+//
+// - BatchSize / BatchSizeBytes are int64 (parity with proto and Java
+// long, avoiding overflow on 32-bit platforms).
+// - BatchSizeBytes is limited to primitive value types ([]byte,
+// string, numeric, bool) in this release; opaque V types panic at
+// build time if BatchSizeBytes > 0.
+package batch
+
+import (
+ "bytes"
+ "context"
+ "encoding/binary"
+ "fmt"
+ "reflect"
+ "sync"
+ "sync/atomic"
+ "time"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/funcx"
+ beamcoder "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/state"
+ "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/reflectx"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/register"
+ "github.com/google/uuid"
+)
+
+// ShardedKey pairs a user key with an opaque shard identifier. It is
+// the key type of the PCollection produced by
+// GroupIntoBatchesWithShardedKey.
+type ShardedKey[K any] struct {
+ Key K
+ ShardID []byte
+}
+
+// RegisterShardedKeyType registers a ShardedKey[K] instantiation so
+// its coder survives cross-worker serialization. Common key types
+// (string, []byte, int, int64) are registered automatically at init.
+// Users of other K types must call this at init time.
+func RegisterShardedKeyType[K any]() {
+ var zero K
+ keyT := reflect.TypeOf(zero)
+ skT := reflect.TypeOf(ShardedKey[K]{})
+
+ register.DoFn3x0[K, typex.V, func(ShardedKey[K],
typex.V)](&wrapShardedKeyFn[K]{})
+ register.Emitter2[ShardedKey[K], typex.V]()
+ beam.RegisterType(skT)
+
+ keyEnc := beam.NewElementEncoder(keyT)
+ keyDec := beam.NewElementDecoder(keyT)
+
+ enc := func(sk ShardedKey[K]) []byte {
+ var buf bytes.Buffer
+ writeVarInt(&buf, int64(len(sk.ShardID)))
+ buf.Write(sk.ShardID)
+ if err := keyEnc.Encode(sk.Key, &buf); err != nil {
+ panic(err)
+ }
+ return buf.Bytes()
+ }
+ dec := func(b []byte) ShardedKey[K] {
+ r := bytes.NewReader(b)
+ n := readVarInt(r)
+ shardID := make([]byte, n)
+ if n > 0 {
+ if _, err := r.Read(shardID); err != nil {
+ panic(err)
+ }
+ }
+ k, err := keyDec.Decode(r)
+ if err != nil {
+ panic(err)
+ }
+ return ShardedKey[K]{Key: k.(K), ShardID: shardID}
+ }
+
+ // Closures inside generic functions share the same compiler
+ // symbol name for every type instantiation. We wrap them with a
+ // type-qualified name so the cross-worker deserializer resolves
+ // the correct enc/dec for each ShardedKey[K].
+ encName := fmt.Sprintf("batch.encShardedKey[%v]", keyT)
+ decName := fmt.Sprintf("batch.decShardedKey[%v]", keyT)
+
+ encFn := reflectx.MakeFuncWithName(encName, enc)
+ decFn := reflectx.MakeFuncWithName(decName, dec)
+
+ // Register in the runtime cache under the qualified name so
+ // ResolveFunction finds them at deserialization time.
+ runtime.RegisterFunctionWithName(encName, enc)
+ runtime.RegisterFunctionWithName(decName, dec)
+
+ encWrapped, err := funcx.New(encFn)
+ if err != nil {
+ panic(fmt.Sprintf("RegisterShardedKeyType: bad enc for %v: %v",
skT, err))
+ }
+ decWrapped, err := funcx.New(decFn)
+ if err != nil {
+ panic(fmt.Sprintf("RegisterShardedKeyType: bad dec for %v: %v",
skT, err))
+ }
+
+ beamcoder.RegisterDeterministicCoderWithFuncs(skT, encWrapped,
decWrapped)
+}
+
+// Params configures GroupIntoBatches and
+// GroupIntoBatchesWithShardedKey.
+//
+// At least one of BatchSize or BatchSizeBytes must be > 0.
+type Params struct {
+ // BatchSize is the target maximum number of elements per batch. A
+ // batch is emitted as soon as it holds BatchSize elements. Zero
+ // disables the count-based trigger.
+ BatchSize int64
+
+ // BatchSizeBytes is the target maximum cumulative byte size per
+ // batch. A batch is emitted as soon as adding another element
+ // would exceed BatchSizeBytes. Zero disables the byte-based
+ // trigger.
+ BatchSizeBytes int64
+
+ // MaxBufferingDuration, when > 0, triggers emission of a partial
+ // batch after this much processing time has elapsed since the
+ // first element of the current batch was buffered.
+ MaxBufferingDuration time.Duration
+}
+
+func (p Params) validate() error {
+ if p.BatchSize < 0 {
+ return fmt.Errorf("Params.BatchSize must be >= 0; got %d",
p.BatchSize)
+ }
+ if p.BatchSizeBytes < 0 {
+ return fmt.Errorf("Params.BatchSizeBytes must be >= 0; got %d",
p.BatchSizeBytes)
+ }
+ if p.BatchSize == 0 && p.BatchSizeBytes == 0 {
+ return fmt.Errorf("Params: at least one of BatchSize or
BatchSizeBytes must be > 0")
+ }
+ if p.MaxBufferingDuration < 0 {
+ return fmt.Errorf("Params.MaxBufferingDuration must be >= 0;
got %s", p.MaxBufferingDuration)
+ }
+ return nil
+}
+
+const (
+ sizerNone int32 = 0
+ sizerPrimitive int32 = 1
+)
+
+// codecCache keeps a per-value-type ElementEncoder/Decoder pair.
+type codecCache struct {
+ once sync.Once
+ enc beam.ElementEncoder
+ dec beam.ElementDecoder
+}
+
+func (c *codecCache) init(t reflect.Type) {
+ c.once.Do(func() {
+ c.enc = beam.NewElementEncoder(t)
+ c.dec = beam.NewElementDecoder(t)
+ })
+}
+
+func (c *codecCache) encode(v any) []byte {
+ var buf bytes.Buffer
+ if err := c.enc.Encode(v, &buf); err != nil {
+ panic(err)
+ }
+ return buf.Bytes()
+}
+
+func (c *codecCache) decode(b []byte) any {
+ v, err := c.dec.Decode(bytes.NewReader(b))
+ if err != nil {
+ panic(err)
+ }
+ return v
+}
+
+// groupIntoBatchesFn is the stateful DoFn without a processing-time
+// buffering timer.
+type groupIntoBatchesFn struct {
+ Buffer state.Bag[[]byte]
+ Count state.Value[int64]
+ ByteSize state.Value[int64]
+ WindowEnd timers.EventTime
+
+ ValueType beam.EncodedType
+
+ BatchSize int64
+ BatchSizeBytes int64
+ AllowedLatenessMs int64
+ SizerKind int32
+
+ codec codecCache
+}
+
+func (fn *groupIntoBatchesFn) ProcessElement(
+ w beam.Window, sp state.Provider, tp timers.Provider,
+ key typex.T, value typex.V, emit func(typex.T, []typex.V),
+) {
+ fn.codec.init(fn.ValueType.T)
+
+ count, _, err := fn.Count.Read(sp)
+ if err != nil {
+ panic(err)
+ }
+
+ if w.MaxTimestamp() < mtime.MaxTimestamp {
+ windowEnd := w.MaxTimestamp().ToTime()
+ if fn.AllowedLatenessMs > 0 {
+ windowEnd =
windowEnd.Add(time.Duration(fn.AllowedLatenessMs) * time.Millisecond)
+ }
+ fn.WindowEnd.Set(tp, windowEnd, timers.WithNoOutputTimestamp())
+ }
+
+ if err := fn.Buffer.Add(sp, fn.codec.encode(value)); err != nil {
+ panic(err)
+ }
+ count++
+ if err := fn.Count.Write(sp, count); err != nil {
+ panic(err)
+ }
+
+ newBytes := int64(0)
+ if fn.BatchSizeBytes > 0 {
+ cur, _, err := fn.ByteSize.Read(sp)
+ if err != nil {
+ panic(err)
+ }
+ cur += sizeOf(fn.SizerKind, value)
+ if err := fn.ByteSize.Write(sp, cur); err != nil {
+ panic(err)
+ }
+ newBytes = cur
+ }
+
+ if fn.BatchSize > 0 && count >= fn.BatchSize {
+ fn.flush(sp, key, emit)
+ return
+ }
+ if fn.BatchSizeBytes > 0 && newBytes >= fn.BatchSizeBytes {
+ fn.flush(sp, key, emit)
+ return
+ }
+}
+
+func (fn *groupIntoBatchesFn) OnTimer(
+ ctx context.Context, ts beam.EventTime, sp state.Provider, tp
timers.Provider,
+ key typex.T, timer timers.Context, emit func(typex.T, []typex.V),
+) {
+ if timer.Family != fn.WindowEnd.Family {
+ panic(fmt.Sprintf("batch.groupIntoBatchesFn: unexpected timer
family %q", timer.Family))
+ }
+ fn.codec.init(fn.ValueType.T)
+ fn.flush(sp, key, emit)
+}
+
+func (fn *groupIntoBatchesFn) flush(
+ sp state.Provider, key typex.T, emit func(typex.T, []typex.V),
+) {
+ buf, ok, err := fn.Buffer.Read(sp)
+ if err != nil {
+ panic(err)
+ }
+ if !ok || len(buf) == 0 {
+ return
+ }
+
+ out := make([]typex.V, len(buf))
+ for i, b := range buf {
+ out[i] = fn.codec.decode(b)
+ }
+ emit(key, out)
+
+ if err := fn.Buffer.Clear(sp); err != nil {
+ panic(err)
+ }
+ if err := fn.Count.Clear(sp); err != nil {
+ panic(err)
+ }
+ if fn.BatchSizeBytes > 0 {
+ if err := fn.ByteSize.Clear(sp); err != nil {
+ panic(err)
+ }
+ }
+}
+
+// groupIntoBatchesBufferedFn adds a processing-time buffering timer.
+type groupIntoBatchesBufferedFn struct {
+ Buffer state.Bag[[]byte]
+ Count state.Value[int64]
+ ByteSize state.Value[int64]
+ TimerSet state.Value[bool]
+ Buffering timers.ProcessingTime
+ WindowEnd timers.EventTime
+
+ ValueType beam.EncodedType
+
+ BatchSize int64
+ BatchSizeBytes int64
+ MaxBufferingMs int64
+ AllowedLatenessMs int64
+ SizerKind int32
+
+ codec codecCache
+}
+
+func (fn *groupIntoBatchesBufferedFn) ProcessElement(
+ w beam.Window, sp state.Provider, tp timers.Provider,
+ key typex.T, value typex.V, emit func(typex.T, []typex.V),
+) {
+ fn.codec.init(fn.ValueType.T)
+
+ count, _, err := fn.Count.Read(sp)
+ if err != nil {
+ panic(err)
+ }
+
+ if w.MaxTimestamp() < mtime.MaxTimestamp {
+ windowEnd := w.MaxTimestamp().ToTime()
+ if fn.AllowedLatenessMs > 0 {
+ windowEnd =
windowEnd.Add(time.Duration(fn.AllowedLatenessMs) * time.Millisecond)
+ }
+ fn.WindowEnd.Set(tp, windowEnd, timers.WithNoOutputTimestamp())
+ }
+
+ if err := fn.Buffer.Add(sp, fn.codec.encode(value)); err != nil {
+ panic(err)
+ }
+ count++
+ if err := fn.Count.Write(sp, count); err != nil {
+ panic(err)
+ }
+
+ newBytes := int64(0)
+ if fn.BatchSizeBytes > 0 {
+ cur, _, err := fn.ByteSize.Read(sp)
+ if err != nil {
+ panic(err)
+ }
+ cur += sizeOf(fn.SizerKind, value)
+ if err := fn.ByteSize.Write(sp, cur); err != nil {
+ panic(err)
+ }
+ newBytes = cur
+ }
+
+ if count == 1 {
+ fn.Buffering.Set(tp,
time.Now().Add(time.Duration(fn.MaxBufferingMs)*time.Millisecond))
+ if err := fn.TimerSet.Write(sp, true); err != nil {
+ panic(err)
+ }
+ }
+
+ if fn.BatchSize > 0 && count >= fn.BatchSize {
+ fn.flush(sp, tp, key, emit)
+ return
+ }
+ if fn.BatchSizeBytes > 0 && newBytes >= fn.BatchSizeBytes {
+ fn.flush(sp, tp, key, emit)
+ return
+ }
+}
+
+func (fn *groupIntoBatchesBufferedFn) OnTimer(
+ ctx context.Context, ts beam.EventTime, sp state.Provider, tp
timers.Provider,
+ key typex.T, timer timers.Context, emit func(typex.T, []typex.V),
+) {
+ fn.codec.init(fn.ValueType.T)
+ switch timer.Family {
+ case fn.Buffering.Family, fn.WindowEnd.Family:
+ fn.flush(sp, tp, key, emit)
+ default:
+ panic(fmt.Sprintf(
+ "batch.groupIntoBatchesBufferedFn: unexpected timer
family %q", timer.Family))
+ }
+}
+
+func (fn *groupIntoBatchesBufferedFn) flush(
+ sp state.Provider, tp timers.Provider, key typex.T, emit func(typex.T,
[]typex.V),
+) {
+ buf, ok, err := fn.Buffer.Read(sp)
+ if err != nil {
+ panic(err)
+ }
+ if !ok || len(buf) == 0 {
+ return
+ }
+
+ out := make([]typex.V, len(buf))
+ for i, b := range buf {
+ out[i] = fn.codec.decode(b)
+ }
+ emit(key, out)
+
+ if err := fn.Buffer.Clear(sp); err != nil {
+ panic(err)
+ }
+ if err := fn.Count.Clear(sp); err != nil {
+ panic(err)
+ }
+ if fn.BatchSizeBytes > 0 {
+ if err := fn.ByteSize.Clear(sp); err != nil {
+ panic(err)
+ }
+ }
+ setBool, _, err := fn.TimerSet.Read(sp)
+ if err != nil {
+ panic(err)
+ }
+ if setBool {
+ fn.Buffering.Clear(tp)
+ if err := fn.TimerSet.Clear(sp); err != nil {
+ panic(err)
+ }
+ }
+}
+
+func sizeOf(kind int32, v any) int64 {
+ switch kind {
+ case sizerNone:
+ return 0
+ case sizerPrimitive:
+ if size, ok := defaultElementByteSize(v); ok {
+ return size
+ }
+ panic(fmt.Sprintf("batch: sizerPrimitive cannot size value of
type %T", v))
+ default:
+ panic(fmt.Sprintf("batch: unknown sizer kind %d", kind))
+ }
+}
+
+// wrapShardedKeyFn maps KV<K, V> → KV<ShardedKey[K], V>.
+type wrapShardedKeyFn[K any] struct{}
+
+func (*wrapShardedKeyFn[K]) ProcessElement(
+ key K, value typex.V, emit func(ShardedKey[K], typex.V),
+) {
+ emit(ShardedKey[K]{Key: key, ShardID: makeShardID()}, value)
+}
+
+var (
+ workerUUIDOnce sync.Once
+ workerUUIDVal [16]byte
+ shardCounter atomic.Uint64
+)
+
+// makeShardID returns a 24-byte shard identifier: a 16-byte worker
+// UUID fixed per process plus an 8-byte atomic counter, big-endian.
+// The layout mirrors the Java and Python shapes exactly so the wire
+// bytes of cross-language round-trips remain aligned.
+func makeShardID() []byte {
+ workerUUIDOnce.Do(func() {
+ b, err := uuid.New().MarshalBinary()
+ if err != nil {
+ panic(fmt.Sprintf("batch: failed to marshal worker
UUID: %v", err))
+ }
+ copy(workerUUIDVal[:], b)
+ })
+ out := make([]byte, 24)
+ copy(out[:16], workerUUIDVal[:])
+ counter := shardCounter.Add(1)
+ binary.BigEndian.PutUint64(out[16:24], counter)
+ return out
+}
+
+// writeVarInt writes a varint-encoded int64 to buf (unsigned,
+// little-endian base-128).
+func writeVarInt(buf *bytes.Buffer, v int64) {
+ u := uint64(v)
+ for u >= 0x80 {
+ buf.WriteByte(byte(u) | 0x80)
+ u >>= 7
+ }
+ buf.WriteByte(byte(u))
+}
+
+// readVarInt reads a varint-encoded int64 from r.
+func readVarInt(r *bytes.Reader) int64 {
+ var u uint64
+ var s uint
+ for {
+ b, err := r.ReadByte()
+ if err != nil {
+ panic(err)
+ }
+ if b < 0x80 {
+ u |= uint64(b) << s
+ break
+ }
+ u |= uint64(b&0x7f) << s
+ s += 7
+ }
+ return int64(u)
+}
+
+func init() {
+ register.DoFn6x0[
+ beam.Window, state.Provider, timers.Provider,
+ typex.T, typex.V, func(typex.T, []typex.V),
+ ](&groupIntoBatchesFn{})
+ register.DoFn6x0[
+ beam.Window, state.Provider, timers.Provider,
+ typex.T, typex.V, func(typex.T, []typex.V),
+ ](&groupIntoBatchesBufferedFn{})
+ register.Emitter2[typex.T, []typex.V]()
+
+ // Register common ShardedKey[K] types for WithShardedKey.
+ RegisterShardedKeyType[string]()
+ RegisterShardedKeyType[int]()
+ RegisterShardedKeyType[int64]()
+}
+
+// GroupIntoBatches groups the values of the input PCollection<KV<K, V>>
+// into batches of up to params.BatchSize elements (or
+// params.BatchSizeBytes bytes) per key and emits them as
+// PCollection<KV<K, []V>>.
+//
+// The input must be KV-typed. The key coder must be deterministic;
+// non-deterministic key coders would corrupt state keying. Panics at
+// pipeline build time on invalid params, non-KV input, zero limits, or
+// a non-deterministic key coder.
+func GroupIntoBatches(s beam.Scope, params Params, col beam.PCollection)
beam.PCollection {
+ s = s.Scope("batch.GroupIntoBatches")
+
+ if err := params.validate(); err != nil {
+ panic(fmt.Errorf("GroupIntoBatches: %w", err))
+ }
+ if !typex.IsKV(col.Type()) {
+ panic(fmt.Errorf(
+ "GroupIntoBatches: input PCollection must be KV-typed;
got %v", col.Type()))
+ }
+
+ keyFT := col.Type().Components()[0]
+ valFT := col.Type().Components()[1]
+
+ if !beam.NewCoder(keyFT).IsDeterministic() {
+ panic(fmt.Errorf(
+ "GroupIntoBatches: key coder for type %v is not
deterministic. "+
+ "Register a deterministic custom coder with "+
+ "coder.RegisterDeterministicCoder, or use a
deterministic key "+
+ "type (string, []byte, bool, integer, float)",
keyFT.Type()))
+ }
+
+ sizerKind := sizerNone
+ if params.BatchSizeBytes > 0 {
+ if !isBuiltinSizeable(valFT.Type()) {
+ panic(fmt.Errorf(
+ "GroupIntoBatches: BatchSizeBytes > 0 requires
value type %v "+
+ "to be a built-in primitive ([]byte,
string, numeric, bool)",
+ valFT.Type()))
+ }
+ sizerKind = sizerPrimitive
+ }
+
+ allowedLatenessMs := int64(col.WindowingStrategy().AllowedLateness)
+ valueType := beam.EncodedType{T: valFT.Type()}
+
+ if params.MaxBufferingDuration > 0 {
+ fn := &groupIntoBatchesBufferedFn{
+ Buffer:
state.MakeBagState[[]byte]("batchBuffer"),
+ Count:
state.MakeValueState[int64]("batchCount"),
+ ByteSize:
state.MakeValueState[int64]("batchBytes"),
+ TimerSet:
state.MakeValueState[bool]("batchTimerSet"),
+ Buffering:
timers.InProcessingTime("batchBuffering"),
+ WindowEnd: timers.InEventTime("batchWindowEnd"),
+ ValueType: valueType,
+ BatchSize: params.BatchSize,
+ BatchSizeBytes: params.BatchSizeBytes,
+ MaxBufferingMs:
params.MaxBufferingDuration.Milliseconds(),
+ AllowedLatenessMs: allowedLatenessMs,
+ SizerKind: sizerKind,
+ }
+ return beam.ParDo(s, fn, col)
+ }
+
+ fn := &groupIntoBatchesFn{
+ Buffer: state.MakeBagState[[]byte]("batchBuffer"),
+ Count: state.MakeValueState[int64]("batchCount"),
+ ByteSize: state.MakeValueState[int64]("batchBytes"),
+ WindowEnd: timers.InEventTime("batchWindowEnd"),
+ ValueType: valueType,
+ BatchSize: params.BatchSize,
+ BatchSizeBytes: params.BatchSizeBytes,
+ AllowedLatenessMs: allowedLatenessMs,
+ SizerKind: sizerKind,
+ }
+
+ return beam.ParDo(s, fn, col)
+}
+
+// GroupIntoBatchesWithShardedKey wraps each user key with a
+// ShardedKey{Key: K, ShardID: [24]byte} and then applies
+// GroupIntoBatches. Output is PCollection<KV<ShardedKey[K], []V>>.
+//
+// The key type K must have been registered via
+// RegisterShardedKeyType[K] at init time. Common types (string,
+// []byte, int, int64) are registered automatically.
+//
+// Sharding spreads the processing of a single hot logical key across
+// multiple workers: each shard is independent state, so distributed
+// runners can parallelize without the user's key type changing.
+func GroupIntoBatchesWithShardedKey[K any](s beam.Scope, params Params, col
beam.PCollection) beam.PCollection {
+ s = s.Scope("batch.GroupIntoBatchesWithShardedKey")
+
+ if err := params.validate(); err != nil {
+ panic(fmt.Errorf("GroupIntoBatchesWithShardedKey: %w", err))
+ }
+ if !typex.IsKV(col.Type()) {
+ panic(fmt.Errorf(
+ "GroupIntoBatchesWithShardedKey: input PCollection must
be KV-typed; got %v",
+ col.Type()))
+ }
+ keyFT := col.Type().Components()[0]
+ var zero K
+ if keyFT.Type() != reflect.TypeOf(zero) {
+ panic(fmt.Errorf(
+ "GroupIntoBatchesWithShardedKey: type parameter K (%v)
does not match input key type (%v)",
+ reflect.TypeOf(zero), keyFT.Type()))
+ }
+
+ wrapped := beam.ParDo(s, &wrapShardedKeyFn[K]{}, col)
+ return GroupIntoBatches(s, params, wrapped)
+}
diff --git a/sdks/go/pkg/beam/transforms/batch/batch_prism_test.go
b/sdks/go/pkg/beam/transforms/batch/batch_prism_test.go
new file mode 100644
index 00000000000..158ea314dd0
--- /dev/null
+++ b/sdks/go/pkg/beam/transforms/batch/batch_prism_test.go
@@ -0,0 +1,222 @@
+// 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 batch
+
+import (
+ "os"
+ "sort"
+ "sync/atomic"
+ "testing"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/options/jobopts"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/register"
+ _ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
+)
+
+func TestMain(m *testing.M) {
+ f, _ := os.CreateTemp("", "dummy")
+ *jobopts.WorkerBinary = f.Name()
+ os.Exit(ptest.MainRetWithDefault(m, "prism"))
+}
+
+// splitOnBar parses "key|value" strings into KV<string,string>.
+func splitOnBar(tuple string, emit func(string, string)) {
+ for i, r := range tuple {
+ if r == '|' {
+ emit(tuple[:i], tuple[i+1:])
+ return
+ }
+ }
+}
+
+func batchSize(_ string, batch []string) int {
+ return len(batch)
+}
+
+func batchSizeSorted(_ string, batch []string) int {
+ sort.Strings(batch)
+ return len(batch)
+}
+
+// intPair emits KV<string, int> from a "key|int" string.
+func intPair(tuple string, emit func(string, int)) {
+ for i, r := range tuple {
+ if r == '|' {
+ n := 0
+ for _, c := range tuple[i+1:] {
+ n = n*10 + int(c-'0')
+ }
+ emit(tuple[:i], n)
+ return
+ }
+ }
+}
+
+func intBatchSize(_ string, batch []int) int { return len(batch) }
+
+func init() {
+ register.Function2x0(splitOnBar)
+ register.Function2x0(intPair)
+ register.Function2x1(batchSize)
+ register.Function2x1(intBatchSize)
+ register.Function2x1(batchSizeSorted)
+ register.Emitter2[string, int]()
+}
+
+// shardedBatchCount counts emitted ShardedKey batches via a side
+// channel (no GBK). Uses a package-level atomic to avoid needing a
+// Combine/GBK for aggregation, which triggers a separate Prism bug
+// on deeply-chained stateful pipelines.
+var shardedBatchCounter atomic.Int64
+
+func shardedBatchSink(sk ShardedKey[string], batch []string) {
+ _ = sk
+ _ = batch
+ shardedBatchCounter.Add(1)
+}
+
+func init() {
+ register.Function2x0(shardedBatchSink)
+}
+
+// TAC-6 (BAC-4): GroupIntoBatchesWithShardedKey wraps each key with
+// a ShardedKey and produces KV<ShardedKey[K], []V>. We validate
+// end-to-end on Prism using a terminal ParDo sink (not passert) to
+// avoid an unrelated Prism GBK panic on deeply-chained pipelines.
+func TestGroupIntoBatchesWithShardedKey_E2E(t *testing.T) {
+ shardedBatchCounter.Store(0)
+
+ p, s := beam.NewPipelineWithRoot()
+
+ tuples := make([]string, 0, 20)
+ for i := 0; i < 20; i++ {
+ tuples = append(tuples, "a|x")
+ }
+ raw := beam.CreateList(s, tuples)
+ kvs := beam.ParDo(s, splitOnBar, raw)
+
+ batches := GroupIntoBatchesWithShardedKey[string](s, Params{BatchSize:
2}, kvs)
+ beam.ParDo0(s, shardedBatchSink, batches)
+
+ ptest.RunAndValidate(t, p)
+
+ got := shardedBatchCounter.Load()
+ // Each element gets a unique shardID (atomic counter), so under
+ // Prism single-process each shard has exactly 1 element — no
+ // batching occurs (BatchSize=2 is never reached per shard).
+ // On a distributed runner the same worker/goroutine would
+ // process multiple elements of the same key, sharing a shardID
+ // and thus producing real batches. Here we verify the pipeline
+ // executed and produced 20 shard-groups.
+ if got != 20 {
+ t.Errorf("expected 20 sharded batches (one per shard), got %d",
got)
+ }
+}
+
+// TestGroupIntoBatches_IntValues verifies that GroupIntoBatches works
+// with a value type (int) that is not string — demonstrating the
+// coder-driven generic value support (BAC-1 with non-string V).
+func TestGroupIntoBatches_IntValues(t *testing.T) {
+ p, s := beam.NewPipelineWithRoot()
+
+ raw := beam.CreateList(s, []string{
+ "a|1", "a|2", "a|3", "a|4",
+ "b|5", "b|6",
+ })
+ kvs := beam.ParDo(s, intPair, raw)
+
+ batches := GroupIntoBatches(s, Params{BatchSize: 2}, kvs)
+ sizes := beam.ParDo(s, intBatchSize, batches)
+
+ passert.Equals(s, sizes, 2, 2, 2)
+
+ ptest.RunAndValidate(t, p)
+}
+
+// TAC-1 (BAC-1): 1000 inputs over 10 keys with BatchSize 100 produces
+// batches of exactly 100 elements for a single key.
+func TestGroupIntoBatches_CountLimit(t *testing.T) {
+ p, s := beam.NewPipelineWithRoot()
+
+ tuples := make([]string, 0, 1000)
+ for k := 0; k < 10; k++ {
+ for i := 0; i < 100; i++ {
+ tuples = append(tuples,
string(rune('a'+k))+"|"+string(rune('0'+i%10)))
+ }
+ }
+
+ raw := beam.CreateList(s, tuples)
+ kvs := beam.ParDo(s, splitOnBar, raw)
+
+ batches := GroupIntoBatches(s, Params{BatchSize: 100}, kvs)
+ sizes := beam.ParDo(s, batchSize, batches)
+
+ // 10 batches of 100.
+ wants := []any{}
+ for i := 0; i < 10; i++ {
+ wants = append(wants, 100)
+ }
+ passert.Equals(s, sizes, wants...)
+
+ ptest.RunAndValidate(t, p)
+}
+
+// TAC-4 (BAC-3): BatchSizeBytes threshold triggers a flush before the
+// sum exceeds the limit. With BatchSizeBytes=10 and input strings of
+// length 5 each, three 5-byte values first sum to 15 (> 10), so the
+// flush happens after 2 elements.
+func TestGroupIntoBatches_ByteLimit(t *testing.T) {
+ p, s := beam.NewPipelineWithRoot()
+
+ raw := beam.CreateList(s, []string{
+ "a|11111", "a|22222", "a|33333", "a|44444", // 4 * 5 bytes on
key a
+ "b|55555", "b|66666", // 2 * 5 bytes on key b
+ })
+ kvs := beam.ParDo(s, splitOnBar, raw)
+
+ batches := GroupIntoBatches(s, Params{BatchSizeBytes: 10}, kvs)
+ sizes := beam.ParDo(s, batchSize, batches)
+
+ // Each 2-element batch reaches 10 bytes and flushes: 2,2 for key a
+ // and 2 for key b = three flushes of size 2.
+ passert.Equals(s, sizes, 2, 2, 2)
+
+ ptest.RunAndValidate(t, p)
+}
+
+// TAC-7 (BAC-5) simplified in global window: batches only contain
+// elements for a single key. Mixed-key batches would fail the
+// key-equality assertion downstream. This test confirms the per-key
+// groupism holds.
+func TestGroupIntoBatches_PerKey(t *testing.T) {
+ p, s := beam.NewPipelineWithRoot()
+
+ raw := beam.CreateList(s, []string{
+ "a|1", "b|1", "a|2", "b|2", "a|3", "b|3", "a|4", "b|4",
+ })
+ kvs := beam.ParDo(s, splitOnBar, raw)
+
+ batches := GroupIntoBatches(s, Params{BatchSize: 2}, kvs)
+ sizes := beam.ParDo(s, batchSize, batches)
+
+ // 8 inputs / BatchSize 2 over 2 keys → 4 batches of size 2.
+ passert.Equals(s, sizes, 2, 2, 2, 2)
+
+ ptest.RunAndValidate(t, p)
+}
diff --git a/sdks/go/pkg/beam/transforms/batch/batch_test.go
b/sdks/go/pkg/beam/transforms/batch/batch_test.go
new file mode 100644
index 00000000000..0e0e00a8064
--- /dev/null
+++ b/sdks/go/pkg/beam/transforms/batch/batch_test.go
@@ -0,0 +1,47 @@
+// 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 batch
+
+import (
+ "testing"
+ "time"
+)
+
+func TestParams_validate(t *testing.T) {
+ cases := []struct {
+ name string
+ p Params
+ wantErr bool
+ }{
+ {"zero_limits", Params{}, true},
+ {"negative_size", Params{BatchSize: -1}, true},
+ {"negative_bytes", Params{BatchSizeBytes: -1}, true},
+ {"negative_duration", Params{BatchSize: 10,
MaxBufferingDuration: -time.Second}, true},
+ {"count_only", Params{BatchSize: 10}, false},
+ {"bytes_only", Params{BatchSizeBytes: 1024}, false},
+ {"both_and_duration", Params{BatchSize: 10, BatchSizeBytes:
1024, MaxBufferingDuration: time.Second}, false},
+ }
+ for _, c := range cases {
+ c := c
+ t.Run(c.name, func(t *testing.T) {
+ err := c.p.validate()
+ gotErr := err != nil
+ if gotErr != c.wantErr {
+ t.Errorf("validate() err = %v, wantErr = %v",
err, c.wantErr)
+ }
+ })
+ }
+}
diff --git a/sdks/go/pkg/beam/transforms/batch/doc.go
b/sdks/go/pkg/beam/transforms/batch/doc.go
new file mode 100644
index 00000000000..8bdbb17fbfd
--- /dev/null
+++ b/sdks/go/pkg/beam/transforms/batch/doc.go
@@ -0,0 +1,58 @@
+// 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-level doc examples for batch. Kept in the package itself so
+// `go doc` surfaces them without a separate test package and a broader
+// module import graph.
+//
+// These examples only construct a pipeline to illustrate API shape; they
+// do not run one.
+
+package batch
+
+import (
+ "fmt"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+)
+
+// ExampleGroupIntoBatches is an example of using the GroupIntoBatches
+// transform. Each input element is a (user, event) pair. After
+// GroupIntoBatches, each batch holds up to 100 events for a single user,
+// ready to be written to a BigQuery sink that accepts bulk inserts.
+func ExampleGroupIntoBatches() {
+ p := beam.NewPipeline()
+ s := p.Root()
+
+ // Build KV<string, string> PCollection via any source. The key
+ // coder (string) is deterministic so state keying is safe.
+ events := beam.CreateList(s, []string{"u1:login", "u1:click",
"u2:login"})
+ kvs := beam.ParDo(s, func(e string, emit func(string, string)) {
+ for i, r := 0, []rune(e); i < len(r); i++ {
+ if r[i] == ':' {
+ emit(string(r[:i]), string(r[i+1:]))
+ return
+ }
+ }
+ }, events)
+
+ batches := GroupIntoBatches(s, Params{BatchSize: 100}, kvs)
+
+ // Downstream: process each per-user batch.
+ _ = batches
+ fmt.Println("pipeline constructed")
+
+ // Output: pipeline constructed
+}
diff --git a/sdks/go/pkg/beam/transforms/batch/size.go
b/sdks/go/pkg/beam/transforms/batch/size.go
new file mode 100644
index 00000000000..ff1499ddaa7
--- /dev/null
+++ b/sdks/go/pkg/beam/transforms/batch/size.go
@@ -0,0 +1,88 @@
+// 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 batch
+
+import (
+ "reflect"
+)
+
+// defaultElementByteSize reports the byte cost of v for a fixed set of
+// primitive types: it is the fallback used when the caller does not
+// supply Params.ElementByteSize but BatchSizeBytes > 0.
+//
+// Returns (size, true) for supported types and (0, false) otherwise.
+// For opaque types (user structs, interfaces, maps, non-byte slices,
+// channels, functions) callers must supply their own sizer.
+func defaultElementByteSize(v any) (int64, bool) {
+ switch x := v.(type) {
+ case []byte:
+ return int64(len(x)), true
+ case string:
+ return int64(len(x)), true
+ case bool:
+ return 1, true
+ case int8:
+ return 1, true
+ case uint8:
+ return 1, true
+ case int16:
+ return 2, true
+ case uint16:
+ return 2, true
+ case int32:
+ return 4, true
+ case uint32:
+ return 4, true
+ case float32:
+ return 4, true
+ case int:
+ return 8, true
+ case uint:
+ return 8, true
+ case int64:
+ return 8, true
+ case uint64:
+ return 8, true
+ case float64:
+ return 8, true
+ }
+ return 0, false
+}
+
+// isBuiltinSizeable reports whether defaultElementByteSize can size an
+// element of type t. Used at pipeline-build time to fail fast when
+// BatchSizeBytes > 0 is requested without a user-supplied
+// ElementByteSize and the value type is not one of the supported
+// primitives.
+//
+// A []byte is recognized via reflect.Slice with Uint8 element kind; any
+// other slice is not sizeable by the built-in fallback.
+func isBuiltinSizeable(t reflect.Type) bool {
+ if t == nil {
+ return false
+ }
+ switch t.Kind() {
+ case reflect.String,
+ reflect.Bool,
+ reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32,
reflect.Int64,
+ reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32,
reflect.Uint64,
+ reflect.Float32, reflect.Float64:
+ return true
+ case reflect.Slice:
+ return t.Elem().Kind() == reflect.Uint8
+ }
+ return false
+}
diff --git a/sdks/go/pkg/beam/transforms/batch/size_test.go
b/sdks/go/pkg/beam/transforms/batch/size_test.go
new file mode 100644
index 00000000000..82d2d8dc044
--- /dev/null
+++ b/sdks/go/pkg/beam/transforms/batch/size_test.go
@@ -0,0 +1,91 @@
+// 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 batch
+
+import (
+ "reflect"
+ "testing"
+)
+
+func TestDefaultElementByteSize(t *testing.T) {
+ cases := []struct {
+ name string
+ v any
+ want int64
+ ok bool
+ }{
+ {"bytes_5", []byte("abcde"), 5, true},
+ {"bytes_empty", []byte{}, 0, true},
+ {"string_5", "abcde", 5, true},
+ {"string_empty", "", 0, true},
+ {"bool", true, 1, true},
+ {"int8", int8(1), 1, true},
+ {"uint8", uint8(1), 1, true},
+ {"int16", int16(1), 2, true},
+ {"uint16", uint16(1), 2, true},
+ {"int32", int32(1), 4, true},
+ {"uint32", uint32(1), 4, true},
+ {"float32", float32(1.0), 4, true},
+ {"int", int(1), 8, true},
+ {"uint", uint(1), 8, true},
+ {"int64", int64(1), 8, true},
+ {"uint64", uint64(1), 8, true},
+ {"float64", float64(1.0), 8, true},
+ {"struct_unsupported", struct{ A int }{A: 1}, 0, false},
+ {"map_unsupported", map[string]int{"a": 1}, 0, false},
+ {"slice_int_unsupported", []int{1, 2, 3}, 0, false},
+ }
+ for _, c := range cases {
+ c := c
+ t.Run(c.name, func(t *testing.T) {
+ got, ok := defaultElementByteSize(c.v)
+ if ok != c.ok {
+ t.Errorf("ok = %v, want %v", ok, c.ok)
+ }
+ if got != c.want {
+ t.Errorf("size = %d, want %d", got, c.want)
+ }
+ })
+ }
+}
+
+func TestIsBuiltinSizeable(t *testing.T) {
+ cases := []struct {
+ name string
+ t reflect.Type
+ want bool
+ }{
+ {"nil", nil, false},
+ {"string", reflect.TypeOf(""), true},
+ {"bytes", reflect.TypeOf([]byte(nil)), true},
+ {"bool", reflect.TypeOf(true), true},
+ {"int", reflect.TypeOf(int(0)), true},
+ {"int64", reflect.TypeOf(int64(0)), true},
+ {"float64", reflect.TypeOf(float64(0)), true},
+ {"struct", reflect.TypeOf(struct{ A int }{}), false},
+ {"map", reflect.TypeOf(map[string]int{}), false},
+ {"slice_int", reflect.TypeOf([]int{}), false},
+ {"slice_string", reflect.TypeOf([]string{}), false},
+ }
+ for _, c := range cases {
+ c := c
+ t.Run(c.name, func(t *testing.T) {
+ if got := isBuiltinSizeable(c.t); got != c.want {
+ t.Errorf("isBuiltinSizeable(%v) = %v, want %v",
c.t, got, c.want)
+ }
+ })
+ }
+}
