From ea83047621db3f941995e5b0515ff413cdb68629 Mon Sep 17 00:00:00 2001 From: Yashvardhan Thanvi Date: Mon, 19 Oct 2026 01:17:58 +0530 Subject: [PATCH] [Go SDK] Fix generic registration of anonymous struct types Generic registration registers every struct type that a function or method takes or returns, with both the runtime type registry and the schema registry. runtime.RegisterType panics for unnamed types, so generic registration could not be used for a function with an anonymous struct parameter, and it stopped before the caller and struct wrappers were set up. registerMethodTypes now adds a struct type to the runtime type registry only when it has a type key, and still adds every struct type to the schema registry, which maps a row schema back to the user's Go type on the worker. graphx already encodes types without a type key field by field. This also enables the castFn registration in the BigQuery integration test that was waiting on this issue. Fixes #21789 --- CHANGES.md | 1 + sdks/go/pkg/beam/register/register.go | 16 ++- sdks/go/pkg/beam/register/register.tmpl | 16 ++- sdks/go/pkg/beam/register/register_test.go | 135 ++++++++++++++++++ .../io/xlang/bigquery/bigquery_test.go | 3 +- 5 files changed, 161 insertions(+), 10 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index c9350da11d69..d0f3c41cd7ed 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -138,6 +138,7 @@ * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)). * (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)). * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)). +* (Go) Fixed generic registration (for example `register.Function1x1` or `register.DoFn1x1`) panicking on functions and methods that take or return an anonymous struct. Unnamed struct types are now only registered with the schema registry, since they have no type key for the runtime type registry ([#21789](https://github.com/apache/beam/issues/21789)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/go/pkg/beam/register/register.go b/sdks/go/pkg/beam/register/register.go index fd0570949f8e..5488a65f93aa 100644 --- a/sdks/go/pkg/beam/register/register.go +++ b/sdks/go/pkg/beam/register/register.go @@ -8480,19 +8480,27 @@ func registerDoFnTypes(doFn any) { registerMethodTypes(reflect.ValueOf(doFn).MethodByName("ProcessElement").Type()) } +// registerStructType registers t with the schema registry, and with the +// runtime type registry when t has a type key. Unnamed types, such as anonymous +// structs, have no type key, and graphx encodes them field by field instead. +func registerStructType(t reflect.Type) { + if _, ok := runtime.TypeKey(t); ok { + runtime.RegisterType(t) + } + schema.RegisterType(t) +} + func registerMethodTypes(fn reflect.Type) { for i := 0; i < fn.NumIn(); i++ { in := reflectx.SkipPtr(fn.In(i)) if in.Kind() == reflect.Struct { - runtime.RegisterType(in) - schema.RegisterType(in) + registerStructType(in) } } for i := 0; i < fn.NumOut(); i++ { out := reflectx.SkipPtr(fn.Out(i)) if out.Kind() == reflect.Struct { - runtime.RegisterType(out) - schema.RegisterType(out) + registerStructType(out) } } } diff --git a/sdks/go/pkg/beam/register/register.tmpl b/sdks/go/pkg/beam/register/register.tmpl index 0914780d6d43..441203aec51f 100644 --- a/sdks/go/pkg/beam/register/register.tmpl +++ b/sdks/go/pkg/beam/register/register.tmpl @@ -567,19 +567,27 @@ func registerDoFnTypes(doFn any) { registerMethodTypes(reflect.ValueOf(doFn).MethodByName("ProcessElement").Type()) } +// registerStructType registers t with the schema registry, and with the +// runtime type registry when t has a type key. Unnamed types, such as anonymous +// structs, have no type key, and graphx encodes them field by field instead. +func registerStructType(t reflect.Type) { + if _, ok := runtime.TypeKey(t); ok { + runtime.RegisterType(t) + } + schema.RegisterType(t) +} + func registerMethodTypes(fn reflect.Type) { for i := 0; i < fn.NumIn(); i++ { in := reflectx.SkipPtr(fn.In(i)) if in.Kind() == reflect.Struct { - runtime.RegisterType(in) - schema.RegisterType(in) + registerStructType(in) } } for i := 0; i < fn.NumOut(); i++ { out := reflectx.SkipPtr(fn.Out(i)) if out.Kind() == reflect.Struct { - runtime.RegisterType(out) - schema.RegisterType(out) + registerStructType(out) } } } \ No newline at end of file diff --git a/sdks/go/pkg/beam/register/register_test.go b/sdks/go/pkg/beam/register/register_test.go index 1d00a4fad03d..f3f14a8a16d3 100644 --- a/sdks/go/pkg/beam/register/register_test.go +++ b/sdks/go/pkg/beam/register/register_test.go @@ -351,6 +351,141 @@ func TestFunction(t *testing.T) { } } +// anonymousStruct is a type alias, so the struct type stays unnamed and has no +// type key to register it under in the runtime type registry. +type anonymousStruct = struct { + Counter *int64 `beam:"counter"` + RandData *struct { + Flip *bool `beam:"flip"` + Word *string `beam:"word"` + } `beam:"rand_data"` + Meta AnonymousStructMeta `beam:"meta"` +} + +// AnonymousStructMeta is only used as a field of anonymousStruct. +type AnonymousStructMeta struct { + Source string +} + +type AnonymousStructReturn struct { + Count int64 +} + +func anonymousStructFunction(a *anonymousStruct) AnonymousStructReturn { + return AnonymousStructReturn{Count: *a.Counter} +} + +// anonymousStructResult is unnamed too, with a different shape from +// anonymousStruct, so that its schema registration comes from its own test. +type anonymousStructResult = struct { + Word string `json:"word"` + Count int64 `json:"count"` +} + +func anonymousStructResultFunction(word string) anonymousStructResult { + return anonymousStructResult{Word: word, Count: int64(len(word))} +} + +type AnonymousStructDoFn1x1 struct{} + +func (fn *AnonymousStructDoFn1x1) ProcessElement(a *anonymousStruct) AnonymousStructReturn { + return AnonymousStructReturn{Count: *a.Counter} +} + +func registerWithoutPanic(t *testing.T, name string, register func()) { + t.Helper() + defer func() { + if r := recover(); r != nil { + t.Fatalf("%v panicked: %v", name, r) + } + }() + register() +} + +func checkTypeRegistered(t *testing.T, rt reflect.Type) { + t.Helper() + key, ok := runtime.TypeKey(rt) + if !ok { + t.Fatalf("runtime.TypeKey(%v) failed, want a registrable type", rt) + } + if _, ok := runtime.LookupType(key); !ok { + t.Errorf("runtime.LookupType(%q) = false, want true", key) + } + checkSchemaRegistered(t, rt) +} + +func checkSchemaRegistered(t *testing.T, rt reflect.Type) { + t.Helper() + // Need to call FromType so that the registry will reconcile its registrations + if _, err := schema.FromType(rt); err != nil { + t.Errorf("schema.FromType(%v) failed: %v", rt, err) + } + if !schema.Registered(rt) { + t.Errorf("schema.Registered(%v) = false, want true", rt) + } +} + +func TestFunction_AnonymousStructParameter(t *testing.T) { + registerWithoutPanic(t, "Function1x1", func() { + Function1x1[*anonymousStruct, AnonymousStructReturn](anonymousStructFunction) + }) + + // The unnamed parameter type has no type key, but it is still registered + // with the schema registry, along with the named struct type it contains. + checkSchemaRegistered(t, reflect.TypeOf(anonymousStruct{})) + checkSchemaRegistered(t, reflect.TypeOf(AnonymousStructMeta{})) + checkTypeRegistered(t, reflect.TypeOf(AnonymousStructReturn{})) + + f := reflectx.MakeFunc(anonymousStructFunction) + if _, ok := f.(*caller1x1[*anonymousStruct, AnonymousStructReturn]); !ok { + t.Fatalf("reflectx.MakeFunc(anonymousStructFunction) = %T, want *caller1x1", f) + } + counter := int64(7) + if got, want := f.Call([]any{&anonymousStruct{Counter: &counter}})[0].(AnonymousStructReturn).Count, counter; got != want { + t.Errorf("Registered function call: got %v, want %v", got, want) + } +} + +func TestFunction_AnonymousStructReturn(t *testing.T) { + registerWithoutPanic(t, "Function1x1", func() { + Function1x1[string, anonymousStructResult](anonymousStructResultFunction) + }) + + checkSchemaRegistered(t, reflect.TypeOf(anonymousStructResult{})) + + f := reflectx.MakeFunc(anonymousStructResultFunction) + if _, ok := f.(*caller1x1[string, anonymousStructResult]); !ok { + t.Fatalf("reflectx.MakeFunc(anonymousStructResultFunction) = %T, want *caller1x1", f) + } + want := anonymousStructResult{Word: "beam", Count: 4} + if got := f.Call([]any{"beam"})[0].(anonymousStructResult); got != want { + t.Errorf("Registered function call: got %v, want %v", got, want) + } +} + +func TestRegister_AnonymousStructParameter(t *testing.T) { + registerWithoutPanic(t, "DoFn1x1", func() { + DoFn1x1[*anonymousStruct, AnonymousStructReturn](&AnonymousStructDoFn1x1{}) + }) + + checkTypeRegistered(t, reflect.TypeOf(AnonymousStructDoFn1x1{})) + checkTypeRegistered(t, reflect.TypeOf(AnonymousStructReturn{})) + checkSchemaRegistered(t, reflect.TypeOf(anonymousStruct{})) + + m, ok := reflectx.WrapMethods(&AnonymousStructDoFn1x1{}) + if !ok { + t.Fatalf("reflectx.WrapMethods(AnonymousStructDoFn1x1{}), no registered entry found") + } + pe, ok := m["ProcessElement"] + if !ok { + t.Fatalf("reflectx.WrapMethods(AnonymousStructDoFn1x1{}), no registered ProcessElement method found") + } + counter := int64(7) + if got, want := pe.Call([]any{&anonymousStruct{Counter: &counter}})[0].(AnonymousStructReturn).Count, counter; got != want { + t.Errorf("Wrapped ProcessElement call: got %v, want %v", got, want) + } +} + type elementProcessor struct { inFV exec.FullValue } diff --git a/sdks/go/test/integration/io/xlang/bigquery/bigquery_test.go b/sdks/go/test/integration/io/xlang/bigquery/bigquery_test.go index 884d7662da7b..a9148f81c92e 100644 --- a/sdks/go/test/integration/io/xlang/bigquery/bigquery_test.go +++ b/sdks/go/test/integration/io/xlang/bigquery/bigquery_test.go @@ -38,8 +38,7 @@ import ( func init() { register.DoFn2x0[[]byte, func(TestRow)](&CreateTestRowsFn{}) register.Emitter1[TestRow]() - // TODO(https://github.com/apache/beam/issues/21789): Uncomment once this register no longer panics. - //register.Function1x1(castFn) + register.Function1x1(castFn) } var expansionAddr string // Populate with expansion address labelled "gcpio".