diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue.go new file mode 100644 index 000000000000..a8e8ffaed091 --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue.go @@ -0,0 +1,156 @@ +// 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 schema + +import ( + "math" + "reflect" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx" + "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" + pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" +) + +// atomicFieldType returns the schema FieldType of an atomic Go type. +// The Go type must be one of the built in types that map directly to an +// AtomicType: bool, uint8, int16, int32, int64, float32, float64, string +// and []byte. +func atomicFieldType(t reflect.Type) (*pipepb.FieldType, error) { + var at pipepb.AtomicType + if t == reflectx.ByteSlice { + at = pipepb.AtomicType_BYTES + } else if enum, ok := reflectTypeToAtomicTypeMap[t.Kind()]; ok && reflectKindToTypeMap[t.Kind()] == t { + at = enum + } else { + return nil, errors.Errorf("%v is not an atomic schema type", t) + } + return &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: at, + }, + }, nil +} + +// atomicValueToProto converts a Go value of an atomic type to its schema +// FieldType and FieldValue. The BYTE schema type is signed, so a uint8 is +// converted through int8. +func atomicValueToProto(rv reflect.Value) (*pipepb.FieldType, *pipepb.FieldValue, error) { + if !rv.IsValid() { + return nil, nil, errors.New("invalid value") + } + ft, err := atomicFieldType(rv.Type()) + if err != nil { + return nil, nil, err + } + var av *pipepb.AtomicTypeValue + switch ft.GetAtomicType() { + case pipepb.AtomicType_BYTE: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Byte{Byte: int32(int8(rv.Uint()))}} + case pipepb.AtomicType_INT16: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int16{Int16: int32(rv.Int())}} + case pipepb.AtomicType_INT32: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int32{Int32: int32(rv.Int())}} + case pipepb.AtomicType_INT64: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int64{Int64: rv.Int()}} + case pipepb.AtomicType_FLOAT: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Float{Float: float32(rv.Float())}} + case pipepb.AtomicType_DOUBLE: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Double{Double: rv.Float()}} + case pipepb.AtomicType_STRING: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_String_{String_: rv.String()}} + case pipepb.AtomicType_BOOLEAN: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Boolean{Boolean: rv.Bool()}} + case pipepb.AtomicType_BYTES: + av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Bytes{Bytes: rv.Bytes()}} + default: + return nil, nil, errors.Errorf("unhandled atomic type %v", ft.GetAtomicType()) + } + return ft, &pipepb.FieldValue{FieldValue: &pipepb.FieldValue_AtomicValue{AtomicValue: av}}, nil +} + +// atomicValueFromProto converts a schema FieldValue of the atomic FieldType ft +// to a Go value of the matching built in type. +func atomicValueFromProto(ft *pipepb.FieldType, fv *pipepb.FieldValue) (reflect.Value, error) { + at, ok := ft.GetTypeInfo().(*pipepb.FieldType_AtomicType) + if !ok { + return reflect.Value{}, errors.Errorf("field type %v is not atomic", ft) + } + t, ok := atomicTypeToReflectType[at.AtomicType] + if !ok { + return reflect.Value{}, errors.Errorf("unknown atomic type: %v", at.AtomicType) + } + av := fv.GetAtomicValue() + if av == nil { + return reflect.Value{}, errors.Errorf("field value %v is not atomic", fv) + } + if got := atomicValueType(av); got != at.AtomicType { + return reflect.Value{}, errors.Errorf("atomic value %v does not match field type %v", av, at.AtomicType) + } + rv := reflect.New(t).Elem() + switch v := av.GetValue().(type) { + case *pipepb.AtomicTypeValue_Byte: + if v.Byte < math.MinInt8 || v.Byte > math.MaxInt8 { + return reflect.Value{}, errors.Errorf("BYTE value %v is out of range", v.Byte) + } + rv.SetUint(uint64(uint8(v.Byte))) + case *pipepb.AtomicTypeValue_Int16: + if v.Int16 < math.MinInt16 || v.Int16 > math.MaxInt16 { + return reflect.Value{}, errors.Errorf("INT16 value %v is out of range", v.Int16) + } + rv.SetInt(int64(v.Int16)) + case *pipepb.AtomicTypeValue_Int32: + rv.SetInt(int64(v.Int32)) + case *pipepb.AtomicTypeValue_Int64: + rv.SetInt(v.Int64) + case *pipepb.AtomicTypeValue_Float: + rv.SetFloat(float64(v.Float)) + case *pipepb.AtomicTypeValue_Double: + rv.SetFloat(v.Double) + case *pipepb.AtomicTypeValue_String_: + rv.SetString(v.String_) + case *pipepb.AtomicTypeValue_Boolean: + rv.SetBool(v.Boolean) + case *pipepb.AtomicTypeValue_Bytes: + rv.SetBytes(v.Bytes) + } + return rv, nil +} + +// atomicValueType returns the atomic type of the value set in av, or +// UNSPECIFIED when no value is set. +func atomicValueType(av *pipepb.AtomicTypeValue) pipepb.AtomicType { + switch av.GetValue().(type) { + case *pipepb.AtomicTypeValue_Byte: + return pipepb.AtomicType_BYTE + case *pipepb.AtomicTypeValue_Int16: + return pipepb.AtomicType_INT16 + case *pipepb.AtomicTypeValue_Int32: + return pipepb.AtomicType_INT32 + case *pipepb.AtomicTypeValue_Int64: + return pipepb.AtomicType_INT64 + case *pipepb.AtomicTypeValue_Float: + return pipepb.AtomicType_FLOAT + case *pipepb.AtomicTypeValue_Double: + return pipepb.AtomicType_DOUBLE + case *pipepb.AtomicTypeValue_String_: + return pipepb.AtomicType_STRING + case *pipepb.AtomicTypeValue_Boolean: + return pipepb.AtomicType_BOOLEAN + case *pipepb.AtomicTypeValue_Bytes: + return pipepb.AtomicType_BYTES + } + return pipepb.AtomicType_UNSPECIFIED +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue_test.go new file mode 100644 index 000000000000..a2c0a1c66146 --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue_test.go @@ -0,0 +1,163 @@ +// 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 schema + +import ( + "math" + "reflect" + "testing" + + pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" + "github.com/google/go-cmp/cmp" +) + +func TestAtomicValueProtoRoundTrip(t *testing.T) { + tests := []struct { + value any + want pipepb.AtomicType + }{ + {true, pipepb.AtomicType_BOOLEAN}, + {uint8(200), pipepb.AtomicType_BYTE}, + {int16(-300), pipepb.AtomicType_INT16}, + {int32(70000), pipepb.AtomicType_INT32}, + {int64(-1 << 40), pipepb.AtomicType_INT64}, + {float32(1.5), pipepb.AtomicType_FLOAT}, + {float64(2.25), pipepb.AtomicType_DOUBLE}, + {"text", pipepb.AtomicType_STRING}, + {[]byte("bytes"), pipepb.AtomicType_BYTES}, + } + for _, test := range tests { + t.Run(reflect.TypeOf(test.value).String(), func(t *testing.T) { + ft, fv, err := atomicValueToProto(reflect.ValueOf(test.value)) + if err != nil { + t.Fatalf("atomicValueToProto(%v) = %v, want nil error", test.value, err) + } + if got := ft.GetAtomicType(); got != test.want { + t.Errorf("atomicValueToProto(%v) field type = %v, want %v", test.value, got, test.want) + } + got, err := atomicValueFromProto(ft, fv) + if err != nil { + t.Fatalf("atomicValueFromProto(%v, %v) = %v, want nil error", ft, fv, err) + } + if d := cmp.Diff(test.value, got.Interface()); d != "" { + t.Errorf("round trip of %v: diff (-want, +got): %v", test.value, d) + } + }) + } +} + +func TestAtomicValueToProto_Unsupported(t *testing.T) { + type named int32 + for _, value := range []any{int(1), uint32(1), named(1), struct{ A int32 }{}, []int32{1}} { + if _, _, err := atomicValueToProto(reflect.ValueOf(value)); err == nil { + t.Errorf("atomicValueToProto(%T) = nil error, want error", value) + } + } +} + +func TestAtomicValueFromProto_Mismatch(t *testing.T) { + stringType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_STRING}, + } + int32Value := &pipepb.FieldValue{ + FieldValue: &pipepb.FieldValue_AtomicValue{ + AtomicValue: &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int32{Int32: 3}}, + }, + } + if _, err := atomicValueFromProto(stringType, int32Value); err == nil { + t.Errorf("atomicValueFromProto(STRING, int32 value) = nil error, want error") + } + rowType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_RowType{RowType: &pipepb.RowType{Schema: &pipepb.Schema{}}}, + } + if _, err := atomicValueFromProto(rowType, int32Value); err == nil { + t.Errorf("atomicValueFromProto(ROW, int32 value) = nil error, want error") + } + if _, err := atomicValueFromProto(stringType, &pipepb.FieldValue{}); err == nil { + t.Errorf("atomicValueFromProto(STRING, null value) = nil error, want error") + } +} + +func TestAtomicValueToProto_ByteIsSigned(t *testing.T) { + _, fv, err := atomicValueToProto(reflect.ValueOf(uint8(200))) + if err != nil { + t.Fatalf("atomicValueToProto(uint8(200)) = %v, want nil error", err) + } + if got := fv.GetAtomicValue().GetByte(); got != -56 { + t.Errorf("atomicValueToProto(uint8(200)) byte = %v, want -56", got) + } + byteType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_BYTE}, + } + got, err := atomicValueFromProto(byteType, fv) + if err != nil { + t.Fatalf("atomicValueFromProto(BYTE, -56) = %v, want nil error", err) + } + if got.Interface() != uint8(200) { + t.Errorf("atomicValueFromProto(BYTE, -56) = %v, want uint8(200)", got) + } +} + +func TestAtomicValueFromProto_Range(t *testing.T) { + byteType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_BYTE}, + } + int16Type := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_INT16}, + } + byteValue := func(v int32) *pipepb.FieldValue { + return &pipepb.FieldValue{FieldValue: &pipepb.FieldValue_AtomicValue{ + AtomicValue: &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Byte{Byte: v}}, + }} + } + int16Value := func(v int32) *pipepb.FieldValue { + return &pipepb.FieldValue{FieldValue: &pipepb.FieldValue_AtomicValue{ + AtomicValue: &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int16{Int16: v}}, + }} + } + tests := []struct { + name string + ft *pipepb.FieldType + fv *pipepb.FieldValue + want any // nil when the value is out of range. + }{ + {"BYTE min", byteType, byteValue(math.MinInt8), uint8(0x80)}, + {"BYTE max", byteType, byteValue(math.MaxInt8), uint8(0x7f)}, + {"BYTE below min", byteType, byteValue(math.MinInt8 - 1), nil}, + {"BYTE above max", byteType, byteValue(256), nil}, + {"INT16 min", int16Type, int16Value(math.MinInt16), int16(math.MinInt16)}, + {"INT16 max", int16Type, int16Value(math.MaxInt16), int16(math.MaxInt16)}, + {"INT16 below min", int16Type, int16Value(math.MinInt16 - 1), nil}, + {"INT16 above max", int16Type, int16Value(65537), nil}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := atomicValueFromProto(test.ft, test.fv) + if test.want == nil { + if err == nil { + t.Fatalf("atomicValueFromProto(%v) = %v, want error", test.fv, got) + } + return + } + if err != nil { + t.Fatalf("atomicValueFromProto(%v) = %v, want nil error", test.fv, err) + } + if got.Interface() != test.want { + t.Errorf("atomicValueFromProto(%v) = %v, want %v", test.fv, got, test.want) + } + }) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go index 413b1f379a9d..ee08bbb8f549 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go @@ -61,12 +61,12 @@ type Registry struct { logicalTypeProviders map[reflect.Type]LogicalTypeProvider logicalTypeInterfaces []reflect.Type - // Maps logical type identifiers to their reflect.Type and the schema representation. - // the type identifier is the reflect.Type name, and included in the proto as well. + // logicalTypes maps registered logical types by identifier and argument. + // The identifier is included in the proto as the logical type URN. // We don't treat all types as "logical" types. - // ... why don't we treat all types as Logical types? - logicalTypes map[string]LogicalType - logicalTypeIdentifiers map[reflect.Type]string + logicalTypes map[logicalTypeKey]LogicalType + // logicalTypeByGoType maps Go types to their registered logical type. + logicalTypeByGoType map[reflect.Type]LogicalType // toReconcile contains a list of types that have been registered // but not yet processed. Registration actually happens on first @@ -74,6 +74,13 @@ type Registry struct { toReconcile []reflect.Type } +// logicalTypeKey identifies a registered logical type. argument is nil for +// logical types without an argument, and the argument value otherwise. +type logicalTypeKey struct { + identifier string + argument any +} + // NewRegistry creates an initialized LogicalTypeRegistry. func NewRegistry() *Registry { return &Registry{ @@ -81,9 +88,9 @@ func NewRegistry() *Registry { idToType: map[string]reflect.Type{}, syntheticToUser: map[reflect.Type]reflect.Type{}, - logicalTypes: map[string]LogicalType{}, - logicalTypeIdentifiers: map[reflect.Type]string{}, - logicalTypeProviders: map[reflect.Type]LogicalTypeProvider{}, + logicalTypes: map[logicalTypeKey]LogicalType{}, + logicalTypeByGoType: map[reflect.Type]LogicalType{}, + logicalTypeProviders: map[reflect.Type]LogicalTypeProvider{}, } } @@ -98,9 +105,30 @@ func (r *Registry) RegisterLogicalType(lt LogicalType) { if len(lt.ID()) == 0 { panic(fmt.Sprintf("invalid logical type, bad id: %v -> %v", lt.GoType(), lt.StorageType())) } + if lt.argT != nil { + if _, _, err := atomicValueToProto(lt.argV); err != nil { + panic(fmt.Sprintf("LogicalType[%v] has an invalid argument %v: %v", lt.ID(), lt.argV, err)) + } + if !lt.argV.Comparable() || !lt.argV.Equal(lt.argV) { + panic(fmt.Sprintf("LogicalType[%v] has the argument %v which does not compare equal to itself", lt.ID(), lt.argV)) + } + } // TODO add duplication checks. - r.logicalTypeIdentifiers[lt.GoType()] = lt.ID() - r.logicalTypes[lt.ID()] = lt + r.logicalTypeByGoType[lt.GoType()] = lt + r.logicalTypes[lt.key()] = lt +} + +// lookupLogicalType returns the logical type registered for identifier and, +// when argument is valid, that argument. A logical type registered without an +// argument matches any argument. +func (r *Registry) lookupLogicalType(identifier string, argument reflect.Value) (LogicalType, bool) { + if argument.IsValid() && argument.Comparable() { + if lt, ok := r.logicalTypes[logicalTypeKey{identifier: identifier, argument: argument.Interface()}]; ok { + return lt, true + } + } + lt, ok := r.logicalTypes[logicalTypeKey{identifier: identifier}] + return lt, ok } // RegisterLogicalTypeProvider allows registration of providers for interface types. @@ -162,6 +190,29 @@ func ToLogicalType(identifier string, goType, storageType reflect.Type) LogicalT return LogicalType{identifier: identifier, goT: goType, storageT: storageType} } +// ToLogicalTypeWithArgument creates a LogicalType with an argument, for +// parameterized logical types such as a fixed length string and its length. +// The argument is included in the schema representation of the logical type, +// and must be a value of a built in Go type with an atomic schema type other +// than []byte. It is matched by equality, so a floating point NaN is not a +// valid argument. A nil argument creates a logical type without an argument. +func ToLogicalTypeWithArgument(identifier string, goType, storageType reflect.Type, argument any) LogicalType { + lt := ToLogicalType(identifier, goType, storageType) + if argument != nil { + lt.argT = reflect.TypeOf(argument) + lt.argV = reflect.ValueOf(argument) + } + return lt +} + +// key returns the registry key of the logical type. +func (l *LogicalType) key() logicalTypeKey { + if l.argT == nil { + return logicalTypeKey{identifier: l.identifier} + } + return logicalTypeKey{identifier: l.identifier, argument: l.argV.Interface()} +} + func preRegLogicalTypes(r *Registry) { r.RegisterLogicalType(ToLogicalType("int", reflectx.Int, reflectx.Int64)) r.RegisterLogicalType(ToLogicalType("int8", reflectx.Int8, reflectx.Int64)) diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go index 0d44e68285b5..9bd7a287d202 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go @@ -218,8 +218,7 @@ func (r *Registry) registerType(ut reflect.Type, seen map[reflect.Type]struct{}) // Lets do some recursion to register fundamental type parts. t := ut - if lID, ok := r.logicalTypeIdentifiers[t]; ok { - lt := r.logicalTypes[lID] + if lt, ok := r.logicalTypeByGoType[t]; ok { r.addToMaps(lt.StorageType(), t) return nil } @@ -327,16 +326,17 @@ func (r *Registry) FromType(ot reflect.Type) (*pipepb.Schema, error) { return r.fromType(ot) } -func (r *Registry) logicalTypeToFieldType(t reflect.Type) (*pipepb.FieldType, string, error) { +// logicalTypeToFieldType returns the storage field type and the logical type +// of t, or a nil field type if t is not a logical type. +func (r *Registry) logicalTypeToFieldType(t reflect.Type) (*pipepb.FieldType, LogicalType, error) { // Check if a logical type was registered that matches this struct type directly // and if so, extract the schema from it for use. - if lID, ok := r.logicalTypeIdentifiers[t]; ok { - lt := r.logicalTypes[lID] + if lt, ok := r.logicalTypeByGoType[t]; ok { ftype, err := r.reflectTypeToFieldType(lt.StorageType()) if err != nil { - return nil, "", errors.Wrapf(err, "unable to convert LogicalType[%v]'s storage type %v for Go type of %v to a schema", lID, lt.StorageType(), lt.GoType()) + return nil, LogicalType{}, errors.Wrapf(err, "unable to convert LogicalType[%v]'s storage type %v for Go type of %v to a schema", lt.ID(), lt.StorageType(), lt.GoType()) } - return ftype, lID, nil + return ftype, lt, nil } for _, lti := range r.logicalTypeInterfaces { if !t.Implements(lti) { @@ -345,18 +345,42 @@ func (r *Registry) logicalTypeToFieldType(t reflect.Type) (*pipepb.FieldType, st p := r.logicalTypeProviders[lti] st, err := p(t) if err != nil { - return nil, "", errors.Wrapf(err, "unable to convert LogicalType[%v] using provider for %v schema field", t, lti) + return nil, LogicalType{}, errors.Wrapf(err, "unable to convert LogicalType[%v] using provider for %v schema field", t, lti) } if st == nil { continue } ftype, err := r.reflectTypeToFieldType(st) if err != nil { - return nil, "", errors.Wrapf(err, "unable to convert LogicalType[%v]'s storage type %v for Go type of %v to a schema", "interface", st, t) + return nil, LogicalType{}, errors.Wrapf(err, "unable to convert LogicalType[%v]'s storage type %v for Go type of %v to a schema", "interface", st, t) } - return ftype, t.String(), nil + return ftype, ToLogicalType(t.String(), t, st), nil } - return nil, "", nil + return nil, LogicalType{}, nil +} + +// logicalStorageSchema returns a copy of the storage row schema of a logical +// type used as a top level type. +func logicalStorageSchema(ftype *pipepb.FieldType, lt LogicalType, ot reflect.Type) (*pipepb.Schema, error) { + schm := ftype.GetRowType().GetSchema() + if schm == nil { + return nil, errors.Errorf("LogicalType[%v] for %v has the non-row storage type %v and cannot be used as a top level schema", lt.ID(), ot, lt.StorageType()) + } + return proto.Clone(schm).(*pipepb.Schema), nil +} + +// appendLogicalOptions appends the options that let ToType recover the +// logical type of a top level schema. +func appendLogicalOptions(schm *pipepb.Schema, lt LogicalType) error { + schm.Options = append(schm.Options, logicalOption(lt.ID())) + if lt.ArgumentType() != nil { + opt, err := logicalArgumentOption(lt) + if err != nil { + return errors.Wrapf(err, "unable to convert the argument of LogicalType[%v]", lt.ID()) + } + schm.Options = append(schm.Options, opt) + } + return nil } // fromType handles if the initial type is a pointer or not WRT lookups against @@ -367,18 +391,20 @@ func (r *Registry) fromType(ot reflect.Type) (*pipepb.Schema, error) { if schm, ok := r.typeToSchema[ot]; ok { return schm, nil } - ftype, lID, err := r.logicalTypeToFieldType(ot) + ftype, lt, err := r.logicalTypeToFieldType(ot) if err != nil { return nil, err } if ftype != nil { - schm := ftype.GetRowType().GetSchema() - schm = proto.Clone(schm).(*pipepb.Schema) + schm, err := logicalStorageSchema(ftype, lt, ot) + if err != nil { + return nil, err + } if ot.Kind() == reflect.Ptr { schm.Options = append(schm.Options, optGoNillable()) } - if lID != "" { - schm.Options = append(schm.Options, logicalOption(lID)) + if err := appendLogicalOptions(schm, lt); err != nil { + return nil, err } schm.Id = getUUID(ot) r.typeToSchema[ot] = schm @@ -413,6 +439,9 @@ const ( // optGoLogical indicates that this top level schema has a logical type equivalent that need to be looked up. // It has a value type of String representing the URN for the logical type to look up. optGoLogicalUrn = "beam:schema:go:logical:v1" + // optGoLogicalArgument holds the argument of the logical type of this top level schema. + // Its value type is the argument type of the logical type. + optGoLogicalArgumentUrn = "beam:schema:go:logical_argument:v1" ) func optGoNillable() *pipepb.Option { @@ -498,6 +527,43 @@ func fromLogicalOption(opts []*pipepb.Option) (string, bool) { return lID, true } +// logicalArgumentOption returns the option holding the argument of a +// parameterized logical type. +func logicalArgumentOption(lt LogicalType) (*pipepb.Option, error) { + ft, fv, err := atomicValueToProto(lt.ArgumentValue()) + if err != nil { + return nil, err + } + return &pipepb.Option{ + Name: optGoLogicalArgumentUrn, + Type: ft, + Value: fv, + }, nil +} + +// fromLogicalArgumentOption returns the logical type argument of this top +// level type, or an invalid Value when the schema has no argument option. +func fromLogicalArgumentOption(opts []*pipepb.Option) (reflect.Value, error) { + o := checkOptions(opts, optGoLogicalArgumentUrn) + if o == nil { + return reflect.Value{}, nil + } + return atomicValueFromProto(o.GetType(), o.GetValue()) +} + +// logicalTypeArgument returns the argument of a logical type proto as a Go +// value. It returns an invalid Value when the logical type has no argument, +// or when the argument is not of an atomic type. +func logicalTypeArgument(lst *pipepb.LogicalType) (reflect.Value, error) { + if lst.GetArgumentType() == nil || lst.GetArgument().GetFieldValue() == nil { + return reflect.Value{}, nil + } + if _, ok := lst.GetArgumentType().GetTypeInfo().(*pipepb.FieldType_AtomicType); !ok { + return reflect.Value{}, nil + } + return atomicValueFromProto(lst.GetArgumentType(), lst.GetArgument()) +} + func (r *Registry) structToSchema(t reflect.Type) (*pipepb.Schema, error) { if t.Kind() != reflect.Struct { return nil, errors.Errorf("non struct type received in structToSchema: %v is kind %v", t, t.Kind()) @@ -506,14 +572,18 @@ func (r *Registry) structToSchema(t reflect.Type) (*pipepb.Schema, error) { return schm, nil } - ftype, lID, err := r.logicalTypeToFieldType(t) + ftype, lt, err := r.logicalTypeToFieldType(t) if err != nil { return nil, err } if ftype != nil { - schm := ftype.GetRowType().GetSchema() - schm = proto.Clone(schm).(*pipepb.Schema) - schm.Options = append(schm.Options, logicalOption(lID)) + schm, err := logicalStorageSchema(ftype, lt, t) + if err != nil { + return nil, err + } + if err := appendLogicalOptions(schm, lt); err != nil { + return nil, err + } schm.Id = getUUID(t) r.typeToSchema[t] = schm r.idToType[schm.GetId()] = t @@ -566,18 +636,26 @@ func (r *Registry) structFieldToField(sf reflect.StructField) (*pipepb.Field, er } func (r *Registry) reflectTypeToFieldType(ot reflect.Type) (*pipepb.FieldType, error) { - ftype, lID, err := r.logicalTypeToFieldType(ot) + ftype, lt, err := r.logicalTypeToFieldType(ot) if err != nil { return nil, err } if ftype != nil { + lst := &pipepb.LogicalType{ + Urn: lt.ID(), + Representation: ftype, + } + if lt.ArgumentType() != nil { + at, av, err := atomicValueToProto(lt.ArgumentValue()) + if err != nil { + return nil, errors.Wrapf(err, "unable to convert the argument of LogicalType[%v]", lt.ID()) + } + lst.ArgumentType = at + lst.Argument = av + } return &pipepb.FieldType{ TypeInfo: &pipepb.FieldType_LogicalType{ - LogicalType: &pipepb.LogicalType{ - Urn: lID, - Representation: ftype, - // TODO(BEAM-9615): Handle type Arguments. - }, + LogicalType: lst, }, }, nil } @@ -701,8 +779,13 @@ func (r *Registry) toType(s *pipepb.Schema) (reflect.Type, error) { if t, ok := r.idToType[s.GetId()]; ok { return t, nil } - if lID, ok := fromLogicalOption(s.GetOptions()); ok { - if lt, ok := r.logicalTypes[lID]; ok { + lID, isLogical := fromLogicalOption(s.GetOptions()) + if isLogical { + arg, err := fromLogicalArgumentOption(s.GetOptions()) + if err != nil { + return nil, errors.Wrapf(err, "invalid argument option for logical type %v", lID) + } + if lt, ok := r.lookupLogicalType(lID, arg); ok { return lt.GoType(), nil } } @@ -720,7 +803,11 @@ func (r *Registry) toType(s *pipepb.Schema) (reflect.Type, error) { } ret := reflect.StructOf(fields) if ut, ok := r.syntheticToUser[ret]; ok { - ret = ut + // The storage row of an unregistered logical type must not resolve to + // the Go type of a registered one. + if !isLogical || !r.isLogicalGoType(ut) { + ret = ut + } } if t := nillableFromOptions(s.GetOptions(), ret); t != nil { return t, nil @@ -728,6 +815,17 @@ func (r *Registry) toType(s *pipepb.Schema) (reflect.Type, error) { return ret, nil } +// isLogicalGoType reports whether ut is the Go type of a registered logical +// type, after skipping a pointer as addToMaps does. +func (r *Registry) isLogicalGoType(ut reflect.Type) bool { + for gt := range r.logicalTypeByGoType { + if reflectx.SkipPtr(gt) == ut { + return true + } + } + return false +} + func (r *Registry) fieldToStructField(sf *pipepb.Field) (reflect.StructField, error) { name := sf.GetName() rt, err := r.fieldTypeToReflectType(sf.GetType(), sf.Options) @@ -794,8 +892,15 @@ func (r *Registry) fieldTypeToReflectType(sft *pipepb.FieldType, opts []*pipepb. case *pipepb.FieldType_LogicalType: lst := sft.GetLogicalType() identifier := lst.GetUrn() - lt, ok := r.logicalTypes[identifier] + arg, err := logicalTypeArgument(lst) + if err != nil { + return nil, errors.Wrapf(err, "invalid argument for logical type %v", identifier) + } + lt, ok := r.lookupLogicalType(identifier, arg) if !ok { + if arg.IsValid() { + return nil, errors.Errorf("unknown logical type: %v with argument %v", identifier, arg.Interface()) + } return nil, errors.Errorf("unknown logical type: %v", identifier) } t = lt.GoType() diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go index 367d70e81d17..56386954629d 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go @@ -19,10 +19,12 @@ package schema import ( "fmt" + "math" "reflect" "testing" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx" pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" "github.com/google/go-cmp/cmp" "google.golang.org/protobuf/encoding/prototext" @@ -90,7 +92,24 @@ type nonRegisteredLogical struct { k int32 } +// parameterizedString is a logical type with an argument. The tests register +// it with the argument int32(5). +type parameterizedString string + +// parameterizedRow is a logical type with an argument and a row storage type, +// so it can be used as a top level type. +type parameterizedRow struct { + S int64 +} + +// anyString is a logical type registered without an argument. +type anyString string + var ( + parameterizedStringType = reflect.TypeOf((*parameterizedString)(nil)).Elem() + parameterizedRowType = reflect.TypeOf((*parameterizedRow)(nil)).Elem() + anyStringType = reflect.TypeOf((*anyString)(nil)).Elem() + unexportedFieldsType = reflect.TypeOf((*unexportedFields)(nil)).Elem() exportedFuncType = reflect.TypeOf((*exportedFunc)(nil)) anotherType = reflect.TypeOf((*anotherStruct)(nil)).Elem() @@ -754,6 +773,53 @@ func TestSchemaConversion(t *testing.T) { rt: reflect.TypeOf(struct { U nonRegisteredLogical }{}), + }, { + st: &pipepb.Schema{ + Fields: []*pipepb.Field{ + { + Name: "P", + Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_LogicalType{ + LogicalType: &pipepb.LogicalType{ + Urn: "schema.parameterizedString", + Representation: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: pipepb.AtomicType_STRING, + }, + }, + ArgumentType: int32FieldType(), + Argument: int32FieldValue(5), + }, + }, + }, + }, + }, + }, + rt: reflect.TypeOf(struct { + P parameterizedString + }{}), + }, { + st: &pipepb.Schema{ + Options: []*pipepb.Option{ + logicalOption("schema.parameterizedRow"), + { + Name: optGoLogicalArgumentUrn, + Type: int32FieldType(), + Value: int32FieldValue(3), + }, + }, + Fields: []*pipepb.Field{ + { + Name: "S", + Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: pipepb.AtomicType_INT64, + }, + }, + }, + }, + }, + rt: parameterizedRowType, }, } @@ -764,6 +830,8 @@ func TestSchemaConversion(t *testing.T) { preRegLogicalTypes(reg) reg.RegisterLogicalType(ToLogicalType(exportedFuncType.Elem().String(), exportedFuncType.Elem(), reflect.TypeOf(struct{ V int16 }{}))) reg.RegisterLogicalType(ToLogicalType(nonRegisteredLogicalType.String(), nonRegisteredLogicalType, reflect.TypeOf(struct{ K int32 }{}))) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.parameterizedString", parameterizedStringType, reflectx.String, int32(5))) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.parameterizedRow", parameterizedRowType, reflect.TypeOf(struct{ S int64 }{}), int32(3))) reg.RegisterType(reflect.TypeOf((*sRegisteredType)(nil))) reg.RegisterLogicalTypeProvider(reflect.TypeOf((*testInterface)(nil)).Elem(), func(t reflect.Type) (reflect.Type, error) { switch t { @@ -829,3 +897,269 @@ func hasEmbeddedField(t reflect.Type) bool { } return false } + +func int32FieldType() *pipepb.FieldType { + return &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: pipepb.AtomicType_INT32, + }, + } +} + +func int32FieldValue(v int32) *pipepb.FieldValue { + return &pipepb.FieldValue{ + FieldValue: &pipepb.FieldValue_AtomicValue{ + AtomicValue: &pipepb.AtomicTypeValue{ + Value: &pipepb.AtomicTypeValue_Int32{Int32: v}, + }, + }, + } +} + +func TestToType_LogicalTypeArguments(t *testing.T) { + reg := NewRegistry() + preRegLogicalTypes(reg) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.parameterizedString", parameterizedStringType, reflectx.String, int32(5))) + reg.RegisterLogicalType(ToLogicalType("schema.anyString", anyStringType, reflectx.String)) + + stringType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: pipepb.AtomicType_STRING, + }, + } + rowType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_RowType{ + RowType: &pipepb.RowType{ + Schema: &pipepb.Schema{ + Fields: []*pipepb.Field{{Name: "precision", Type: int32FieldType()}}, + }, + }, + }, + } + rowValue := &pipepb.FieldValue{ + FieldValue: &pipepb.FieldValue_RowValue{ + RowValue: &pipepb.Row{Values: []*pipepb.FieldValue{int32FieldValue(5)}}, + }, + } + + tests := []struct { + name string + lt *pipepb.LogicalType + want reflect.Type + wantErr bool + }{{ + name: "matching argument", + lt: &pipepb.LogicalType{Urn: "schema.parameterizedString", Representation: stringType, ArgumentType: int32FieldType(), Argument: int32FieldValue(5)}, + want: parameterizedStringType, + }, { + name: "unknown argument", + lt: &pipepb.LogicalType{Urn: "schema.parameterizedString", Representation: stringType, ArgumentType: int32FieldType(), Argument: int32FieldValue(7)}, + wantErr: true, + }, { + name: "missing argument", + lt: &pipepb.LogicalType{Urn: "schema.parameterizedString", Representation: stringType}, + wantErr: true, + }, { + name: "argument ignored for logical type registered without one", + lt: &pipepb.LogicalType{Urn: "schema.anyString", Representation: stringType, ArgumentType: int32FieldType(), Argument: int32FieldValue(7)}, + want: anyStringType, + }, { + name: "non atomic argument ignored", + lt: &pipepb.LogicalType{Urn: "schema.anyString", Representation: stringType, ArgumentType: rowType, Argument: rowValue}, + want: anyStringType, + }, { + name: "unknown urn", + lt: &pipepb.LogicalType{Urn: "schema.unknown", Representation: stringType}, + wantErr: true, + }} + for i, test := range tests { + t.Run(test.name, func(t *testing.T) { + s := &pipepb.Schema{ + Id: fmt.Sprintf("arguments-%d", i), + Fields: []*pipepb.Field{{ + Name: "f", + Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_LogicalType{LogicalType: test.lt}, + }, + }}, + } + got, err := reg.ToType(s) + if test.wantErr { + if err == nil { + t.Fatalf("ToType(%v) = %v, want error", prototext.Format(s), got) + } + return + } + if err != nil { + t.Fatalf("ToType(%v) = %v, want nil error", prototext.Format(s), err) + } + if ft := got.Field(0).Type; ft != test.want { + t.Errorf("ToType(%v) field type = %v, want %v", prototext.Format(s), ft, test.want) + } + }) + } +} + +// The multi types share one URN with different arguments, and multiAny is +// registered for the URN without an argument. +type multiOne struct{ S int64 } +type multiTwo struct{ S int64 } +type multiAny struct{ S int64 } + +// multiOnePtr is a defined pointer type, used as a logical Go type. +type multiOnePtr *multiOne + +func TestLogicalTypeArgumentPrecedence(t *testing.T) { + reg := NewRegistry() + preRegLogicalTypes(reg) + storage := reflect.TypeOf(struct{ S int64 }{}) + oneType := reflect.TypeOf(multiOne{}) + twoType := reflect.TypeOf(multiTwo{}) + anyType := reflect.TypeOf(multiAny{}) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.multi", oneType, storage, int32(1))) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.multi", twoType, storage, int32(2))) + reg.RegisterLogicalType(ToLogicalType("schema.multi", anyType, storage)) + + rowType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_RowType{RowType: &pipepb.RowType{Schema: &pipepb.Schema{ + Fields: []*pipepb.Field{{Name: "S", Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_INT64}, + }}}, + }}}, + } + fieldSchema := func(id string, argument *pipepb.FieldValue) *pipepb.Schema { + lt := &pipepb.LogicalType{Urn: "schema.multi", Representation: rowType} + if argument != nil { + lt.ArgumentType = int32FieldType() + lt.Argument = argument + } + return &pipepb.Schema{Id: id, Fields: []*pipepb.Field{{Name: "f", Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_LogicalType{LogicalType: lt}, + }}}} + } + topLevelSchema := func(id string, argument *pipepb.FieldValue) *pipepb.Schema { + opts := []*pipepb.Option{logicalOption("schema.multi")} + if argument != nil { + opts = append(opts, &pipepb.Option{Name: optGoLogicalArgumentUrn, Type: int32FieldType(), Value: argument}) + } + return &pipepb.Schema{Id: id, Options: opts, Fields: rowType.GetRowType().GetSchema().GetFields()} + } + + tests := []struct { + name string + s *pipepb.Schema + want reflect.Type + }{ + {"field argument 1", fieldSchema("field-1", int32FieldValue(1)), oneType}, + {"field argument 2", fieldSchema("field-2", int32FieldValue(2)), twoType}, + {"field argument 3 falls back", fieldSchema("field-3", int32FieldValue(3)), anyType}, + {"field without argument", fieldSchema("field-none", nil), anyType}, + {"top level argument 2", topLevelSchema("top-2", int32FieldValue(2)), twoType}, + {"top level argument 3 falls back", topLevelSchema("top-3", int32FieldValue(3)), anyType}, + {"top level without argument", topLevelSchema("top-none", nil), anyType}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := reg.ToType(test.s) + if err != nil { + t.Fatalf("ToType(%v) = %v, want nil error", prototext.Format(test.s), err) + } + if len(test.s.GetOptions()) == 0 { + got = got.Field(0).Type + } + if got != test.want { + t.Errorf("ToType(%v) = %v, want exactly %v", prototext.Format(test.s), got, test.want) + } + }) + } + for _, rt := range []reflect.Type{oneType, twoType, anyType} { + schm, err := reg.FromType(rt) + if err != nil { + t.Fatalf("FromType(%v) = %v, want nil error", rt, err) + } + got, err := reg.ToType(schm) + if err != nil { + t.Fatalf("ToType(FromType(%v)) = %v, want nil error", rt, err) + } + if got != rt { + t.Errorf("ToType(FromType(%v)) = %v, want exactly %v", rt, got, rt) + } + } +} + +// A top level schema whose logical type is not registered keeps its storage +// row type, even when a registered logical type has the same storage. +func TestToType_UnregisteredTopLevelLogicalType(t *testing.T) { + storage := reflect.TypeOf(struct{ S int64 }{}) + fields := []*pipepb.Field{{Name: "S", Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_INT64}, + }}} + tests := []struct { + name string + opts []*pipepb.Option + }{ + {"unknown argument", []*pipepb.Option{ + logicalOption("schema.multi"), + {Name: optGoLogicalArgumentUrn, Type: int32FieldType(), Value: int32FieldValue(2)}, + }}, + {"missing argument", []*pipepb.Option{logicalOption("schema.multi")}}, + {"unknown urn", []*pipepb.Option{logicalOption("schema.unknown")}}, + } + for _, goType := range []reflect.Type{reflect.TypeOf(multiOne{}), reflect.TypeOf(&multiOne{}), reflect.TypeOf(multiOnePtr(nil))} { + reg := NewRegistry() + preRegLogicalTypes(reg) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.multi", goType, storage, int32(1))) + reg.RegisterType(goType) + for _, test := range tests { + for _, nillable := range []bool{false, true} { + name := fmt.Sprintf("%v/%v/nillable=%v", goType, test.name, nillable) + t.Run(name, func(t *testing.T) { + opts, want := test.opts, storage + if nillable { + opts, want = append(append([]*pipepb.Option{}, opts...), optGoNillable()), reflect.PtrTo(storage) + } + s := &pipepb.Schema{Id: "unregistered-" + name, Options: opts, Fields: fields} + got, err := reg.ToType(s) + if err != nil { + t.Fatalf("ToType(%v) = %v, want nil error", prototext.Format(s), err) + } + if got != want { + t.Errorf("ToType(%v) = %v, want %v", prototext.Format(s), got, want) + } + }) + } + } + } +} + +type bytesBacked struct{ B []byte } + +func TestRegisterLogicalType_Validation(t *testing.T) { + reg := NewRegistry() + preRegLogicalTypes(reg) + bytesBackedType := reflect.TypeOf(bytesBacked{}) + reg.RegisterLogicalType(ToLogicalType("schema.bytesBacked", bytesBackedType, reflectx.ByteSlice)) + if schm, err := reg.FromType(bytesBackedType); err == nil { + t.Errorf("FromType(%v) = %v, want error for a top level logical type with non row storage", bytesBackedType, prototext.Format(schm)) + } + if _, err := reg.FromType(reflect.TypeOf(struct{ F bytesBacked }{})); err != nil { + t.Errorf("FromType(struct{F bytesBacked}) = %v, want nil error", err) + } + + for name, argument := range map[string]any{ + "NaN": math.NaN(), + "float32 NaN": float32(math.NaN()), + "[]byte": []byte("x"), + "int": int(1), + "struct": struct{ A int32 }{}, + } { + t.Run(name, func(t *testing.T) { + defer func() { + if recover() == nil { + t.Errorf("RegisterLogicalType with argument %v did not panic", argument) + } + }() + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.invalid", anyStringType, reflectx.String, argument)) + }) + } +}