lidavidm commented on code in PR #395: URL: https://github.com/apache/arrow-go/pull/395#discussion_r2119773219
########## arrow/extensions/variant.go: ########## @@ -0,0 +1,1536 @@ +// 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 extensions + +import ( + "fmt" + "math" + "reflect" + "strings" + "sync" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/internal/debug" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/json" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" +) + +// VariantType is the arrow extension type for representing Variant values as +// defined by the Parquet Variant specification for encoding and shredding values. +// The underlying storage must be a struct type with a minimum of two fields +// ("metadata" and "value") and an optional third field ("typed_value"). +// +// See the documentation for [NewVariantType] for the rules for creating a variant +// type. +type VariantType struct { + arrow.ExtensionBase + + metadataFieldIdx int + valueFieldIdx int + typedValueFieldIdx int +} + +// NewDefaultVariantType creates a basic, non-shredded variant type. The underlying +// storage type will be struct<metadata: binary required, value: binary required>. +func NewDefaultVariantType() *VariantType { + s := arrow.StructOf( + arrow.Field{Name: "metadata", Type: arrow.BinaryTypes.Binary, Nullable: false}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: false}) + + vt, _ := NewVariantType(s) + return vt +} + +// NewVariantType creates a new variant type based on the provided storage type. +// +// The rules for a variant storage type are: +// 1. MUST be a struct +// 2. MUST have required field named "metadata" that is binary/largebinary/binary_view +// 3. Must satisfy exactly one of the following: +// a. MUST have required field named "value" that is binary/largebinary/binary_view +// b. MUST have an optional field named "value" that is binary/largebinary/binary_view +// and another optional field named "typed_value" that is either a primitive type or +// a list/large_list/list_view or struct which also satisfies the following requirements: +// i. The elements must be REQUIRED +// ii. There must either be a single REQUIRED field named "value" which is +// binary/largebinary/binary_view or havve an optional "value" field and an optional +// "typed_value" field that follows the rules laid out in (b). +// +// The metadata field may also be dictionary encoded +func NewVariantType(storage arrow.DataType) (*VariantType, error) { + s, ok := storage.(*arrow.StructType) + if !ok { + return nil, fmt.Errorf("%w: bad storage type %s for variant type", arrow.ErrInvalid, storage) + } + + var ( + metadataFieldIdx = -1 + valueFieldIdx = -1 + typedValueFieldIdx = -1 + ) + + if metadataFieldIdx, ok = s.FieldIdx("metadata"); !ok { + return nil, fmt.Errorf("%w: missing required field 'metadata' in variant storage type %s", arrow.ErrInvalid, storage) + } + + if valueFieldIdx, ok = s.FieldIdx("value"); !ok { + return nil, fmt.Errorf("%w: missing required field 'value' in variant storage type %s", arrow.ErrInvalid, storage) + } + + if s.NumFields() > 3 { + return nil, fmt.Errorf("%w: too many fields in variant storage type %s, expected 2 or 3", arrow.ErrInvalid, storage) + } + + if s.NumFields() == 3 { + if typedValueFieldIdx, ok = s.FieldIdx("typed_value"); !ok { + return nil, fmt.Errorf("%w: has 3 fields, but missing 'typed_value' field, %s", arrow.ErrInvalid, storage) + } + } + + mdField, valField := s.Field(metadataFieldIdx), s.Field(valueFieldIdx) + if mdField.Nullable { + return nil, fmt.Errorf("%w: metadata field must be non-nullable binary type, got %s", arrow.ErrInvalid, mdField.Type) + } + + if !isBinary(mdField.Type) { + if mdField.Type.ID() != arrow.DICTIONARY || !isBinary(mdField.Type.(*arrow.DictionaryType).ValueType) { + return nil, fmt.Errorf("%w: metadata field must be non-nullable binary type, got %s", arrow.ErrInvalid, mdField.Type) + } + } + + if !isBinary(valField.Type) || (valField.Nullable && typedValueFieldIdx == -1) { + return nil, fmt.Errorf("%w: value field must be non-nullable binary type, got %s", arrow.ErrInvalid, valField.Type) + } + + if typedValueFieldIdx == -1 { + return &VariantType{ + ExtensionBase: arrow.ExtensionBase{Storage: storage}, + metadataFieldIdx: metadataFieldIdx, + valueFieldIdx: valueFieldIdx, + typedValueFieldIdx: -1, + }, nil + } + + valueField := s.Field(valueFieldIdx) + if !valueField.Nullable { + return nil, fmt.Errorf("%w: value field must be nullable if typed_value is present, got %s", arrow.ErrInvalid, valueField.Type) + } + + typedValueField := s.Field(typedValueFieldIdx) + if !typedValueField.Nullable { + return nil, fmt.Errorf("%w: typed_value field must be nullable, got %s", arrow.ErrInvalid, typedValueField.Type) + } + + if nt, ok := typedValueField.Type.(arrow.NestedType); ok { + if !validNestedType(nt) { + return nil, fmt.Errorf("%w: typed_value field must be a valid nested type, got %s", arrow.ErrInvalid, typedValueField.Type) + } + } + + return &VariantType{ + ExtensionBase: arrow.ExtensionBase{Storage: storage}, + metadataFieldIdx: metadataFieldIdx, + valueFieldIdx: valueFieldIdx, + typedValueFieldIdx: typedValueFieldIdx, + }, nil +} + +func (*VariantType) ArrayType() reflect.Type { + return reflect.TypeOf(VariantArray{}) +} + +func (v *VariantType) Metadata() arrow.Field { + return v.StorageType().(*arrow.StructType).Field(v.metadataFieldIdx) +} + +func (v *VariantType) Value() arrow.Field { + return v.StorageType().(*arrow.StructType).Field(v.valueFieldIdx) +} + +func (*VariantType) ExtensionName() string { return "parquet.variant" } + +func (v *VariantType) String() string { + return fmt.Sprintf("extension<%s>", v.ExtensionName()) +} + +func (v *VariantType) ExtensionEquals(other arrow.ExtensionType) bool { + return v.ExtensionName() == other.ExtensionName() && + arrow.TypeEqual(v.Storage, other.StorageType()) +} + +func (*VariantType) Serialize() string { return "" } +func (*VariantType) Deserialize(storageType arrow.DataType, _ string) (arrow.ExtensionType, error) { + return NewVariantType(storageType) +} + +func (*VariantType) ParquetLogicalType() schema.LogicalType { + return schema.VariantLogicalType{} +} + +func (v *VariantType) NewBuilder(mem memory.Allocator) array.Builder { + return NewVariantBuilder(mem, v) +} + +func isBinary(dt arrow.DataType) bool { + return dt.ID() == arrow.BINARY || dt.ID() == arrow.LARGE_BINARY || + dt.ID() == arrow.BINARY_VIEW +} Review Comment: There's IsBinaryLike already, but that doesn't cover BINARY_VIEW (should it?) ########## parquet/variant/builder.go: ########## @@ -848,3 +882,27 @@ func (b *Builder) Build() (Value, error) { }, }, nil } + +type variantPrimitiveType interface { + constraints.Integer | constraints.Float | string | []byte | + arrow.Date32 | arrow.Time64 | arrow.Timestamp | bool | + uuid.UUID | DecimalValue[decimal.Decimal32] | + DecimalValue[decimal.Decimal64] | DecimalValue[decimal.Decimal128] +} + +// Encode is a convenience function that produces the encoded bytes for a primitive +// variant value. At the moment this is just delegating to the [Builder.Append] method, +// but in the future it will be optimized to avoid the extra overhead and reduce allocations. +func Encode[T variantPrimitiveType](v T, opt ...AppendOpt) ([]byte, error) { Review Comment: Should variantPrimitiveType be public if it's used in a public signature? ########## arrow/array/dictionary.go: ########## @@ -1652,6 +1652,32 @@ func UnifyTableDicts(alloc memory.Allocator, table arrow.Table) (arrow.Table, er return NewTable(table.Schema(), cols, table.NumRows()), nil } +type dictWrapper[T arrow.ValueType] struct { + *Dictionary + + typedDict arrow.TypedArray[T] +} + +// NewDictWrapper creates a simple wrapper around a Dictionary array provides the +// a Value method which will use the underlying dictionary to return the value Review Comment: ```suggestion // NewDictWrapper creates a simple wrapper around a Dictionary array that provides // a Value method which will use the underlying dictionary to return the value ``` ########## arrow/extensions/variant.go: ########## @@ -0,0 +1,1536 @@ +// 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 extensions + +import ( + "fmt" + "math" + "reflect" + "strings" + "sync" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/internal/debug" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/json" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" +) + +// VariantType is the arrow extension type for representing Variant values as +// defined by the Parquet Variant specification for encoding and shredding values. +// The underlying storage must be a struct type with a minimum of two fields +// ("metadata" and "value") and an optional third field ("typed_value"). +// +// See the documentation for [NewVariantType] for the rules for creating a variant +// type. +type VariantType struct { + arrow.ExtensionBase + + metadataFieldIdx int + valueFieldIdx int + typedValueFieldIdx int +} + +// NewDefaultVariantType creates a basic, non-shredded variant type. The underlying +// storage type will be struct<metadata: binary required, value: binary required>. +func NewDefaultVariantType() *VariantType { + s := arrow.StructOf( + arrow.Field{Name: "metadata", Type: arrow.BinaryTypes.Binary, Nullable: false}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: false}) + + vt, _ := NewVariantType(s) + return vt +} + +// NewVariantType creates a new variant type based on the provided storage type. +// +// The rules for a variant storage type are: +// 1. MUST be a struct +// 2. MUST have required field named "metadata" that is binary/largebinary/binary_view +// 3. Must satisfy exactly one of the following: +// a. MUST have required field named "value" that is binary/largebinary/binary_view +// b. MUST have an optional field named "value" that is binary/largebinary/binary_view +// and another optional field named "typed_value" that is either a primitive type or +// a list/large_list/list_view or struct which also satisfies the following requirements: +// i. The elements must be REQUIRED +// ii. There must either be a single REQUIRED field named "value" which is +// binary/largebinary/binary_view or havve an optional "value" field and an optional +// "typed_value" field that follows the rules laid out in (b). +// +// The metadata field may also be dictionary encoded +func NewVariantType(storage arrow.DataType) (*VariantType, error) { + s, ok := storage.(*arrow.StructType) + if !ok { + return nil, fmt.Errorf("%w: bad storage type %s for variant type", arrow.ErrInvalid, storage) + } + + var ( + metadataFieldIdx = -1 + valueFieldIdx = -1 + typedValueFieldIdx = -1 + ) + + if metadataFieldIdx, ok = s.FieldIdx("metadata"); !ok { + return nil, fmt.Errorf("%w: missing required field 'metadata' in variant storage type %s", arrow.ErrInvalid, storage) + } + + if valueFieldIdx, ok = s.FieldIdx("value"); !ok { + return nil, fmt.Errorf("%w: missing required field 'value' in variant storage type %s", arrow.ErrInvalid, storage) + } + + if s.NumFields() > 3 { + return nil, fmt.Errorf("%w: too many fields in variant storage type %s, expected 2 or 3", arrow.ErrInvalid, storage) + } + + if s.NumFields() == 3 { + if typedValueFieldIdx, ok = s.FieldIdx("typed_value"); !ok { + return nil, fmt.Errorf("%w: has 3 fields, but missing 'typed_value' field, %s", arrow.ErrInvalid, storage) + } + } + + mdField, valField := s.Field(metadataFieldIdx), s.Field(valueFieldIdx) + if mdField.Nullable { + return nil, fmt.Errorf("%w: metadata field must be non-nullable binary type, got %s", arrow.ErrInvalid, mdField.Type) + } + + if !isBinary(mdField.Type) { + if mdField.Type.ID() != arrow.DICTIONARY || !isBinary(mdField.Type.(*arrow.DictionaryType).ValueType) { + return nil, fmt.Errorf("%w: metadata field must be non-nullable binary type, got %s", arrow.ErrInvalid, mdField.Type) + } + } + + if !isBinary(valField.Type) || (valField.Nullable && typedValueFieldIdx == -1) { + return nil, fmt.Errorf("%w: value field must be non-nullable binary type, got %s", arrow.ErrInvalid, valField.Type) + } + + if typedValueFieldIdx == -1 { + return &VariantType{ + ExtensionBase: arrow.ExtensionBase{Storage: storage}, + metadataFieldIdx: metadataFieldIdx, + valueFieldIdx: valueFieldIdx, + typedValueFieldIdx: -1, + }, nil + } + + valueField := s.Field(valueFieldIdx) + if !valueField.Nullable { + return nil, fmt.Errorf("%w: value field must be nullable if typed_value is present, got %s", arrow.ErrInvalid, valueField.Type) + } + + typedValueField := s.Field(typedValueFieldIdx) + if !typedValueField.Nullable { + return nil, fmt.Errorf("%w: typed_value field must be nullable, got %s", arrow.ErrInvalid, typedValueField.Type) + } + + if nt, ok := typedValueField.Type.(arrow.NestedType); ok { + if !validNestedType(nt) { + return nil, fmt.Errorf("%w: typed_value field must be a valid nested type, got %s", arrow.ErrInvalid, typedValueField.Type) + } + } + + return &VariantType{ + ExtensionBase: arrow.ExtensionBase{Storage: storage}, + metadataFieldIdx: metadataFieldIdx, + valueFieldIdx: valueFieldIdx, + typedValueFieldIdx: typedValueFieldIdx, + }, nil +} + +func (*VariantType) ArrayType() reflect.Type { + return reflect.TypeOf(VariantArray{}) +} + +func (v *VariantType) Metadata() arrow.Field { + return v.StorageType().(*arrow.StructType).Field(v.metadataFieldIdx) +} + +func (v *VariantType) Value() arrow.Field { + return v.StorageType().(*arrow.StructType).Field(v.valueFieldIdx) +} + +func (*VariantType) ExtensionName() string { return "parquet.variant" } + +func (v *VariantType) String() string { + return fmt.Sprintf("extension<%s>", v.ExtensionName()) +} + +func (v *VariantType) ExtensionEquals(other arrow.ExtensionType) bool { + return v.ExtensionName() == other.ExtensionName() && + arrow.TypeEqual(v.Storage, other.StorageType()) +} + +func (*VariantType) Serialize() string { return "" } +func (*VariantType) Deserialize(storageType arrow.DataType, _ string) (arrow.ExtensionType, error) { + return NewVariantType(storageType) +} + +func (*VariantType) ParquetLogicalType() schema.LogicalType { + return schema.VariantLogicalType{} +} + +func (v *VariantType) NewBuilder(mem memory.Allocator) array.Builder { + return NewVariantBuilder(mem, v) +} + +func isBinary(dt arrow.DataType) bool { + return dt.ID() == arrow.BINARY || dt.ID() == arrow.LARGE_BINARY || + dt.ID() == arrow.BINARY_VIEW +} + +func validStruct(s *arrow.StructType) bool { + switch s.NumFields() { + case 1: + f := s.Field(0) + return f.Name == "value" && !f.Nullable && isBinary(f.Type) + case 2: + valField, ok := s.FieldByName("value") + if !ok || !valField.Nullable || !isBinary(valField.Type) { + return false + } + + typedField, ok := s.FieldByName("typed_value") + if !ok || !typedField.Nullable { + return false + } + + if nt, ok := typedField.Type.(arrow.NestedType); ok && nt.Name() != "extension" { + return validNestedType(nt) + } + + return true + default: + return false + } +} + +func validNestedType(dt arrow.NestedType) bool { + switch t := dt.(type) { + case arrow.ListLikeType: + if t.ElemField().Nullable { + return false + } + + s, ok := t.Elem().(*arrow.StructType) + if !ok { + return false + } + + return validStruct(s) + case *arrow.StructType: + if t.NumFields() == 0 { + return false + } + + for i := range t.NumFields() { + f := t.Field(i) + if f.Nullable { + return false + } + + s, ok := f.Type.(*arrow.StructType) + if !ok { + return false + } + + if !validStruct(s) { + return false + } + } + + return true + default: + return false + } +} + +// VariantArray is an extension Array type containing Variant values which may +// potentially be shredded into multiple fields. +type VariantArray struct { + array.ExtensionArrayBase + + rdr variantReader + rdrErr error + initRdr sync.Once +} + +func (v *VariantArray) initReader() { + // initialize a reader that coalesces shredded fields back into a variant + // or just returns the basic variants if the array is not shredded. + v.initRdr.Do(func() { + vt := v.ExtensionType().(*VariantType) + st := v.Storage().(*array.Struct) + metaField := st.Field(vt.metadataFieldIdx) + valueField := st.Field(vt.valueFieldIdx) + + metadata, ok := metaField.(arrow.TypedArray[[]byte]) + if !ok { + // we already validated that if the metadata field isn't a binary + // type directly, it must be a dictionary with a binary value type. + metadata, _ = array.NewDictWrapper[[]byte](metaField.(*array.Dictionary)) + } + + if vt.typedValueFieldIdx == -1 { + v.rdr = &basicVariantReader{ + metadata: metadata, + value: valueField.(arrow.TypedArray[[]byte]), + } + return + } + + ivreader, err := getReader(st.Field(vt.typedValueFieldIdx)) + if err != nil { + v.rdrErr = err + return + } + + v.rdr = &shreddedVariantReader{ + metadata: metadata, + value: valueField.(arrow.TypedArray[[]byte]), + typedValue: ivreader, + } + }) +} + +// Metadata returns the metadata column of the variant array, containing the +// metadata for each variant value. +func (v *VariantArray) Metadata() arrow.TypedArray[[]byte] { + vt := v.ExtensionType().(*VariantType) + return v.Storage().(*array.Struct).Field(vt.metadataFieldIdx).(arrow.TypedArray[[]byte]) +} + +// UntypedValues returns the untyped variant values for each element of the array, +// if the array is not shredded this will contain the variant bytes for each value. +// If the array is shredded, this will contain any variant values that are either +// partially shredded objects or are not shredded at all (e.g. a value that doesnt +// match the types of the shredding). +// +// The shredded array and the untyped values array together are used to encode a +// single value. If this is not encoding shredded object fields, then a given index +// will never be null in both arrays. (A null value will be an encoded null variant value +// in this array with a null in the shredded array). +// +// If both arrays are null for a given index (only valid for shredded object fields), +// it means that the value is missing entirely (as opposed to existing and having a +// value of null). +func (v *VariantArray) UntypedValues() arrow.TypedArray[[]byte] { + vt := v.ExtensionType().(*VariantType) + return v.Storage().(*array.Struct).Field(vt.valueFieldIdx).(arrow.TypedArray[[]byte]) +} + +// Shredded returns the typed array for the shredded values of the variant array, +// following the rules of the Parquet Variant specification. As such, this array will +// always be either a struct, a list, or a primitive array. +// +// The reason for exposing this is to allow users to quickly access one of the shredded +// fields without having to decode the entire variant value. +func (v *VariantArray) Shredded() arrow.Array { + vt := v.ExtensionType().(*VariantType) + if vt.typedValueFieldIdx == -1 { + return nil + } + + return v.Storage().(*array.Struct).Field(vt.typedValueFieldIdx) +} + +// IsShredded returns true if the variant has shredded columns. +func (v *VariantArray) IsShredded() bool { + return v.ExtensionType().(*VariantType).typedValueFieldIdx != -1 +} + +// IsNull will also take into account the special case where there is an +// encoded null variant in the untyped values array for this index and return +// appropriately. +func (v *VariantArray) IsNull(i int) bool { + if v.Storage().IsNull(i) { + return true + } + + vt := v.ExtensionType().(*VariantType) + valArr := v.Storage().(*array.Struct).Field(vt.valueFieldIdx) + if vt.typedValueFieldIdx != -1 { + typedArr := v.Storage().(*array.Struct).Field(vt.typedValueFieldIdx) + if !typedArr.IsNull(i) { + return false + } + } + + b := valArr.(arrow.TypedArray[[]byte]).Value(i) + return len(b) == 1 && b[0] == 0 // variant null +} + +func (v *VariantArray) IsValid(i int) bool { + return !v.IsNull(i) +} + +func (v *VariantArray) String() string { + o := new(strings.Builder) + o.WriteString("VariantArray[") + for i := 0; i < v.Len(); i++ { + if i > 0 { + o.WriteString(" ") + } + if v.IsNull(i) { + o.WriteString(array.NullValueStr) + continue + } + + val, err := v.Value(i) + if err != nil { + o.WriteString(fmt.Sprintf("error: %v", err)) + continue + } + + o.WriteString(val.String()) + } + o.WriteString("]") + return o.String() +} + +func (v *VariantArray) Value(i int) (variant.Value, error) { + v.initReader() + if v.rdrErr != nil { + return variant.Value{}, v.rdrErr + } + + return v.rdr.Value(i) +} + +func (v *VariantArray) Values() ([]variant.Value, error) { + values := make([]variant.Value, v.Len()) + for i := range v.Len() { + val, err := v.Value(i) + if err != nil { + return nil, fmt.Errorf("error getting value at index %d: %w", i, err) + } + values[i] = val + } + return values, nil +} + +func (v *VariantArray) ValueStr(i int) string { + if v.IsNull(i) { + return array.NullValueStr + } + + val, err := v.Value(i) + if err != nil { + return fmt.Sprintf("error: %v", err) + } + + return val.String() +} + +func (v *VariantArray) MarshalJSON() ([]byte, error) { + values := make([]any, v.Len()) + for i := range v.Len() { + if v.IsNull(i) { + values[i] = nil + continue + } + + val, err := v.Value(i) + if err != nil { + values[i] = fmt.Sprintf("error: %v", err) + continue + } + + values[i] = val.Value() + } + return json.Marshal(values) +} + +func (v *VariantArray) GetOneForMarshal(i int) any { + if v.IsNull(i) { + return nil + } + + val, err := v.Value(i) + if err != nil { + return fmt.Sprintf("error: %v", err) + } + + return val.Value() +} + +type variantReader interface { + IsNull(i int) bool + Value(i int) (variant.Value, error) +} + +type basicVariantReader struct { + metadata arrow.TypedArray[[]byte] + value arrow.TypedArray[[]byte] +} + +func (r *basicVariantReader) IsNull(i int) bool { + if r.value.IsNull(i) { + return true + } + + // special case for null variant + val := r.value.Value(i) + return len(val) == 1 && val[0] == 0 +} + +func (r *basicVariantReader) Value(i int) (variant.Value, error) { + if r.IsNull(i) { + return variant.NullValue, nil + } + + meta := r.metadata.Value(i) + val := r.value.Value(i) + + return variant.New(meta, val) +} + +func createPrimitiveVariantReader(arr arrow.Array) (typedValReader, error) { + switch a := arr.(type) { + case *array.Boolean: + return asVariantReader[bool]{typedVal: a}, nil + case *array.Int8: + return asVariantReader[int8]{typedVal: a}, nil + case *array.Uint8: + return asVariantReader[uint8]{typedVal: a}, nil + case *array.Int16: + return asVariantReader[int16]{typedVal: a}, nil + case *array.Uint16: + return asVariantReader[uint16]{typedVal: a}, nil + case *array.Int32: + return asVariantReader[int32]{typedVal: a}, nil + case *array.Uint32: + return asVariantReader[uint32]{typedVal: a}, nil + case *array.Int64: + return asVariantReader[int64]{typedVal: a}, nil + case *array.Float32: + return asVariantReader[float32]{typedVal: a}, nil + case *array.Float64: + return asVariantReader[float64]{typedVal: a}, nil + case array.StringLike: + return asVariantReader[string]{typedVal: a}, nil + case arrow.TypedArray[[]byte]: + return asVariantReader[[]byte]{typedVal: a}, nil + case *array.Date32: + return asVariantReader[arrow.Date32]{typedVal: a}, nil + case *array.Time64: + if a.DataType().(*arrow.Time64Type).Unit != arrow.Microsecond { + return nil, fmt.Errorf("%w: unsupported time64 unit %s for variant", + arrow.ErrInvalid, a.DataType().(*arrow.Time64Type).Unit) + } + return asVariantReader[arrow.Time64]{typedVal: a}, nil + case *array.Timestamp: + var opt variant.AppendOpt + dt := a.DataType().(*arrow.TimestampType) + switch dt.Unit { + case arrow.Microsecond: + case arrow.Nanosecond: + opt |= variant.OptTimestampNano + default: + return nil, fmt.Errorf("%w: unsupported timestamp unit %s for variant", + arrow.ErrInvalid, a.DataType().(*arrow.TimestampType).Unit) + } + + if dt.TimeZone == "UTC" { + opt |= variant.OptTimestampUTC + } + + return asVariantReader[arrow.Timestamp]{typedVal: a, opts: opt}, nil + case *UUIDArray: + return asVariantReader[uuid.UUID]{typedVal: a}, nil + case *array.Decimal32: + return asVariantReader[variant.DecimalValue[decimal.Decimal32]]{ + typedVal: decimal32Wrapper{ + Decimal32: a, + scale: uint8(a.DataType().(*arrow.Decimal32Type).Scale), + }, + }, nil + case *array.Decimal64: + return asVariantReader[variant.DecimalValue[decimal.Decimal64]]{ + typedVal: decimal64Wrapper{ + Decimal64: a, + scale: uint8(a.DataType().(*arrow.Decimal64Type).Scale), + }, + }, nil + case *array.Decimal128: + return asVariantReader[variant.DecimalValue[decimal.Decimal128]]{ + typedVal: decimal128Wrapper{ + Decimal128: a, + scale: uint8(a.DataType().(*arrow.Decimal128Type).Scale), + }, + }, nil + } + + return nil, fmt.Errorf("%w: unsupported primitive type %s for variant", + arrow.ErrInvalid, arr.DataType().String()) +} + +type decimal32Wrapper struct { + *array.Decimal32 + + scale uint8 +} + +func (d decimal32Wrapper) Value(i int) variant.DecimalValue[decimal.Decimal32] { + return variant.DecimalValue[decimal.Decimal32]{ + Value: d.Decimal32.Value(i), + Scale: d.scale, + } +} + +type decimal64Wrapper struct { + *array.Decimal64 + + scale uint8 +} + +func (d decimal64Wrapper) Value(i int) variant.DecimalValue[decimal.Decimal64] { + return variant.DecimalValue[decimal.Decimal64]{ + Value: d.Decimal64.Value(i), + Scale: d.scale, + } +} + +type decimal128Wrapper struct { + *array.Decimal128 + + scale uint8 +} + +func (d decimal128Wrapper) Value(i int) variant.DecimalValue[decimal.Decimal128] { + return variant.DecimalValue[decimal.Decimal128]{ + Value: d.Decimal128.Value(i), + Scale: d.scale, + } +} + +type arrowVariantPrimitiveType interface { + arrow.NumericType | bool | string | []byte | uuid.UUID | + variant.DecimalValue[decimal.Decimal32] | + variant.DecimalValue[decimal.Decimal64] | + variant.DecimalValue[decimal.Decimal128] +} + +type typedArr[T arrowVariantPrimitiveType] interface { + arrow.Array + Value(int) T +} + +type asVariantReader[T arrowVariantPrimitiveType] struct { + typedVal typedArr[T] + + opts variant.AppendOpt +} + +func (vr asVariantReader[T]) IsNull(i int) bool { + return vr.typedVal.IsNull(i) +} + +func (vr asVariantReader[T]) Value(_ variant.Metadata, i int) (any, error) { + if vr.typedVal.IsNull(i) { + return nil, nil + } + + return variant.Encode(vr.typedVal.Value(i), vr.opts) +} + +func getReader(typedArr arrow.Array) (typedValReader, error) { + switch arr := typedArr.(type) { + case *array.Struct: + fieldReaders := make(map[string]fieldReaderPair) + fieldList := arr.DataType().(*arrow.StructType).Fields() + for i := range arr.NumField() { + child := arr.Field(i).(*array.Struct) + childType := child.DataType().(*arrow.StructType) + + valueIdx, _ := childType.FieldIdx("value") + valueArr := child.Field(valueIdx).(arrow.TypedArray[[]byte]) + + typedValueIdx, _ := childType.FieldIdx("typed_value") + typedRdr, err := getReader(child.Field(typedValueIdx)) + if err != nil { + return nil, fmt.Errorf("error getting typed value reader for field %s: %w", fieldList[i].Name, err) + } + + fieldReaders[fieldList[i].Name] = fieldReaderPair{ + values: valueArr, + typedVal: typedRdr, + } + } + + return &typedObjReader{ + objArr: arr, + fieldRdrs: fieldReaders, + }, nil + case array.ListLike: + listValues := arr.ListValues().(*array.Struct) + elemType := listValues.DataType().(*arrow.StructType) + valueIdx, _ := elemType.FieldIdx("value") + valueArr := listValues.Field(valueIdx).(arrow.TypedArray[[]byte]) + + typedValueIdx, _ := elemType.FieldIdx("typed_value") + typedRdr, err := getReader(listValues.Field(typedValueIdx)) + if err != nil { + return nil, fmt.Errorf("error getting typed value reader: %w", err) + } + + return &typedListReader{ + listArr: arr, + valueArr: valueArr, + typedVal: typedRdr, + }, nil + default: + return createPrimitiveVariantReader(arr) + } +} + +type typedPair struct { + Value []byte + TypedValue any +} + +func constructVariant(b *variant.Builder, meta variant.Metadata, value []byte, typedVal any) error { + switch v := typedVal.(type) { + case nil: + if len(value) == 0 { + return nil + } + + return b.AppendEncoded(value) + case map[string]typedPair: + fields := make([]variant.FieldEntry, 0, len(v)) + objstart := b.Offset() + for k, pair := range v { + if pair.TypedValue != nil || len(pair.Value) != 0 { + fields = append(fields, b.NextField(objstart, k)) + if err := constructVariant(b, meta, pair.Value, pair.TypedValue); err != nil { + return err + } + } + } + + if len(value) > 0 { + obj, err := variant.NewWithMetadata(meta, value) + if err != nil { + return err + } + + objval, ok := obj.Value().(variant.ObjectValue) + if !ok { + return fmt.Errorf("%w: expected object value, got %T", arrow.ErrInvalid, obj.Value()) + } + + for key, field := range objval.Values() { + fields = append(fields, b.NextField(objstart, key)) + if err := b.AppendEncoded(field.Bytes()); err != nil { + return fmt.Errorf("error appending field %s: %w", key, err) + } + } + } + + return b.FinishObject(objstart, fields) + case []typedPair: + debug.Assert(len(value) == 0, "shredded array must not conflict with variant value") + + elems := make([]int, 0, len(v)) + arrstart := b.Offset() + for _, pair := range v { + elems = append(elems, b.NextElement(arrstart)) + if err := constructVariant(b, meta, pair.Value, pair.TypedValue); err != nil { + return err + } + } + + return b.FinishArray(arrstart, elems) + case []byte: + return b.AppendEncoded(v) + default: + return fmt.Errorf("%w: unsupported typed value type %T for variant", arrow.ErrInvalid, v) + } +} + +type typedValReader interface { + Value(meta variant.Metadata, i int) (any, error) + IsNull(i int) bool +} + +type fieldReaderPair struct { + values arrow.TypedArray[[]byte] + typedVal typedValReader +} + +type typedObjReader struct { + objArr *array.Struct + fieldRdrs map[string]fieldReaderPair +} + +func (v *typedObjReader) IsNull(i int) bool { + return v.objArr.IsNull(i) +} + +func (v *typedObjReader) Value(meta variant.Metadata, i int) (any, error) { + if v.objArr.IsNull(i) { + return nil, nil + } + + result := make(map[string]typedPair) + for name, rdr := range v.fieldRdrs { + typedValue, err := rdr.typedVal.Value(meta, i) + if err != nil { + return nil, fmt.Errorf("error reading typed value for field %s at index %d: %w", name, i, err) + } + result[name] = typedPair{ + Value: rdr.values.Value(i), + TypedValue: typedValue, + } + } + return result, nil +} + +type typedListReader struct { + listArr array.ListLike + + valueArr arrow.TypedArray[[]byte] + typedVal typedValReader +} + +func (v *typedListReader) IsNull(i int) bool { + return v.listArr.IsNull(i) +} + +func (v *typedListReader) Value(meta variant.Metadata, i int) (any, error) { + if v.listArr.IsNull(i) { + return nil, nil + } + + start, end := v.listArr.ValueOffsets(i) + if start == end { + return []typedPair{}, nil + } + + result := make([]typedPair, 0, end-start) + for j := start; j < end; j++ { + val := v.valueArr.Value(int(j)) + typedValue, err := v.typedVal.Value(meta, int(j)) + if err != nil { + return nil, fmt.Errorf("error reading typed value at index %d: %w", j, err) + } + + result = append(result, typedPair{ + Value: val, + TypedValue: typedValue, + }) + } + + return result, nil +} + +type shreddedVariantReader struct { + metadata arrow.TypedArray[[]byte] + value arrow.TypedArray[[]byte] + + typedValue typedValReader +} + +func (v *shreddedVariantReader) IsNull(i int) bool { + if !v.typedValue.IsNull(i) { + return false + } + + if v.value.IsNull(i) { + return true + } + + val := v.value.Value(i) + return len(val) == 1 && val[0] == 0 // variant null +} + +func (v *shreddedVariantReader) Value(i int) (variant.Value, error) { + metaBytes := v.metadata.Value(i) + meta, err := variant.NewMetadata(metaBytes) + if err != nil { + return variant.NullValue, fmt.Errorf("error reading metadata at index %d: %w", i, err) + } + + b := variant.NewBuilderFromMeta(meta) + typed, err := v.typedValue.Value(meta, i) + if err != nil { + return variant.NullValue, fmt.Errorf("error reading typed value at index %d: %w", i, err) + } + + if err := constructVariant(b, meta, v.value.Value(i), typed); err != nil { + return variant.NullValue, fmt.Errorf("error constructing variant at index %d: %w", i, err) + } + return b.Build() +} + +// VariantBuilder is an array builder for both shredded or non-shredded variant extension +// arrays. It allows you to append variant values, and will appropriately shred them +// if it is able to do so based on the underlying storage type. +type VariantBuilder struct { + *array.ExtensionBuilder + shreddedSchema shreddedSchema + + structBldr *array.StructBuilder + metaBldr array.BinaryLikeBuilder + valueBldr array.BinaryLikeBuilder + typedBldr shreddedBuilder +} + +type binaryDictBuilderAdapter struct { + *array.BinaryDictionaryBuilder +} + +func (b *binaryDictBuilderAdapter) ReserveData(int) {} + +func (b *binaryDictBuilderAdapter) AppendValues(v [][]byte, valids []bool) { + if len(valids) == 0 { + for _, val := range v { + b.Append(val) + } + return + } + + for i, valid := range valids { + if !valid { + b.AppendNull() + } else { + b.Append(v[i]) + } + } +} + +func (b *binaryDictBuilderAdapter) UnsafeAppend(v []byte) { + b.Append(v) +} + +func (b *binaryDictBuilderAdapter) Append(v []byte) { + if err := b.BinaryDictionaryBuilder.Append(v); err != nil { + panic(fmt.Sprintf("error appending value %s to binary dictionary builder: %v", string(v), err)) Review Comment: Why does Append panic? Shouldn't it return an error and let UnsafeAppend panic? ########## arrow/extensions/variant.go: ########## @@ -0,0 +1,1536 @@ +// 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 extensions + +import ( + "fmt" + "math" + "reflect" + "strings" + "sync" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/internal/debug" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/json" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" +) + +// VariantType is the arrow extension type for representing Variant values as +// defined by the Parquet Variant specification for encoding and shredding values. +// The underlying storage must be a struct type with a minimum of two fields +// ("metadata" and "value") and an optional third field ("typed_value"). +// +// See the documentation for [NewVariantType] for the rules for creating a variant +// type. +type VariantType struct { + arrow.ExtensionBase + + metadataFieldIdx int + valueFieldIdx int + typedValueFieldIdx int +} + +// NewDefaultVariantType creates a basic, non-shredded variant type. The underlying +// storage type will be struct<metadata: binary required, value: binary required>. +func NewDefaultVariantType() *VariantType { + s := arrow.StructOf( + arrow.Field{Name: "metadata", Type: arrow.BinaryTypes.Binary, Nullable: false}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: false}) + + vt, _ := NewVariantType(s) + return vt +} + +// NewVariantType creates a new variant type based on the provided storage type. +// +// The rules for a variant storage type are: +// 1. MUST be a struct +// 2. MUST have required field named "metadata" that is binary/largebinary/binary_view +// 3. Must satisfy exactly one of the following: +// a. MUST have required field named "value" that is binary/largebinary/binary_view +// b. MUST have an optional field named "value" that is binary/largebinary/binary_view +// and another optional field named "typed_value" that is either a primitive type or +// a list/large_list/list_view or struct which also satisfies the following requirements: +// i. The elements must be REQUIRED +// ii. There must either be a single REQUIRED field named "value" which is +// binary/largebinary/binary_view or havve an optional "value" field and an optional Review Comment: ```suggestion // binary/largebinary/binary_view or have an optional "value" field and an optional ``` ########## parquet/variant/builder.go: ########## @@ -777,6 +791,18 @@ func (b *Builder) FinishObject(start int, fields []FieldEntry) error { return nil } +// AppendEncoded is a special case where we directly append a pre-encoded variant +// value. Its keys must already be in the dictionary and v must already be +// a properly encoded variant value. No checking is performed here currently, so +// be careful as this can easily lead to an invalid variant result. +func (b *Builder) AppendEncoded(v []byte) error { Review Comment: Should this be named UnsafeApppendEncoded perhaps? I guess it still returns an error, but only in the case of Write itself failing... ########## arrow/extensions/variant.go: ########## @@ -0,0 +1,1536 @@ +// 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 extensions + +import ( + "fmt" + "math" + "reflect" + "strings" + "sync" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/internal/debug" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/json" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" +) + +// VariantType is the arrow extension type for representing Variant values as +// defined by the Parquet Variant specification for encoding and shredding values. +// The underlying storage must be a struct type with a minimum of two fields +// ("metadata" and "value") and an optional third field ("typed_value"). +// +// See the documentation for [NewVariantType] for the rules for creating a variant +// type. +type VariantType struct { + arrow.ExtensionBase + + metadataFieldIdx int + valueFieldIdx int + typedValueFieldIdx int +} + +// NewDefaultVariantType creates a basic, non-shredded variant type. The underlying +// storage type will be struct<metadata: binary required, value: binary required>. +func NewDefaultVariantType() *VariantType { + s := arrow.StructOf( + arrow.Field{Name: "metadata", Type: arrow.BinaryTypes.Binary, Nullable: false}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: false}) + + vt, _ := NewVariantType(s) + return vt +} + +// NewVariantType creates a new variant type based on the provided storage type. +// +// The rules for a variant storage type are: +// 1. MUST be a struct +// 2. MUST have required field named "metadata" that is binary/largebinary/binary_view +// 3. Must satisfy exactly one of the following: +// a. MUST have required field named "value" that is binary/largebinary/binary_view +// b. MUST have an optional field named "value" that is binary/largebinary/binary_view +// and another optional field named "typed_value" that is either a primitive type or +// a list/large_list/list_view or struct which also satisfies the following requirements: +// i. The elements must be REQUIRED +// ii. There must either be a single REQUIRED field named "value" which is Review Comment: I think Arrow tends to use "non-nullable" instead of "required" -- 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]
