This is an automated email from the ASF dual-hosted git repository.
shunping 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 927c5b23029 [Go SDK] Fix generic registration of anonymous struct
types (#40649)
927c5b23029 is described below
commit 927c5b230299bf1c5b49b867f0261448ad56696c
Author: Yashvardhan Thanvi <[email protected]>
AuthorDate: Mon Oct 26 22:28:57 2026 +0530
[Go SDK] Fix generic registration of anonymous struct types (#40649)
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 +++++++++++++++++++++
.../integration/io/xlang/bigquery/bigquery_test.go | 3 +-
5 files changed, 161 insertions(+), 10 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 8fa59e7d5ea..04ac961fcd9 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -141,6 +141,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 fd0570949f8..5488a65f93a 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 0914780d6d4..441203aec51 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 1d00a4fad03..f3f14a8a16d 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 884d7662da7..a9148f81c92 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".
