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".