diff --git a/arrow/extensions/variant.go b/arrow/extensions/variant.go index fee2e046a..0a691452d 100644 --- a/arrow/extensions/variant.go +++ b/arrow/extensions/variant.go @@ -122,10 +122,10 @@ func createShreddedField(dt arrow.DataType) arrow.DataType { // ), Nullable: true}) // // This is intended to be a convenient way to create a shredded variant type from a definition -// of the fields to shred. If the provided data type is nil, it will create a default -// variant type. +// of the fields to shred. If the provided data type is nil or Null (including an +// extension whose storage is Null), it will create a default variant type. func NewShreddedVariantType(dt arrow.DataType) *VariantType { - if dt == nil { + if dt == nil || isNullType(dt) { return NewDefaultVariantType() } @@ -220,9 +220,9 @@ func NewVariantType(storage arrow.DataType) (*VariantType, error) { return nil, fmt.Errorf("%w: typed_value field must be nullable, got %s", arrow.ErrInvalid, typedValueField.Type) } - dt := typedValueField.Type - if dt.ID() == arrow.EXTENSION { - dt = dt.(arrow.ExtensionType).StorageType() + dt := storageType(typedValueField.Type) + if dt.ID() == arrow.NULL { + return nil, fmt.Errorf("%w: typed_value field must not be null type", arrow.ErrInvalid) } if nt, ok := dt.(arrow.NestedType); ok { @@ -291,11 +291,25 @@ func isBinary(dt arrow.DataType) bool { dt.ID() == arrow.BINARY_VIEW } +func storageType(dt arrow.DataType) arrow.DataType { + if ext, ok := dt.(arrow.ExtensionType); ok { + return ext.StorageType() + } + return dt +} + +func isNullType(dt arrow.DataType) bool { + return storageType(dt).ID() == arrow.NULL +} + func validStruct(s *arrow.StructType) bool { switch s.NumFields() { case 1: f := s.Field(0) - return (f.Name == "value" && isBinary(f.Type)) || f.Name == "typed_value" + if f.Name == "value" { + return isBinary(f.Type) + } + return f.Name == "typed_value" && !isNullType(f.Type) case 2: valField, ok := s.FieldByName("value") if !ok || !valField.Nullable || !isBinary(valField.Type) { @@ -311,7 +325,7 @@ func validStruct(s *arrow.StructType) bool { return validNestedType(nt) } - return true + return !isNullType(typedField.Type) default: return false } diff --git a/arrow/extensions/variant_test.go b/arrow/extensions/variant_test.go index a39fd5131..d2cc01249 100644 --- a/arrow/extensions/variant_test.go +++ b/arrow/extensions/variant_test.go @@ -83,6 +83,16 @@ func TestVariantExtensionType(t *testing.T) { arrow.Field{Name: "metadata", Type: arrow.BinaryTypes.String, Nullable: false}, arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: false}), "metadata field must be non-nullable binary type, got utf8"}, + {arrow.StructOf( + arrow.Field{Name: "metadata", Type: arrow.BinaryTypes.Binary, Nullable: false}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: true}, + arrow.Field{Name: "typed_value", Type: arrow.Null, Nullable: true}), + "typed_value field must not be null type"}, + {arrow.StructOf( + arrow.Field{Name: "metadata", Type: arrow.BinaryTypes.Binary, Nullable: false}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: true}, + arrow.Field{Name: "typed_value", Type: extensions.NewOpaqueType(arrow.Null, "null", "test"), Nullable: true}), + "typed_value field must not be null type"}, } for _, tt := range tests { @@ -113,6 +123,20 @@ func TestVariantExtensionBadNestedTypes(t *testing.T) { ), Nullable: false})}, {"empty struct elem", arrow.StructOf( arrow.Field{Name: "foobar", Type: arrow.StructOf(), Nullable: false})}, + {"null typed_value in shredded field", arrow.StructOf( + arrow.Field{Name: "foobar", Type: arrow.StructOf( + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: true}, + arrow.Field{Name: "typed_value", Type: arrow.Null, Nullable: true}, + ), Nullable: false})}, + {"null typed_value extension in one-field shredded field", arrow.StructOf( + arrow.Field{Name: "foobar", Type: arrow.StructOf( + arrow.Field{Name: "typed_value", Type: extensions.NewOpaqueType(arrow.Null, "null", "test"), Nullable: true}, + ), Nullable: false})}, + {"null typed_value extension in two-field shredded field", arrow.StructOf( + arrow.Field{Name: "foobar", Type: arrow.StructOf( + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: true}, + arrow.Field{Name: "typed_value", Type: extensions.NewOpaqueType(arrow.Null, "null", "test"), Nullable: true}, + ), Nullable: false})}, {"non-nullable two elem struct", arrow.StructOf( arrow.Field{Name: "foobar", Type: arrow.StructOf( arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: true}, @@ -1560,6 +1584,10 @@ func TestVariantBuilderUnmarshalJSON(t *testing.T) { func TestNewSimpleShreddedVariantType(t *testing.T) { assert.True(t, arrow.TypeEqual(extensions.NewDefaultVariantType(), extensions.NewShreddedVariantType(nil))) + assert.True(t, arrow.TypeEqual(extensions.NewDefaultVariantType(), + extensions.NewShreddedVariantType(arrow.Null))) + assert.True(t, arrow.TypeEqual(extensions.NewDefaultVariantType(), + extensions.NewShreddedVariantType(extensions.NewOpaqueType(arrow.Null, "null", "test")))) vt := extensions.NewShreddedVariantType(arrow.PrimitiveTypes.Float32) s := arrow.StructOf(