diff --git a/arrow/util/protobuf_reflect.go b/arrow/util/protobuf_reflect.go index e8b80c683..434f42c8f 100644 --- a/arrow/util/protobuf_reflect.go +++ b/arrow/util/protobuf_reflect.go @@ -251,10 +251,19 @@ func (pfr *ProtobufFieldReflection) asUnion() protobufUnionReflection { } func (pur protobufUnionReflection) isThisOne() bool { - for pur.rValue.Kind() == reflect.Ptr || pur.rValue.Kind() == reflect.Interface { - pur.rValue = pur.rValue.Elem() + rValue := pur.rValue + for rValue.IsValid() && (rValue.Kind() == reflect.Ptr || rValue.Kind() == reflect.Interface) { + if rValue.IsNil() { + return false + } + rValue = rValue.Elem() } - return pur.rValue.Field(0).String() == pur.prValue.String() + + if !rValue.IsValid() || !pur.prValue.IsValid() || rValue.Kind() != reflect.Struct || rValue.NumField() == 0 { + return false + } + + return rValue.Field(0).String() == pur.prValue.String() } func (pur protobufUnionReflection) whichOne() arrow.UnionTypeCode { @@ -366,7 +375,10 @@ func (pfr *ProtobufFieldReflection) asMap() protobufMapReflection { func (pmr protobufMapReflection) getDataType() arrow.DataType { for kvp := range pmr.generateKeyValuePairs() { - return kvp.getDataType() + if kvp.err != nil { + continue + } + return kvp.kvp.getDataType() } return protobufMapKeyValuePairReflection{ k: ProtobufFieldReflection{ @@ -387,12 +399,17 @@ type protobufMapKeyValuePairReflection struct { v ProtobufFieldReflection } +type protobufMapKeyValuePairResult struct { + kvp protobufMapKeyValuePairReflection + err error +} + func (pmr protobufMapKeyValuePairReflection) getDataType() arrow.DataType { return arrow.MapOf(pmr.k.getDataType(), pmr.v.getDataType()) } -func (pmr protobufMapReflection) generateKeyValuePairs() chan protobufMapKeyValuePairReflection { - out := make(chan protobufMapKeyValuePairReflection) +func (pmr protobufMapReflection) generateKeyValuePairs() chan protobufMapKeyValuePairResult { + out := make(chan protobufMapKeyValuePairResult) go func() { defer close(out) @@ -409,45 +426,50 @@ func (pmr protobufMapReflection) generateKeyValuePairs() chan protobufMapKeyValu schemaOptions: pmr.schemaOptions, }, } - out <- kvp + out <- protobufMapKeyValuePairResult{kvp: kvp} return } for _, k := range pmr.rValue.MapKeys() { + mapKey, err := getMapKey(k) + if err != nil { + out <- protobufMapKeyValuePairResult{err: err} + continue + } kvp := protobufMapKeyValuePairReflection{ k: ProtobufFieldReflection{ parent: pmr.parent, descriptor: pmr.descriptor.MapKey(), - prValue: getMapKey(k), + prValue: mapKey, rValue: k, schemaOptions: pmr.schemaOptions, }, v: ProtobufFieldReflection{ parent: pmr.parent, descriptor: pmr.descriptor.MapValue(), - prValue: pmr.prValue.Map().Get(protoreflect.MapKey(getMapKey(k))), + prValue: pmr.prValue.Map().Get(protoreflect.MapKey(mapKey)), rValue: pmr.rValue.MapIndex(k), schemaOptions: pmr.schemaOptions, }, } - out <- kvp + out <- protobufMapKeyValuePairResult{kvp: kvp} } }() return out } -func getMapKey(v reflect.Value) protoreflect.Value { +func getMapKey(v reflect.Value) (protoreflect.Value, error) { switch v.Kind() { case reflect.String: - return protoreflect.ValueOf(v.String()) + return protoreflect.ValueOf(v.String()), nil case reflect.Int32, reflect.Int64: - return protoreflect.ValueOf(v.Int()) + return protoreflect.ValueOf(v.Int()), nil case reflect.Bool: - return protoreflect.ValueOf(v.Bool()) + return protoreflect.ValueOf(v.Bool()), nil case reflect.Uint32, reflect.Uint64: - return protoreflect.ValueOf(v.Uint()) + return protoreflect.ValueOf(v.Uint()), nil default: - panic("Unmapped protoreflect map key type") + return protoreflect.Value{}, fmt.Errorf("unsupported map key kind %s", v.Type()) } } @@ -520,14 +542,20 @@ func (psr ProtobufMessageReflection) getFieldByName(n string) *ProtobufFieldRefl if fv.IsValid() { if !fv.IsZero() { for fv.Kind() == reflect.Ptr || fv.Kind() == reflect.Interface { + if fv.IsNil() { + fv = reflect.Value{} + break + } fv = fv.Elem() } - if fd.ContainingOneof() != nil { - n = string(fd.ContainingOneof().Name()) - } - fv = fv.FieldByName(strcase.UpperCamelCase(n)) - for fv.Kind() == reflect.Ptr { - fv = fv.Elem() + if fv.IsValid() { + if fd.ContainingOneof() != nil { + n = string(fd.ContainingOneof().Name()) + } + fv = fv.FieldByName(strcase.UpperCamelCase(n)) + for fv.Kind() == reflect.Ptr { + fv = fv.Elem() + } } } } @@ -858,12 +886,16 @@ func (f ProtobufMessageFieldReflection) AppendValueOrNull(b array.Builder, mem m Field: f.Field.Type.(*arrow.MapType).ItemField(), } for kvp := range f.asMap().generateKeyValuePairs() { - k.protobufReflection = &kvp.k + if kvp.err != nil { + return fmt.Errorf("failed to append map field %q (%s): %w", f.name(), f.Field.Type, kvp.err) + } + + k.protobufReflection = &kvp.kvp.k err := k.AppendValueOrNull(mb.KeyBuilder(), mem) if err != nil { return err } - v.protobufReflection = &kvp.v + v.protobufReflection = &kvp.kvp.v err = v.AppendValueOrNull(mb.ItemBuilder(), mem) if err != nil { return err diff --git a/arrow/util/protobuf_reflect_test.go b/arrow/util/protobuf_reflect_test.go index 35ef15692..8638d29c3 100644 --- a/arrow/util/protobuf_reflect_test.go +++ b/arrow/util/protobuf_reflect_test.go @@ -19,6 +19,7 @@ package util import ( "encoding/json" "fmt" + "reflect" "testing" "github.com/apache/arrow-go/v18/arrow" @@ -495,3 +496,85 @@ func TestAppendValueOrNull(t *testing.T) { want := "not able to appendValueOrNull for type TIME32" assert.EqualErrorf(t, got, want, "Error is: %v, want: %v", got, want) } + +func TestGetMapKeyRejectsUnsupportedType(t *testing.T) { + _, err := getMapKey(reflect.ValueOf(struct{}{})) + require.Error(t, err) + assert.ErrorContains(t, err, "unsupported map key kind") +} + +func TestMapAppendReturnsContextualErrorForUnsupportedKey(t *testing.T) { + msg := util_message.AllTheTypesNoAny{SimpleMap: map[int32]string{1: "value"}} + pmr := NewProtobufMessageReflection(&msg) + + var field *ProtobufMessageFieldReflection + for i := range pmr.fields { + if pmr.fields[i].name() == "simple_map" { + field = &pmr.fields[i] + break + } + } + require.NotNil(t, field) + fr := field.protobufReflection.(*ProtobufFieldReflection) + + badMapReflection := &protobufMapReflection{ + ProtobufFieldReflection: ProtobufFieldReflection{ + parent: fr.parent, + descriptor: fr.descriptor, + prValue: fr.prValue, + rValue: reflect.ValueOf(map[struct{}]string{struct{}{}: "value"}), + schemaOptions: fr.schemaOptions, + }, + } + badField := ProtobufMessageFieldReflection{ + parent: field.parent, + protobufReflection: badMapReflection, + Field: field.Field, + } + + schema := arrow.NewSchema([]arrow.Field{badField.Field}, nil) + mem := memory.NewGoAllocator() + recordBuilder := array.NewRecordBuilder(mem, schema) + + err := badField.AppendValueOrNull(recordBuilder.Field(0), mem) + require.Error(t, err) + assert.ErrorContains(t, err, "failed to append map field") + assert.ErrorContains(t, err, "simple_map") +} + +func TestUnionReflectionHandlesUnsetAndNilOneofValues(t *testing.T) { + tests := []struct { + name string + msg proto.Message + expectedActive arrow.UnionTypeCode + }{ + {"unset", &util_message.AllTheTypes{}, -1}, + {"nil pointer", &util_message.AllTheTypes{Oneof: (*util_message.AllTheTypes_Oneofmessage)(nil)}, -1}, + {"valid", &util_message.AllTheTypes{Oneof: &util_message.AllTheTypes_Oneofstring{Oneofstring: "value"}}, 0}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + pmr := NewProtobufMessageReflection(test.msg, WithOneOfHandler(OneOfDenseUnion)) + var field *ProtobufMessageFieldReflection + for i := range pmr.fields { + if pmr.fields[i].name() == "oneof" { + field = &pmr.fields[i] + break + } + } + require.NotNil(t, field) + + fr := field.protobufReflection.(*ProtobufFieldReflection) + assert.NotPanics(t, func() { + gotActive := fr.asUnion().whichOne() + assert.Equal(t, test.expectedActive, gotActive) + if test.expectedActive == -1 { + assert.Nil(t, fr.asUnion().getField()) + } else { + assert.NotNil(t, fr.asUnion().getField()) + } + }) + }) + } +}