zeroshade commented on code in PR #395: URL: https://github.com/apache/arrow-go/pull/395#discussion_r2121114276
########## 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: The primary reason why the `BinaryDictionaryBuilder` doesn't meet the usual builder interface is that the standard builders *don't* return an error for their `Append` methods, the dictionary builder is unique in that way. Technically the only way you could really end up with an error coming from the dictionary builder is if something goes *really* wrong with the internal map used for the dictionary building. In practice, unless you run out of memory (leading to further problems) it's unlikely that you'd hit this at all. But in order to meet the Append interface without returning an error I didn't want to throw away the error if a user somehow hit this condition. -- 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]
