Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Go] Ensure iterated and emitted types are registered. #23890

Merged
merged 3 commits into from
Nov 2, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGES.md
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@
* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
* Fixed JmsIO acknowledgment issue (https://github.com/apache/beam/issues/20814)
* Fixed Beam SQL CalciteUtils (Java) and Cross-language JdbcIO (Python) did not support JDBC CHAR/VARCHAR, BINARY/VARBINARY logical types ([#23747](https://github.com/apache/beam/issues/23747), [#23526](https://github.com/apache/beam/issues/23526)).
* Ensure iterated and emitted types are used with the generic register package are registered with the type and schema registries.(Go) ([#23889](https://github.com/apache/beam/pull/23889))

# [2.43.0] - Unreleased

Expand Down
6 changes: 5 additions & 1 deletion sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ func getUUID(ut reflect.Type) string {
// Registered returns whether the given type has been registered with
// the schema package.
func (r *Registry) Registered(ut reflect.Type) bool {
r.reconcileRegistrations()
_, ok := r.syntheticToUser[ut]
return ok
}
Expand All @@ -118,7 +119,10 @@ func (r *Registry) reconcileRegistrations() (deferedErr error) {
check := func(ut reflect.Type) bool {
return coder.LookupCustomCoder(ut) != nil
}
if check(ut) || check(reflect.PtrTo(ut)) {
// We could have either a pointer or non pointer here,
// so we strip pointerness and then check both.
vT := reflectx.SkipPtr(ut)
if check(vT) && check(reflect.PtrTo(vT)) {
continue
}
if err := r.registerType(ut, map[reflect.Type]struct{}{}); err != nil {
Expand Down
2 changes: 2 additions & 0 deletions sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (
pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
"github.com/golang/protobuf/proto"
"github.com/google/go-cmp/cmp"
"google.golang.org/protobuf/encoding/prototext"
"google.golang.org/protobuf/testing/protocmp"
)

Expand Down Expand Up @@ -792,6 +793,7 @@ func TestSchemaConversion(t *testing.T) {
// real embedded type.
if !hasEmbeddedField(test.rt) && !test.rt.AssignableTo(got) {
t.Errorf("%v not assignable to %v", test.rt, got)
t.Errorf("%v for schema %v", test.rt, prototext.Format(test.st))
if d := cmp.Diff(reflect.New(test.rt).Elem().Interface(), reflect.New(got).Elem().Interface()); d != "" {
t.Errorf("diff (-want, +got): %v", d)
}
Expand Down
23 changes: 16 additions & 7 deletions sdks/go/pkg/beam/register/emitter.go
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,9 @@ func Emitter1[T1 any]() {
registerFunc := func(n exec.ElementProcessor) exec.ReusableEmitter {
return &emit1[T1]{n: n}
}
exec.RegisterEmitter(reflect.TypeOf(e).Elem(), registerFunc)
eT := reflect.TypeOf(e).Elem()
registerType(eT.In(0))
exec.RegisterEmitter(eT, registerFunc)
}

// Emitter2 registers parameters from your DoFn with a
Expand All @@ -147,18 +149,25 @@ func Emitter2[T1, T2 any]() {
return &emit1WithTimestamp[T2]{n: n}
}
}
exec.RegisterEmitter(reflect.TypeOf(e).Elem(), registerFunc)
eT := reflect.TypeOf(e).Elem()
registerType(eT.In(0))
registerType(eT.In(1))
exec.RegisterEmitter(eT, registerFunc)
}

// Emitter3 registers parameters from your DoFn with a
// signature func(T1, T2, T3) and optimizes their execution.
// signature func(beam.EventTime, T2, T3) and optimizes their execution.
// This must be done by passing in type parameters of all inputs as constraints,
// aka: register.Emitter3[beam.EventTime, T1, T2](), where T1 is the type of
// your key and T2 is the type of your value.
func Emitter3[T1 typex.EventTime, T2, T3 any]() {
e := (*func(T1, T2, T3))(nil)
func Emitter3[ET typex.EventTime, T1, T2 any]() {
e := (*func(ET, T1, T2))(nil)
registerFunc := func(n exec.ElementProcessor) exec.ReusableEmitter {
return &emit2WithTimestamp[T2, T3]{n: n}
return &emit2WithTimestamp[T1, T2]{n: n}
}
exec.RegisterEmitter(reflect.TypeOf(e).Elem(), registerFunc)
eT := reflect.TypeOf(e).Elem()
// No need to register event time.
registerType(eT.In(1))
registerType(eT.In(2))
exec.RegisterEmitter(eT, registerFunc)
}
166 changes: 166 additions & 0 deletions sdks/go/pkg/beam/register/emitter_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,166 @@
// 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 register

import (
"context"
"reflect"
"testing"

"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
)

type myTestTypeEmitter1 struct {
Int int
}

func TestEmitter1(t *testing.T) {
Emitter1[int]()
e1T := reflect.TypeOf((*func(int))(nil)).Elem()
if !exec.IsEmitterRegistered(e1T) {
t.Fatalf("exec.IsEmitterRegistered(%v) = false, want true", e1T)
}

Emitter1[myTestTypeEmitter1]()
rt := reflect.TypeOf((*myTestTypeEmitter1)(nil)).Elem()
checkRegisterations(t, rt)
}

type myTestTypeEmitter2A struct {
Int int
}

type myTestTypeEmitter2B struct {
String string
}

func TestEmitter2(t *testing.T) {
Emitter2[int, string]()
e2isT := reflect.TypeOf((*func(int, string))(nil)).Elem()
if !exec.IsEmitterRegistered(e2isT) {
t.Fatalf("exec.IsEmitterRegistered(%v) = false, want true", e2isT)
}

Emitter2[*myTestTypeEmitter2A, myTestTypeEmitter2B]()
e2ABT := reflect.TypeOf((*func(*myTestTypeEmitter2A, myTestTypeEmitter2B))(nil)).Elem()
if !exec.IsEmitterRegistered(e2ABT) {
t.Fatalf("exec.IsEmitterRegistered(%v) = false, want true", e2ABT)
}

tA := reflect.TypeOf((*myTestTypeEmitter2A)(nil)).Elem()
checkRegisterations(t, tA)
tB := reflect.TypeOf((*myTestTypeEmitter2B)(nil)).Elem()
checkRegisterations(t, tB)
}

func TestEmitter2_WithTimestamp(t *testing.T) {
Emitter2[typex.EventTime, string]()
e2tssT := reflect.TypeOf((*func(typex.EventTime, string))(nil)).Elem()
if !exec.IsEmitterRegistered(e2tssT) {
t.Fatalf("exec.IsEmitterRegistered(%v) = false, want true", e2tssT)
}
}

type myTestTypeEmitter3A struct {
Int int
}

type myTestTypeEmitter3B struct {
String string
}

func TestEmitter3(t *testing.T) {
Emitter3[typex.EventTime, int, string]()
if !exec.IsEmitterRegistered(reflect.TypeOf((*func(typex.EventTime, int, string))(nil)).Elem()) {
t.Fatalf("exec.IsEmitterRegistered(reflect.TypeOf((*func(typex.EventTime, int, string))(nil)).Elem()) = false, want true")
}

Emitter3[typex.EventTime, myTestTypeEmitter3A, *myTestTypeEmitter3B]()
e3tsABT := reflect.TypeOf((*func(typex.EventTime, myTestTypeEmitter3A, *myTestTypeEmitter3B))(nil)).Elem()
if !exec.IsEmitterRegistered(e3tsABT) {
t.Fatalf("exec.IsEmitterRegistered(%v) = false, want true", e3tsABT)
}
tA := reflect.TypeOf((*myTestTypeEmitter3A)(nil)).Elem()
checkRegisterations(t, tA)
tB := reflect.TypeOf((*myTestTypeEmitter3B)(nil)).Elem()
checkRegisterations(t, tB)
}

func TestEmit1(t *testing.T) {
e := &emit1[int]{n: &elementProcessor{}}
e.Init(context.Background(), []typex.Window{}, mtime.ZeroTimestamp)
fn := e.Value().(func(int))
fn(3)
if got, want := e.n.(*elementProcessor).inFV.Elm, 3; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Elm=%v, want %v", got, want)
}
if got := e.n.(*elementProcessor).inFV.Elm2; got != nil {
t.Errorf("e.Value.(func(int))(3).n.inFV.Elm2=%v, want nil", got)
}
if got, want := e.n.(*elementProcessor).inFV.Timestamp, mtime.ZeroTimestamp; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Timestamp=%v, want %v", got, want)
}
}

func TestEmit2(t *testing.T) {
e := &emit2[int, string]{n: &elementProcessor{}}
e.Init(context.Background(), []typex.Window{}, mtime.ZeroTimestamp)
fn := e.Value().(func(int, string))
fn(3, "hello")
if got, want := e.n.(*elementProcessor).inFV.Elm, 3; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Elm=%v, want %v", got, want)
}
if got, want := e.n.(*elementProcessor).inFV.Elm2, "hello"; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Elm2=%v, want %v", got, want)
}
if got, want := e.n.(*elementProcessor).inFV.Timestamp, mtime.ZeroTimestamp; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Timestamp=%v, want %v", got, want)
}
}

func TestEmit1WithTimestamp(t *testing.T) {
e := &emit1WithTimestamp[int]{n: &elementProcessor{}}
e.Init(context.Background(), []typex.Window{}, mtime.ZeroTimestamp)
fn := e.Value().(func(typex.EventTime, int))
fn(mtime.MaxTimestamp, 3)
if got, want := e.n.(*elementProcessor).inFV.Elm, 3; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Elm=%v, want %v", got, want)
}
if got := e.n.(*elementProcessor).inFV.Elm2; got != nil {
t.Errorf("e.Value.(func(int))(3).n.inFV.Elm2=%v, want nil", got)
}
if got, want := e.n.(*elementProcessor).inFV.Timestamp, mtime.MaxTimestamp; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Timestamp=%v, want %v", got, want)
}
}

func TestEmit2WithTimestamp(t *testing.T) {
e := &emit2WithTimestamp[int, string]{n: &elementProcessor{}}
e.Init(context.Background(), []typex.Window{}, mtime.ZeroTimestamp)
fn := e.Value().(func(typex.EventTime, int, string))
fn(mtime.MaxTimestamp, 3, "hello")
if got, want := e.n.(*elementProcessor).inFV.Elm, 3; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Elm=%v, want %v", got, want)
}
if got, want := e.n.(*elementProcessor).inFV.Elm2, "hello"; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Elm2=%v, want %v", got, want)
}
if got, want := e.n.(*elementProcessor).inFV.Timestamp, mtime.MaxTimestamp; got != want {
t.Errorf("e.Value.(func(int))(3).n.inFV.Timestamp=%v, want %v", got, want)
}
}
22 changes: 20 additions & 2 deletions sdks/go/pkg/beam/register/iter.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,10 @@ import (
"io"
"reflect"

"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/schema"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
)

type iter1[T any] struct {
Expand Down Expand Up @@ -104,6 +107,16 @@ func (v *iter2[T1, T2]) invoke(key *T1, value *T2) bool {
return true
}

func registerType(t reflect.Type) {
// strip the pointer if present.
t = reflectx.SkipPtr(t)
if _, ok := runtime.TypeKey(t); !ok {
return
}
runtime.RegisterType(t)
schema.RegisterType(t)
}

// Iter1 registers parameters from your DoFn with a
// signature func(*T) bool and optimizes their execution.
// This must be done by passing in type parameters of all inputs as constraints,
Expand All @@ -113,7 +126,9 @@ func Iter1[T any]() {
registerFunc := func(s exec.ReStream) exec.ReusableInput {
return &iter1[T]{s: s}
}
exec.RegisterInput(reflect.TypeOf(i).Elem(), registerFunc)
itT := reflect.TypeOf(i).Elem()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @lostluck
I just started a test run with your changes applied with the bigtable io connector (#23411).
Unfortunately I still ran into the following error: reflect.Set: value of type struct { RowKey string; Ops []struct { Family string; Column string; Ts int64; Value []uint8 }; GroupKey string } is not assignable to type bigtableio.Mutation.
I played around with this implementation but wasn't able to make it work. My knowledge on the reflect package is somehow limited though.

registerType(itT.In(0).Elem())
exec.RegisterInput(itT, registerFunc)
}

// Iter1 registers parameters from your DoFn with a
Expand All @@ -125,5 +140,8 @@ func Iter2[T1, T2 any]() {
registerFunc := func(s exec.ReStream) exec.ReusableInput {
return &iter2[T1, T2]{s: s}
}
exec.RegisterInput(reflect.TypeOf(i).Elem(), registerFunc)
itT := reflect.TypeOf(i).Elem()
registerType(itT.In(0).Elem())
registerType(itT.In(1).Elem())
exec.RegisterInput(itT, registerFunc)
}
Loading