Eliaaazzz commented on code in PR #40153:
URL: https://github.com/apache/beam/pull/40153#discussion_r4253557678


##########
sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go:
##########
@@ -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)

Review Comment:
   Thanks for catching this. Fixed in #40544. #40154 also added one call in 
row_test.go, and I replaced it there as well.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to