fix: guard protobuf oneof/union reflection on nil oneof values

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

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())
+				}
+			})
+		})
+	}
+}