diff --git a/arrow/compute/variant_get.go b/arrow/compute/variant_get.go new file mode 100644 index 000000000..8f5a2e2af --- /dev/null +++ b/arrow/compute/variant_get.go @@ -0,0 +1,750 @@ +// 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 compute + +import ( + "context" + "fmt" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/bitutil" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/arrow/extensions" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" +) + +// VariantGetOptions controls VariantGet. +type VariantGetOptions struct { + // Path is the path to extract from every variant value. + Path variant.VariantPath + // AsType, when nil, makes VariantGet return a VariantArray pointing at the path; + // when set, the extracted values are cast to it via the cast kernels. + AsType arrow.DataType + // Strict makes a lossy cast fail; otherwise a value that cannot convert to AsType + // nulls only that row, not the rest of its natural-type group. + Strict bool +} + +// VariantGet extracts opts.Path from every value of input. It follows the shredded +// typed_value columns as far as the path allows - stepping into struct fields +// directly and gathering list elements with the take kernel - then reassembles only +// the residual for any remaining path. With AsType nil it returns a VariantArray of +// the extracted values; otherwise it casts them to AsType with the cast kernels. +func VariantGet(ctx context.Context, input *extensions.VariantArray, opts VariantGetOptions) (arrow.Array, error) { + if input == nil { + return nil, fmt.Errorf("%w: VariantGet requires a non-nil VariantArray", arrow.ErrInvalid) + } + + // Reject nested target storage. Unwrap extension types first: they embed + // ExtensionBase and so satisfy arrow.NestedType even when backed by e.g. UUID. + nestedCheck := opts.AsType + if ext, ok := nestedCheck.(arrow.ExtensionType); ok { + nestedCheck = ext.StorageType() + } + if _, ok := nestedCheck.(arrow.NestedType); ok { + return nil, fmt.Errorf("%w: VariantGet cast to nested type %s", arrow.ErrNotImplemented, opts.AsType) + } + + // Empty path, no cast: the values are returned unchanged. + if opts.Path.Len() == 0 && opts.AsType == nil { + input.Retain() + + return input, nil + } + + return shreddedGetPath(ctx, input, opts) +} + +// shreddingState is a (value?, typed_value?) column pair at one level of a shredded +// variant, mirroring arrow-rs ShreddingState. +type shreddingState struct { + value arrow.TypedArray[[]byte] + typedValue arrow.Array + length int +} + +func stateFromInput(input *extensions.VariantArray) shreddingState { + return shreddingState{ + value: input.UntypedValues(), + typedValue: input.Shredded(), + length: input.Len(), + } +} + +func stateFromFieldStruct(child *array.Struct) shreddingState { + ct := child.DataType().(*arrow.StructType) + + var value arrow.TypedArray[[]byte] + if idx, ok := ct.FieldIdx("value"); ok { + value = child.Field(idx).(arrow.TypedArray[[]byte]) + } + + var typed arrow.Array + if idx, ok := ct.FieldIdx("typed_value"); ok { + typed = child.Field(idx) + } + + return shreddingState{value: value, typedValue: typed, length: child.Len()} +} + +type pathStepKind int + +const ( + stepSuccess pathStepKind = iota + stepMissing +) + +type pathStep struct { + kind pathStepKind + state shreddingState + owned []arrow.Array // intermediate take results the caller must release +} + +// missingStep marks a path step whose typed field is absent. The descent loop breaks +// on any residual before stepping, so a step that reaches here is provably missing. +func (s shreddingState) missingStep() pathStep { + return pathStep{kind: stepMissing} +} + +// hasResidual reports whether any row carries a value in this level's value column. +func (s shreddingState) hasResidual() bool { + return s.value != nil && s.value.NullN() != s.value.Len() +} + +func fieldStep(s shreddingState, name string) (pathStep, error) { + if s.typedValue == nil { + return s.missingStep(), nil + } + st, ok := s.typedValue.(*array.Struct) + if !ok { + // A field step into a non-object shredded value is a type error, matching the + // per-row GetByPath path. Any residual was already diverted before this runs. + return pathStep{}, fmt.Errorf("%w: variant path field %q applied to non-object %s", + arrow.ErrInvalid, name, s.typedValue.DataType()) + } + idx, ok := st.DataType().(*arrow.StructType).FieldIdx(name) + if !ok { + return s.missingStep(), nil + } + child, ok := st.Field(idx).(*array.Struct) + if !ok { + return pathStep{}, fmt.Errorf("%w: expected struct field %q while following path, got %s", + arrow.ErrInvalid, name, st.Field(idx).DataType()) + } + + return pathStep{kind: stepSuccess, state: stateFromFieldStruct(child)}, nil +} + +// indexStep gathers element index from every row of a shredded list with the take +// kernel, producing the shredding state one level deeper. +func indexStep(ctx context.Context, mem memory.Allocator, s shreddingState, index int) (pathStep, error) { + if s.typedValue == nil { + return s.missingStep(), nil + } + list, ok := s.typedValue.(array.ListLike) + if !ok { + return s.missingStep(), nil + } + elems, ok := list.ListValues().(*array.Struct) + if !ok { + return s.missingStep(), nil + } + + ib := array.NewUint64Builder(mem) + defer ib.Release() + ib.Reserve(s.length) + for row := 0; row < s.length; row++ { + start, end := list.ValueOffsets(row) + if list.IsValid(row) && index >= 0 && int64(index) < end-start { + ib.Append(uint64(start + int64(index))) + } else { + ib.AppendNull() + } + } + indices := ib.NewArray() + defer indices.Release() + + et := elems.DataType().(*arrow.StructType) + var owned []arrow.Array + var next shreddingState + next.length = s.length + + if vi, ok := et.FieldIdx("value"); ok { + taken, err := TakeArray(ctx, elems.Field(vi), indices) + if err != nil { + return pathStep{}, err + } + owned = append(owned, taken) + next.value = taken.(arrow.TypedArray[[]byte]) + } + if ti, ok := et.FieldIdx("typed_value"); ok { + taken, err := TakeArray(ctx, elems.Field(ti), indices) + if err != nil { + releaseAll(owned) + + return pathStep{}, err + } + owned = append(owned, taken) + next.typedValue = taken + } + + return pathStep{kind: stepSuccess, state: next, owned: owned}, nil +} + +func releaseAll(arrs []arrow.Array) { + for _, a := range arrs { + a.Release() + } +} + +func shreddedGetPath(ctx context.Context, input *extensions.VariantArray, opts VariantGetOptions) (arrow.Array, error) { + mem := GetAllocator(ctx) + state := stateFromInput(input) + nulls := newNullTracker(input.Len(), mem) + defer nulls.release() + nulls.merge(input.Storage()) + + var owned []arrow.Array + defer func() { releaseAll(owned) }() + + idx := 0 + for idx < opts.Path.Len() { + // residual-backed rows live in the value column, unreachable by the typed_value descent; reassemble per-row. + if state.hasResidual() { + break + } + + name, index, isField := opts.Path.StepAt(idx) + var ( + step pathStep + err error + ) + if isField { + step, err = fieldStep(state, name) + } else { + step, err = indexStep(ctx, mem, state, index) + } + if err != nil { + return nil, err + } + + if step.kind == stepSuccess { + nulls.merge(state.typedValue) + state = step.state + owned = append(owned, step.owned...) + idx++ + + continue + } + + // stepMissing: the typed field is provably absent (no residual, checked above). + return allNullResult(mem, input.Len(), opts.AsType), nil + } + + remaining := subPath(opts.Path, idx) + + // Try to return the typed column directly before building the target array, + // so a perfect shredding does not allocate a struct and bitmap it discards. + if remaining.Len() == 0 && opts.AsType != nil { + if col := perfectShredded(state, nulls, opts.AsType); col != nil { + defer col.Release() + + return CastArray(ctx, col, NewCastOptions(opts.AsType, opts.Strict)) + } + } + + target, err := buildTargetVariant(input, state, nulls, mem) + if err != nil { + return nil, err + } + defer target.Release() + + if remaining.Len() == 0 && opts.AsType == nil { + target.Retain() + + return target, nil + } + + leaves, err := extractLeaves(target, remaining) + if err != nil { + return nil, err + } + if opts.AsType == nil { + return buildLeafVariantArray(mem, leaves), nil + } + + return castLeaves(ctx, mem, leaves, opts.AsType, opts.Strict) +} + +// perfectShredded returns the typed_value column when the path landed on a fully +// shredded value of exactly AsType and no ancestor nulls need merging; otherwise +// the caller's reassembly path produces the same values. +func perfectShredded(s shreddingState, nulls *nullTracker, asType arrow.DataType) arrow.Array { + if s.typedValue == nil || !nulls.allValid() { + return nil + } + if !arrow.TypeEqual(s.typedValue.DataType(), asType) { + return nil + } + if s.value != nil && s.value.NullN() != s.value.Len() { + return nil + } + + s.typedValue.Retain() + + return s.typedValue +} + +func buildTargetVariant(input *extensions.VariantArray, s shreddingState, nulls *nullTracker, mem memory.Allocator) (*extensions.VariantArray, error) { + // Read the raw metadata column rather than input.Metadata(), which asserts plain + // binary and panics on dictionary-encoded metadata; the raw column preserves + // dictionary/large-binary encoding and is decoded by the target's own reader. + storage := input.Storage().(*array.Struct) + mdIdx, ok := storage.DataType().(*arrow.StructType).FieldIdx("metadata") + if !ok { + return nil, fmt.Errorf("%w: variant storage is missing its metadata field", arrow.ErrInvalid) + } + metadata := storage.Field(mdIdx) + + fields := []arrow.Field{{Name: "metadata", Type: metadata.DataType(), Nullable: false}} + cols := []arrow.Array{metadata} + if s.value != nil { + fields = append(fields, arrow.Field{Name: "value", Type: s.value.DataType(), Nullable: true}) + cols = append(cols, s.value) + } + if s.typedValue != nil { + fields = append(fields, arrow.Field{Name: "typed_value", Type: s.typedValue.DataType(), Nullable: true}) + cols = append(cols, s.typedValue) + } + + bitmap, nullCount := nulls.validityBitmap() + st, err := array.NewStructArrayWithFieldsAndNulls(cols, fields, bitmap, nullCount, 0) + if err != nil { + return nil, err + } + defer st.Release() + + vt, err := extensions.NewVariantType(st.DataType()) + if err != nil { + return nil, err + } + + return array.NewExtensionArrayWithStorage(vt, st).(*extensions.VariantArray), nil +} + +// subPath returns the suffix of p starting at from, rebuilt through the opaque API. +func subPath(p variant.VariantPath, from int) variant.VariantPath { + var out variant.VariantPath + for i := from; i < p.Len(); i++ { + if name, index, isField := p.StepAt(i); isField { + out = out.Field(name) + } else { + out = out.Index(index) + } + } + + return out +} + +// variantLeaf is one row's extracted value; present is false when the path is +// absent for that row (or the row is null). +type variantLeaf struct { + value variant.Value + present bool +} + +func extractLeaves(target *extensions.VariantArray, path variant.VariantPath) ([]variantLeaf, error) { + leaves := make([]variantLeaf, target.Len()) + for i := range leaves { + if target.IsNull(i) { + continue + } + v, err := target.Value(i) + if err != nil { + return nil, fmt.Errorf("variant: reassembling row %d: %w", i, err) + } + leaf, found, err := v.GetByPath(path) + if err != nil { + return nil, err + } + leaves[i] = variantLeaf{value: leaf, present: found} + } + + return leaves, nil +} + +func buildLeafVariantArray(mem memory.Allocator, leaves []variantLeaf) arrow.Array { + bldr := extensions.NewVariantBuilder(mem, extensions.NewDefaultVariantType()) + defer bldr.Release() + bldr.Reserve(len(leaves)) + for _, l := range leaves { + if !l.present { + bldr.AppendNull() + + continue + } + bldr.Append(l.value) + } + + return bldr.NewArray() +} + +// castLeaves converts each leaf to asType (arrow-rs variant_get parity): leaves are +// grouped by natural type, each group cast with the cast kernels, then scattered back +// so the result is order-independent. Strict errors on a lossy cast, else null. +func castLeaves(ctx context.Context, mem memory.Allocator, leaves []variantLeaf, asType arrow.DataType, strict bool) (arrow.Array, error) { + type leafGroup struct { + dt arrow.DataType + rows []int + } + groups := make(map[string]*leafGroup) + var order []string + for i, l := range leaves { + if !l.present || l.value.Type() == variant.Null { + continue + } + dt := naturalArrowType(l.value) + if dt == nil { + // Object/array leaves have no primitive natural type. Under Strict this is an + // impossible cast (errors like any other); otherwise the row stays null. + if strict { + return nil, fmt.Errorf("%w: cannot cast non-primitive variant leaf to %s", arrow.ErrInvalid, asType) + } + + continue + } + key := dt.String() + g := groups[key] + if g == nil { + g = &leafGroup{dt: dt} + groups[key] = g + order = append(order, key) + } + g.rows = append(g.rows, i) + } + + perm := make([]uint64, len(leaves)) + valid := make([]bool, len(leaves)) + var casted []arrow.Array + defer func() { releaseAll(casted) }() + + var pos uint64 + for _, key := range order { + g := groups[key] + col := buildTypedColumn(mem, g.dt, leaves, g.rows) + cast, err := CastArray(ctx, col, NewCastOptions(asType, strict)) + col.Release() + if err == nil { + casted = append(casted, cast) + for _, row := range g.rows { + perm[row] = pos + valid[row] = true + pos++ + } + + continue + } + if strict { + return nil, err + } + // Non-strict: retry each row alone so only the inconvertible rows null, not the whole group. + for _, row := range g.rows { + rc := buildTypedColumn(mem, g.dt, leaves, []int{row}) + one, cerr := CastArray(ctx, rc, NewCastOptions(asType, strict)) + rc.Release() + if cerr != nil { + continue // this row alone is inconvertible; it stays null + } + casted = append(casted, one) + perm[row] = pos + valid[row] = true + pos++ + } + } + + if len(casted) == 0 { + return allNullResult(mem, len(leaves), asType), nil + } + + // One natural type covering every row already sits in original order. + if len(casted) == 1 && pos == uint64(len(leaves)) { + casted[0].Retain() + + return casted[0], nil + } + + combined, err := array.Concatenate(casted, mem) + if err != nil { + return nil, err + } + defer combined.Release() + + ib := array.NewUint64Builder(mem) + defer ib.Release() + ib.AppendValues(perm, valid) + indices := ib.NewArray() + defer indices.Release() + + return TakeArray(ctx, combined, indices) +} + +// buildTypedColumn materializes the given leaf rows, all of natural type dt, into a +// homogeneous Arrow array the cast kernels can consume. +func buildTypedColumn(mem memory.Allocator, dt arrow.DataType, leaves []variantLeaf, rows []int) arrow.Array { + bldr := array.NewBuilder(mem, dt) + defer bldr.Release() + bldr.Reserve(len(rows)) + for _, row := range rows { + if !appendNatural(bldr, leaves[row].value) { + bldr.AppendNull() + } + } + + return bldr.NewArray() +} + +// allNullResult builds the all-null output for a provably missing path. +func allNullResult(mem memory.Allocator, n int, asType arrow.DataType) arrow.Array { + if asType != nil { + return array.MakeArrayOfNull(mem, asType, n) + } + + // MakeArrayOfNull cannot build the variant extension type (its storage struct's + // metadata/value are non-nullable), so append encoded variant nulls instead. + bldr := extensions.NewVariantBuilder(mem, extensions.NewDefaultVariantType()) + defer bldr.Release() + for range n { + bldr.AppendNull() + } + + return bldr.NewArray() +} + +// nullTracker accumulates ancestor validity bitmaps with a bitmap AND. +type nullTracker struct { + length int + mem memory.Allocator + buf *memory.Buffer // validity bitmap (1 = valid); nil means all valid +} + +func newNullTracker(length int, mem memory.Allocator) *nullTracker { + return &nullTracker{length: length, mem: mem} +} + +// merge folds arr's validity into the accumulated mask. Arrow validity bits are +// 1=valid, so accumulating ancestor nulls is a bitmap AND (a row is null in the +// result when it is null at any level) - the validity-space equivalent of OR-ing +// null masks. +func (n *nullTracker) merge(arr arrow.Array) { + if arr == nil { + return + } + vb := arr.Data().Buffers()[0] + if vb == nil { + return // all valid + } + off := int64(arr.Data().Offset()) + if n.buf == nil { + n.buf = bitutil.BitmapAndAlloc(n.mem, vb.Bytes(), vb.Bytes(), off, off, int64(n.length), 0) + + return + } + merged := bitutil.BitmapAndAlloc(n.mem, n.buf.Bytes(), vb.Bytes(), 0, off, int64(n.length), 0) + n.buf.Release() + n.buf = merged +} + +func (n *nullTracker) allValid() bool { return n.buf == nil } + +func (n *nullTracker) validityBitmap() (*memory.Buffer, int) { + if n.buf == nil { + return nil, 0 + } + + return n.buf, n.length - bitutil.CountSetBits(n.buf.Bytes(), 0, n.length) +} + +func (n *nullTracker) release() { + if n.buf != nil { + n.buf.Release() + n.buf = nil + } +} + +func naturalArrowType(v variant.Value) arrow.DataType { + switch v.Type() { + case variant.Bool: + return arrow.FixedWidthTypes.Boolean + case variant.Int8: + return arrow.PrimitiveTypes.Int8 + case variant.Int16: + return arrow.PrimitiveTypes.Int16 + case variant.Int32: + return arrow.PrimitiveTypes.Int32 + case variant.Int64: + return arrow.PrimitiveTypes.Int64 + case variant.Float: + return arrow.PrimitiveTypes.Float32 + case variant.Double: + return arrow.PrimitiveTypes.Float64 + case variant.String: + return arrow.BinaryTypes.String + case variant.Binary: + return arrow.BinaryTypes.Binary + case variant.Date: + return arrow.FixedWidthTypes.Date32 + case variant.Time: + return arrow.FixedWidthTypes.Time64us + case variant.TimestampMicros: + return &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"} + case variant.TimestampMicrosNTZ: + return &arrow.TimestampType{Unit: arrow.Microsecond} + case variant.TimestampNanos: + return &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"} + case variant.TimestampNanosNTZ: + return &arrow.TimestampType{Unit: arrow.Nanosecond} + case variant.UUID: + return extensions.NewUUIDType() + case variant.Decimal4, variant.Decimal8, variant.Decimal16: + return &arrow.Decimal128Type{Precision: 38, Scale: int32(decimalScale(v))} + } + + return nil +} + +// appendNatural appends v to bldr when v's value matches bldr's element type, +// reporting whether it did; callers group leaves by natural type first, so it matches. +func appendNatural(bldr array.Builder, v variant.Value) bool { + switch b := bldr.(type) { + case *array.BooleanBuilder: + if x, ok := v.Value().(bool); ok { + b.Append(x) + + return true + } + case *array.Int8Builder: + if x, ok := v.Value().(int8); ok { + b.Append(x) + + return true + } + case *array.Int16Builder: + if x, ok := v.Value().(int16); ok { + b.Append(x) + + return true + } + case *array.Int32Builder: + if x, ok := v.Value().(int32); ok { + b.Append(x) + + return true + } + case *array.Int64Builder: + if x, ok := v.Value().(int64); ok { + b.Append(x) + + return true + } + case *array.Float32Builder: + if x, ok := v.Value().(float32); ok { + b.Append(x) + + return true + } + case *array.Float64Builder: + if x, ok := v.Value().(float64); ok { + b.Append(x) + + return true + } + case *array.StringBuilder: + if x, ok := v.Value().(string); ok { + b.Append(x) + + return true + } + case *array.BinaryBuilder: + if x, ok := v.Value().([]byte); ok { + b.Append(x) + + return true + } + case *array.Date32Builder: + if x, ok := v.Value().(arrow.Date32); ok { + b.Append(x) + + return true + } + case *array.Time64Builder: + if x, ok := v.Value().(arrow.Time64); ok { + b.Append(x) + + return true + } + case *array.TimestampBuilder: + if x, ok := v.Value().(arrow.Timestamp); ok { + b.Append(x) + + return true + } + case *extensions.UUIDBuilder: + if x, ok := v.Value().(uuid.UUID); ok { + b.Append(x) + + return true + } + case *array.Decimal128Builder: + if num, ok := decimalAsNum128(v); ok && int32(decimalScale(v)) == b.Type().(*arrow.Decimal128Type).Scale { + b.Append(num) + + return true + } + } + + return false +} + +func decimalScale(v variant.Value) uint8 { + switch d := v.Value().(type) { + case variant.DecimalValue[decimal.Decimal32]: + return d.Scale + case variant.DecimalValue[decimal.Decimal64]: + return d.Scale + case variant.DecimalValue[decimal.Decimal128]: + return d.Scale + } + + return 0 +} + +func decimalAsNum128(v variant.Value) (decimal128.Num, bool) { + switch d := v.Value().(type) { + case variant.DecimalValue[decimal.Decimal32]: + return decimal128.FromI64(int64(d.Value.(decimal.Decimal32))), true + case variant.DecimalValue[decimal.Decimal64]: + return decimal128.FromI64(int64(d.Value.(decimal.Decimal64))), true + case variant.DecimalValue[decimal.Decimal128]: + return d.Value.(decimal.Decimal128), true + } + + return decimal128.Num{}, false +} diff --git a/arrow/compute/variant_get_test.go b/arrow/compute/variant_get_test.go new file mode 100644 index 000000000..cf841ba28 --- /dev/null +++ b/arrow/compute/variant_get_test.go @@ -0,0 +1,875 @@ +// 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 compute_test + +import ( + "context" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/compute" + "github.com/apache/arrow-go/v18/arrow/compute/exec" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/extensions" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func vgVariant(t *testing.T, v any) variant.Value { + t.Helper() + var b variant.Builder + require.NoError(t, b.Append(v)) + val, err := b.Build() + require.NoError(t, err) + + return val +} + +func vgNonShredded(t *testing.T, mem memory.Allocator, vals ...any) *extensions.VariantArray { + t.Helper() + bldr := extensions.NewVariantBuilder(mem, extensions.NewDefaultVariantType()) + defer bldr.Release() + for _, v := range vals { + if v == nil { + bldr.AppendNull() + + continue + } + bldr.Append(vgVariant(t, v)) + } + + return bldr.NewArray().(*extensions.VariantArray) +} + +func vgShreddedInt(t *testing.T, mem memory.Allocator, vals ...int64) *extensions.VariantArray { + t.Helper() + vt := extensions.NewShreddedVariantType(arrow.PrimitiveTypes.Int64) + bldr := extensions.NewVariantBuilder(mem, vt) + defer bldr.Release() + for _, v := range vals { + bldr.Append(vgVariant(t, v)) + } + + return bldr.NewArray().(*extensions.VariantArray) +} + +func field(name string) variant.VariantPath { return variant.VariantPath{}.Field(name) } + +func TestVariantGetTyped(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, + map[string]any{"a": int64(1)}, + map[string]any{"a": int64(2)}, + nil, + ) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field("a"), AsType: arrow.PrimitiveTypes.Int64}) + require.NoError(t, err) + defer out.Release() + + ints := out.(*array.Int64) + require.Equal(t, 3, ints.Len()) + assert.EqualValues(t, 1, ints.Value(0)) + assert.EqualValues(t, 2, ints.Value(1)) + assert.True(t, ints.IsNull(2)) +} + +func TestVariantGetVariantOutput(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, map[string]any{"a": int64(7)}, map[string]any{"b": int64(9)}) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field("a")}) + require.NoError(t, err) + defer out.Release() + + varr := out.(*extensions.VariantArray) + v, err := varr.Value(0) + require.NoError(t, err) + assert.EqualValues(t, 7, v.Value()) + assert.True(t, varr.IsNull(1)) +} + +func TestVariantGetNestedAndIndex(t *testing.T) { + mem := memory.DefaultAllocator + nested := vgNonShredded(t, mem, map[string]any{"a": map[string]any{"b": int64(5)}}) + defer nested.Release() + out, err := compute.VariantGet(context.Background(), nested, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Field("a").Field("b"), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer out.Release() + assert.EqualValues(t, 5, out.(*array.Int64).Value(0)) + + arrs := vgNonShredded(t, mem, []any{int64(10), int64(20), int64(30)}) + defer arrs.Release() + got, err := compute.VariantGet(context.Background(), arrs, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Index(1), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer got.Release() + assert.EqualValues(t, 20, got.(*array.Int64).Value(0)) + + oob, err := compute.VariantGet(context.Background(), arrs, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Index(9), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer oob.Release() + assert.True(t, oob.(*array.Int64).IsNull(0)) +} + +// TestVariantGetMixedShreddedRows reproduces zeroshade's [1,2] case: row 0 is in +// typed_value, row 1 is in the residual value. Both must come back, not [1,null]. +func TestVariantGetMixedShreddedRows(t *testing.T) { + mem := memory.DefaultAllocator + s := arrow.StructOf( + arrow.Field{Name: "metadata", Type: arrow.BinaryTypes.Binary}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: true}, + arrow.Field{Name: "typed_value", Type: arrow.PrimitiveTypes.Int64, Nullable: true}) + b := array.NewStructBuilder(mem, s) + defer b.Release() + mb := b.FieldBuilder(0).(*array.BinaryBuilder) + vb := b.FieldBuilder(1).(*array.BinaryBuilder) + tb := b.FieldBuilder(2).(*array.Int64Builder) + + b.Append(true) + mb.Append(variant.EmptyMetadataBytes[:]) + vb.AppendNull() + tb.Append(1) + + b.Append(true) + mb.Append(variant.EmptyMetadataBytes[:]) + enc, err := variant.Encode(int64(2)) + require.NoError(t, err) + vb.Append(enc) + tb.AppendNull() + + st := b.NewArray() + defer st.Release() + vt, err := extensions.NewVariantType(s) + require.NoError(t, err) + arr := array.NewExtensionArrayWithStorage(vt, st).(*extensions.VariantArray) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{AsType: arrow.PrimitiveTypes.Int64}) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + require.Equal(t, 2, ints.Len()) + assert.EqualValues(t, 1, ints.Value(0)) + assert.EqualValues(t, 2, ints.Value(1), "residual-value row must be reconstructed, not null") + assert.False(t, ints.IsNull(1)) +} + +// TestVariantGetLenientCast covers zeroshade's :269 examples: ordinary widening +// casts succeed under the default (non-strict) mode. +func TestVariantGetLenientCast(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgShreddedInt(t, mem, 3, 5) + defer arr.Release() + + f64, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{AsType: arrow.PrimitiveTypes.Float64}) + require.NoError(t, err) + defer f64.Release() + assert.EqualValues(t, 3, f64.(*array.Float64).Value(0)) + assert.EqualValues(t, 5, f64.(*array.Float64).Value(1)) + + dec, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{AsType: &arrow.Decimal128Type{Precision: 10, Scale: 0}}) + require.NoError(t, err) + defer dec.Release() + assert.Equal(t, 2, dec.Len()) +} + +func TestVariantGetStrictCastErrors(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgShreddedInt(t, mem, 5_000_000_000) // overflows int8 + defer arr.Release() + + _, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + AsType: arrow.PrimitiveTypes.Int8, Strict: true, + }) + require.Error(t, err, "strict cast of an overflowing value must error") +} + +// TestVariantGetFieldOnScalarErrors covers :323: a field step into a scalar is a +// type error, not a silent null. +func TestVariantGetFieldOnScalarErrors(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, int64(1)) + defer arr.Release() + + _, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field("a")}) + require.ErrorIs(t, err, arrow.ErrInvalid) +} + +// TestVariantGetHugeIndex covers :307: a huge index must not wrap to a valid one. +func TestVariantGetHugeIndex(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, []any{int64(10), int64(20)}) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Index(1 << 40), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer out.Release() + assert.True(t, out.(*array.Int64).IsNull(0)) +} + +func TestVariantGetEmptyPath(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, int64(1), int64(2)) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{}) + require.NoError(t, err) + defer out.Release() + assert.Equal(t, 2, out.Len()) +} + +// TestVariantGetShreddedFieldPushdown drives nested field steps through the shredded +// typed_value columns and the perfect-shredding fast path. +func TestVariantGetShreddedFieldPushdown(t *testing.T) { + mem := memory.DefaultAllocator + vt := extensions.NewShreddedVariantType(arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.StructOf( + arrow.Field{Name: "b", Type: arrow.PrimitiveTypes.Int64})})) + bldr := extensions.NewVariantBuilder(mem, vt) + defer bldr.Release() + bldr.Append(vgVariant(t, map[string]any{"a": map[string]any{"b": int64(5)}})) + bldr.Append(vgVariant(t, map[string]any{"a": map[string]any{"b": int64(6)}})) + arr := bldr.NewArray().(*extensions.VariantArray) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Field("a").Field("b"), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + assert.EqualValues(t, 5, ints.Value(0)) + assert.EqualValues(t, 6, ints.Value(1)) +} + +// TestVariantGetShreddedListIndex drives an index step over a shredded list, which +// gathers elements with the take kernel. +func TestVariantGetShreddedListIndex(t *testing.T) { + mem := memory.DefaultAllocator + vt := extensions.NewShreddedVariantType(arrow.ListOf(arrow.PrimitiveTypes.Int64)) + bldr := extensions.NewVariantBuilder(mem, vt) + defer bldr.Release() + bldr.Append(vgVariant(t, []any{int64(10), int64(20), int64(30)})) + bldr.Append(vgVariant(t, []any{int64(40), int64(50)})) + arr := bldr.NewArray().(*extensions.VariantArray) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Index(1), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + assert.EqualValues(t, 20, ints.Value(0)) + assert.EqualValues(t, 50, ints.Value(1)) +} + +func TestVariantGetNoLeak(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + ctx := exec.WithAllocator(context.Background(), mem) + + vt := extensions.NewShreddedVariantType(arrow.ListOf(arrow.PrimitiveTypes.Int64)) + bldr := extensions.NewVariantBuilder(mem, vt) + bldr.Append(vgVariant(t, []any{int64(10), int64(20)})) + bldr.AppendNull() + arr := bldr.NewArray().(*extensions.VariantArray) + bldr.Release() + + idx, err := compute.VariantGet(ctx, arr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Index(0), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + idx.Release() + arr.Release() + + // A missing field on an object-shredded variant exercises the all-null path. + objVT := extensions.NewShreddedVariantType(arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64})) + ob := extensions.NewVariantBuilder(mem, objVT) + ob.Append(vgVariant(t, map[string]any{"a": int64(1)})) + objArr := ob.NewArray().(*extensions.VariantArray) + ob.Release() + + missing, err := compute.VariantGet(ctx, objArr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Field("nope"), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + missing.Release() + objArr.Release() +} + +// TestVariantGetDictMetadata guards against a panic when the metadata column is +// dictionary-encoded (spec-legal): buildTargetVariant must read the raw column, +// not the plain-binary accessor. +func TestVariantGetDictMetadata(t *testing.T) { + mem := memory.DefaultAllocator + s := arrow.StructOf( + arrow.Field{Name: "metadata", Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Uint8, ValueType: arrow.BinaryTypes.Binary}}, + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: true}, + arrow.Field{Name: "typed_value", Type: arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.StructOf( + arrow.Field{Name: "value", Type: arrow.BinaryTypes.Binary, Nullable: true}, + arrow.Field{Name: "typed_value", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + )}, + ), Nullable: true}) + vt, err := extensions.NewVariantType(s) + require.NoError(t, err) + bldr := vt.NewBuilder(mem).(*extensions.VariantBuilder) + defer bldr.Release() + bldr.Append(vgVariant(t, map[string]any{"a": int64(5), "b": "resid"})) + arr := bldr.NewArray().(*extensions.VariantArray) + defer arr.Release() + + // "b" is not shredded, so this takes the NotShredded -> buildTargetVariant path. + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: variant.VariantPath{}.Field("b")}) + require.NoError(t, err) + defer out.Release() + v, err := out.(*extensions.VariantArray).Value(0) + require.NoError(t, err) + assert.Equal(t, "resid", v.Value()) +} + +// vgMixedResidual builds a two-row shredded array: row 0 shredded (fill populates typed_value, top value null), +// row 1 residual (top typed_value null, top value = resid). Storage comes from NewShreddedVariantType. +func vgMixedResidual(t *testing.T, mem memory.Allocator, shredType arrow.DataType, fill func(b array.Builder), resid variant.Value) *extensions.VariantArray { + t.Helper() + vt := extensions.NewShreddedVariantType(shredType) + s := vt.StorageType().(*arrow.StructType) + mIdx, _ := s.FieldIdx("metadata") + vIdx, _ := s.FieldIdx("value") + tIdx, _ := s.FieldIdx("typed_value") + + b := array.NewStructBuilder(mem, s) + defer b.Release() + mb := b.FieldBuilder(mIdx).(*array.BinaryBuilder) + vb := b.FieldBuilder(vIdx).(*array.BinaryBuilder) + + b.Append(true) + mb.Append(variant.EmptyMetadataBytes[:]) + vb.AppendNull() + fill(b.FieldBuilder(tIdx)) + + b.Append(true) + mb.Append(resid.Metadata().Bytes()) + vb.Append(resid.Bytes()) + b.FieldBuilder(tIdx).AppendNull() + + st := b.NewArray() + defer st.Release() + + return array.NewExtensionArrayWithStorage(vt, st).(*extensions.VariantArray) +} + +// TestVariantGetResidualBackedField covers zeroshade's blocking :233 case for a root +// field path: row 1's {"a":2} lives in the top residual, so $.a must return 2 not null. +func TestVariantGetResidualBackedField(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgMixedResidual(t, mem, arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}), + func(b array.Builder) { + tv := b.(*array.StructBuilder) + tv.Append(true) + a := tv.FieldBuilder(0).(*array.StructBuilder) + a.Append(true) + a.FieldBuilder(0).(*array.BinaryBuilder).AppendNull() + a.FieldBuilder(1).(*array.Int64Builder).Append(1) + }, vgVariant(t, map[string]any{"a": int64(2)})) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field("a"), AsType: arrow.PrimitiveTypes.Int64}) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + require.Equal(t, 2, ints.Len()) + assert.EqualValues(t, 1, ints.Value(0)) + assert.EqualValues(t, 2, ints.Value(1), "residual-backed row must be reassembled, not nulled") +} + +// TestVariantGetResidualBackedNestedField covers the nested-field case: $.a.b where +// row 1's whole {"a":{"b":6}} lives in the top residual. +func TestVariantGetResidualBackedNestedField(t *testing.T) { + mem := memory.DefaultAllocator + shred := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.StructOf(arrow.Field{Name: "b", Type: arrow.PrimitiveTypes.Int64})}) + arr := vgMixedResidual(t, mem, shred, + func(b array.Builder) { + tv := b.(*array.StructBuilder) + tv.Append(true) + a := tv.FieldBuilder(0).(*array.StructBuilder) + a.Append(true) + a.FieldBuilder(0).(*array.BinaryBuilder).AppendNull() + aTV := a.FieldBuilder(1).(*array.StructBuilder) + aTV.Append(true) + bf := aTV.FieldBuilder(0).(*array.StructBuilder) + bf.Append(true) + bf.FieldBuilder(0).(*array.BinaryBuilder).AppendNull() + bf.FieldBuilder(1).(*array.Int64Builder).Append(5) + }, vgVariant(t, map[string]any{"a": map[string]any{"b": int64(6)}})) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Field("a").Field("b"), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + assert.EqualValues(t, 5, ints.Value(0)) + assert.EqualValues(t, 6, ints.Value(1), "residual-backed row must be reassembled, not nulled") +} + +// TestVariantGetResidualBackedListIndex covers the list-index case: [0] where row 1's +// whole [30,40] lives in the top residual. +func TestVariantGetResidualBackedListIndex(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgMixedResidual(t, mem, arrow.ListOf(arrow.PrimitiveTypes.Int64), + func(b array.Builder) { + lb := b.(*array.ListBuilder) + lb.Append(true) + el := lb.ValueBuilder().(*array.StructBuilder) + for _, v := range []int64{10, 20} { + el.Append(true) + el.FieldBuilder(0).(*array.BinaryBuilder).AppendNull() + el.FieldBuilder(1).(*array.Int64Builder).Append(v) + } + }, vgVariant(t, []any{int64(30), int64(40)})) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Index(0), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + assert.EqualValues(t, 10, ints.Value(0)) + assert.EqualValues(t, 30, ints.Value(1), "residual-backed row must be reassembled, not nulled") +} + +// TestVariantGetMixedWidthIntegers covers zeroshade :411: variant ints encode at their +// natural width, so a wider AsType must not drop rows whose leaf shredded narrower. +func TestVariantGetMixedWidthIntegers(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, + map[string]any{"a": int64(1)}, // int8 + map[string]any{"a": int64(1000)}, // int16 + map[string]any{"a": int64(5_000_000_000)}, // int64 + map[string]any{"a": int64(9007199254740993)}) // int64 > 2^53, must stay exact (no float intermediate) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field("a"), AsType: arrow.PrimitiveTypes.Int64}) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + require.Equal(t, 4, ints.Len()) + assert.EqualValues(t, 1, ints.Value(0)) + assert.EqualValues(t, 1000, ints.Value(1), "narrower-width leaf must not be dropped") + assert.EqualValues(t, 5_000_000_000, ints.Value(2)) + assert.EqualValues(t, 9007199254740993, ints.Value(3), "value > 2^53 must be exact, not routed through float64") +} + +// TestVariantGetHeterogeneousLeaves (arrow-rs parity): a valid int64 survives a +// narrower-typed sibling; only a non-numeric string nulls. +func TestVariantGetHeterogeneousLeaves(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, + map[string]any{"a": int64(1)}, // int8 encoding + map[string]any{"a": int64(5_000_000_000)}, // int64, does not fit int8 + map[string]any{"a": "x"}) // non-numeric -> null + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field("a"), AsType: arrow.PrimitiveTypes.Int64}) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + assert.EqualValues(t, 1, ints.Value(0)) + assert.EqualValues(t, 5_000_000_000, ints.Value(1), "valid int64 must survive a narrower-typed sibling") + assert.True(t, ints.IsNull(2), "non-numeric string must be null") +} + +// TestVariantGetEmptyKey covers zeroshade's blocking path.go:42 case: an empty-string +// object key is a field step, not array index 0. +func TestVariantGetEmptyKey(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, map[string]any{"": int64(42)}) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field(""), AsType: arrow.PrimitiveTypes.Int64}) + require.NoError(t, err) + defer out.Release() + assert.EqualValues(t, 42, out.(*array.Int64).Value(0)) +} + +// TestVariantGetResidualBackedDeepField exercises the residual break after a successful +// columnar descent: row 0 is shredded through a.b, row 1's {"b":6} sits in a's residual. +func TestVariantGetResidualBackedDeepField(t *testing.T) { + mem := memory.DefaultAllocator + shred := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.StructOf(arrow.Field{Name: "b", Type: arrow.PrimitiveTypes.Int64})}) + vt := extensions.NewShreddedVariantType(shred) + s := vt.StorageType().(*arrow.StructType) + mIdx, _ := s.FieldIdx("metadata") + vIdx, _ := s.FieldIdx("value") + tIdx, _ := s.FieldIdx("typed_value") + + b := array.NewStructBuilder(mem, s) + defer b.Release() + mb := b.FieldBuilder(mIdx).(*array.BinaryBuilder) + vb := b.FieldBuilder(vIdx).(*array.BinaryBuilder) + tvb := b.FieldBuilder(tIdx).(*array.StructBuilder) // struct{a} + aField := tvb.FieldBuilder(0).(*array.StructBuilder) + aVal := aField.FieldBuilder(0).(*array.BinaryBuilder) + aTyped := aField.FieldBuilder(1).(*array.StructBuilder) // struct{b} + bField := aTyped.FieldBuilder(0).(*array.StructBuilder) + bVal := bField.FieldBuilder(0).(*array.BinaryBuilder) + bTyped := bField.FieldBuilder(1).(*array.Int64Builder) + + // row 0: fully shredded a.b = 5 + b.Append(true) + mb.Append(variant.EmptyMetadataBytes[:]) + vb.AppendNull() + tvb.Append(true) + aField.Append(true) + aVal.AppendNull() + aTyped.Append(true) + bField.Append(true) + bVal.AppendNull() + bTyped.Append(5) + + // row 1: a is residual-backed with {"b":6} (top typed_value present, a.value set) + resid := vgVariant(t, map[string]any{"b": int64(6)}) + b.Append(true) + mb.Append(resid.Metadata().Bytes()) + vb.AppendNull() + tvb.Append(true) + aField.Append(true) + aVal.Append(resid.Bytes()) + aTyped.AppendNull() // a.typed_value null -> recurses bField/bVal/bTyped null + + st := b.NewArray() + defer st.Release() + arr := array.NewExtensionArrayWithStorage(vt, st).(*extensions.VariantArray) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Field("a").Field("b"), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + defer out.Release() + ints := out.(*array.Int64) + require.Equal(t, 2, ints.Len()) + assert.EqualValues(t, 5, ints.Value(0)) + assert.EqualValues(t, 6, ints.Value(1), "mid-level residual row must be reassembled, not nulled") +} + +// TestVariantGetResidualNoLeak guards the residual break path (buildTargetVariant + +// per-row reassembly) against leaks, which TestVariantGetNoLeak does not reach. +func TestVariantGetResidualNoLeak(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + ctx := exec.WithAllocator(context.Background(), mem) + + arr := vgMixedResidual(t, mem, arrow.ListOf(arrow.PrimitiveTypes.Int64), + func(bld array.Builder) { + lb := bld.(*array.ListBuilder) + lb.Append(true) + el := lb.ValueBuilder().(*array.StructBuilder) + for _, v := range []int64{10, 20} { + el.Append(true) + el.FieldBuilder(0).(*array.BinaryBuilder).AppendNull() + el.FieldBuilder(1).(*array.Int64Builder).Append(v) + } + }, vgVariant(t, []any{int64(30), int64(40)})) + + out, err := compute.VariantGet(ctx, arr, compute.VariantGetOptions{ + Path: variant.VariantPath{}.Index(0), AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + out.Release() + arr.Release() +} + +// vgNonShreddedVals builds a non-shredded array from pre-built values, so a test can +// control per-value encoding (e.g. timestamp unit) that vgNonShredded cannot. +func vgNonShreddedVals(t *testing.T, mem memory.Allocator, vals ...variant.Value) *extensions.VariantArray { + t.Helper() + bldr := extensions.NewVariantBuilder(mem, extensions.NewDefaultVariantType()) + defer bldr.Release() + for _, v := range vals { + bldr.Append(v) + } + + return bldr.NewArray().(*extensions.VariantArray) +} + +func vgTimestamp(t *testing.T, ts arrow.Timestamp, nano bool) variant.Value { + t.Helper() + var b variant.Builder + opts := []variant.AppendOpt{variant.OptTimestampUTC} + if nano { + opts = append(opts, variant.OptTimestampNano) + } + require.NoError(t, b.Append(ts, opts...)) + val, err := b.Build() + require.NoError(t, err) + + return val +} + +// TestVariantGetMixedFloatWidths (:411, floats): a Float(32) and a Double(64) leaf +// both survive a Float64 request, independent of order. +func TestVariantGetMixedFloatWidths(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, float32(1.5), float64(2.5)) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{AsType: arrow.PrimitiveTypes.Float64}) + require.NoError(t, err) + defer out.Release() + f := out.(*array.Float64) + require.Equal(t, 2, f.Len()) + assert.InDelta(t, 1.5, f.Value(0), 1e-9) + assert.InDelta(t, 2.5, f.Value(1), 1e-9, "Double leaf must not be dropped by a Float first leaf") +} + +// TestVariantGetIntPlusFloat: a mixed int/float column cast to Float64 widens the int +// through the cast kernels rather than nulling it. +func TestVariantGetIntPlusFloat(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, int64(3), float64(2.5)) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{AsType: arrow.PrimitiveTypes.Float64}) + require.NoError(t, err) + defer out.Release() + f := out.(*array.Float64) + assert.InDelta(t, 3.0, f.Value(0), 1e-9) + assert.InDelta(t, 2.5, f.Value(1), 1e-9) +} + +// TestVariantGetTypeOrderIndependent (:411): the same two values give the same result +// regardless of which row comes first. +func TestVariantGetTypeOrderIndependent(t *testing.T) { + mem := memory.DefaultAllocator + forward := vgNonShredded(t, mem, int64(3), float64(2.5)) + defer forward.Release() + reverse := vgNonShredded(t, mem, float64(2.5), int64(3)) + defer reverse.Release() + + optsF := compute.VariantGetOptions{AsType: arrow.PrimitiveTypes.Float64} + fwd, err := compute.VariantGet(context.Background(), forward, optsF) + require.NoError(t, err) + defer fwd.Release() + rev, err := compute.VariantGet(context.Background(), reverse, optsF) + require.NoError(t, err) + defer rev.Release() + + fa, ra := fwd.(*array.Float64), rev.(*array.Float64) + assert.InDelta(t, fa.Value(0), ra.Value(1), 1e-9) + assert.InDelta(t, fa.Value(1), ra.Value(0), 1e-9) + assert.False(t, fa.IsNull(0) || fa.IsNull(1) || ra.IsNull(0) || ra.IsNull(1), "no leaf dropped in either order") +} + +// TestVariantGetMixedTimestampUnits: a micros leaf and a nanos leaf of the same instant +// both land on it when cast to a nanos target (units are converted, not reinterpreted). +func TestVariantGetMixedTimestampUnits(t *testing.T) { + mem := memory.DefaultAllocator + const micros = arrow.Timestamp(1_600_000_000_000_000) + const nanos = arrow.Timestamp(1_600_000_000_000_000_000) + arr := vgNonShreddedVals(t, mem, vgTimestamp(t, micros, false), vgTimestamp(t, nanos, true)) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + AsType: &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}, + }) + require.NoError(t, err) + defer out.Release() + ts := out.(*array.Timestamp) + require.Equal(t, 2, ts.Len()) + assert.EqualValues(t, nanos, ts.Value(0), "micros leaf must be scaled to nanos, not copied raw") + assert.EqualValues(t, nanos, ts.Value(1)) +} + +// TestVariantGetMixedDecimalScales: leaves shredded at different scales both rescale to +// the requested target scale. +func TestVariantGetMixedDecimalScales(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShreddedVals(t, mem, + vgVariant(t, variant.DecimalValue[decimal.Decimal32]{Scale: 1, Value: decimal.Decimal32(15)}), // 1.5 + vgVariant(t, variant.DecimalValue[decimal.Decimal32]{Scale: 2, Value: decimal.Decimal32(225)})) // 2.25 + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + AsType: &arrow.Decimal128Type{Precision: 38, Scale: 2}, + }) + require.NoError(t, err) + defer out.Release() + d := out.(*array.Decimal128) + require.Equal(t, 2, d.Len()) + assert.InDelta(t, 1.5, d.Value(0).ToFloat64(2), 1e-9, "scale-1 leaf must rescale to scale-2, not drop") + assert.InDelta(t, 2.25, d.Value(1).ToFloat64(2), 1e-9) +} + +// TestVariantGetStrictSlowPathErrors (:411/:269): on the reassembly path a lossy cast +// errors under Strict rather than nulling. +func TestVariantGetStrictSlowPathErrors(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, map[string]any{"a": int64(5_000_000_000)}) // overflows int32 + defer arr.Release() + + _, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: field("a"), AsType: arrow.PrimitiveTypes.Int32, Strict: true, + }) + require.Error(t, err, "Strict must error on an overflowing cast, not null it") +} + +// TestVariantGetMixedTypeNoLeak guards the multi-group scatter path (cast + Concatenate +// + Take), which the single-type leak tests do not reach. +// TestVariantGetNestedTypeNotImplemented pins that a nested AsType is rejected with +// ErrNotImplemented rather than silently producing an all-null array. +func TestVariantGetNestedTypeNotImplemented(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, map[string]any{"a": map[string]any{"x": int64(1)}}) + defer arr.Release() + + for _, nested := range []arrow.DataType{ + arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int64}), + arrow.ListOf(arrow.PrimitiveTypes.Int64), + } { + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field("a"), AsType: nested}) + if out != nil { + out.Release() + } + require.ErrorIs(t, err, arrow.ErrNotImplemented, "nested AsType %s must error, not null", nested) + } +} + +func TestVariantGetMixedTypeNoLeak(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + ctx := exec.WithAllocator(context.Background(), mem) + + arr := vgNonShredded(t, mem, int64(3), float64(2.5), nil, "x") + out, err := compute.VariantGet(ctx, arr, compute.VariantGetOptions{AsType: arrow.PrimitiveTypes.Float64}) + require.NoError(t, err) + out.Release() + arr.Release() +} + +// TestVariantGetInterleavedScatter exercises the multi-group scatter (Concatenate + +// TakeArray) with a NON-IDENTITY permutation: two same-typed leaves straddle a +// differently-typed one, so group order [3,7,2.5] must scatter back to row order +// [3,2.5,7]. Every other multi-group test lands in identity order. +func TestVariantGetInterleavedScatter(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgNonShredded(t, mem, int64(3), float64(2.5), int64(7)) // Int8{0,2}, Float64{1} + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{AsType: arrow.PrimitiveTypes.Float64}) + require.NoError(t, err) + defer out.Release() + f := out.(*array.Float64) + require.Equal(t, 3, f.Len()) + assert.InDelta(t, 3.0, f.Value(0), 1e-9) + assert.InDelta(t, 2.5, f.Value(1), 1e-9, "interleaved leaf must scatter back to its row, not stay in group order") + assert.InDelta(t, 7.0, f.Value(2), 1e-9) +} + +// TestVariantGetStrictObjectLeafErrors pins that under Strict an object/array leaf cast +// to a primitive errors (impossible cast) rather than silently nulling. +func TestVariantGetStrictObjectLeafErrors(t *testing.T) { + mem := memory.DefaultAllocator + for _, v := range []any{ + map[string]any{"a": map[string]any{"x": int64(1)}}, // object leaf at $.a + map[string]any{"a": []any{int64(1), int64(2)}}, // array leaf at $.a + } { + arr := vgNonShredded(t, mem, v) + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: field("a"), AsType: arrow.PrimitiveTypes.Int64, Strict: true, + }) + if out != nil { + out.Release() + } + arr.Release() + require.ErrorIs(t, err, arrow.ErrInvalid, "strict cast of a non-primitive leaf to Int64 must error") + } +} + +// TestVariantGetShreddedFieldOnScalarErrors pins that a field step into a shredded +// scalar column errors on the columnar path, matching the per-row GetByPath path. +func TestVariantGetShreddedFieldOnScalarErrors(t *testing.T) { + mem := memory.DefaultAllocator + arr := vgShreddedInt(t, mem, 1, 2, 3) // typed_value is a scalar Int64 column + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{Path: field("a"), AsType: arrow.PrimitiveTypes.Int64}) + if out != nil { + out.Release() + } + require.ErrorIs(t, err, arrow.ErrInvalid, "field access on a shredded scalar must error, not return null") +} + +// TestVariantGetUUIDTarget pins that a UUID AsType reaches the UUID cast, not the nested-type reject. +func TestVariantGetUUIDTarget(t *testing.T) { + mem := memory.DefaultAllocator + u := uuid.MustParse("00112233-4455-6677-8899-aabbccddeeff") + arr := vgNonShredded(t, mem, map[string]any{"id": u}) + defer arr.Release() + + out, err := compute.VariantGet(context.Background(), arr, compute.VariantGetOptions{ + Path: field("id"), AsType: extensions.NewUUIDType(), + }) + require.NoError(t, err, "UUID target must not be rejected as a nested type") + defer out.Release() + + uarr, ok := out.(*extensions.UUIDArray) + require.True(t, ok, "expected *extensions.UUIDArray, got %T", out) + require.Equal(t, 1, uarr.Len()) + require.False(t, uarr.IsNull(0)) + assert.Equal(t, u, uarr.Value(0)) +} + +// TestVariantGetPartialCastFailure pins that non-strict nulls only the inconvertible row: ["1","bad"]->int64 is [1,null]. +func TestVariantGetPartialCastFailure(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + ctx := exec.WithAllocator(context.Background(), mem) + arr := vgNonShredded(t, mem, "1", "bad") + + out, err := compute.VariantGet(ctx, arr, compute.VariantGetOptions{ + AsType: arrow.PrimitiveTypes.Int64, + }) + require.NoError(t, err) + + ints := out.(*array.Int64) + require.Equal(t, 2, ints.Len()) + assert.False(t, ints.IsNull(0), "valid row must survive a sibling row's cast failure") + assert.EqualValues(t, 1, ints.Value(0)) + assert.True(t, ints.IsNull(1), "only the inconvertible row is null") + + out.Release() + arr.Release() +} diff --git a/arrow/extensions/variant.go b/arrow/extensions/variant.go index fee2e046a..a0cc74edc 100644 --- a/arrow/extensions/variant.go +++ b/arrow/extensions/variant.go @@ -458,6 +458,11 @@ func (v *VariantArray) IsShredded() bool { return v.ExtensionType().(*VariantType).typedValueFieldIdx != -1 } +// VariantType returns the array's extension type without the ExtensionType cast. +func (v *VariantArray) VariantType() *VariantType { + return v.ExtensionType().(*VariantType) +} + // UnshredVariant returns an equivalent VariantArray in the non-shredded layout // (a struct of metadata and value), reassembling each row's value from the // shredded typed_value and value columns. If the array is already non-shredded @@ -513,6 +518,11 @@ func (v *VariantArray) IsNull(i int) bool { } } + if vt.valueFieldIdx == -1 { + // No residual value column: a null typed_value means the value is missing. + return true + } + valArr := v.Storage().(*array.Struct).Field(vt.valueFieldIdx) b := valArr.(arrow.TypedArray[[]byte]).Value(i) return len(b) == 1 && b[0] == 0 // variant null diff --git a/parquet/variant/path.go b/parquet/variant/path.go new file mode 100644 index 000000000..bc46da77d --- /dev/null +++ b/parquet/variant/path.go @@ -0,0 +1,109 @@ +// 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 variant + +import ( + "errors" + "fmt" + + "github.com/apache/arrow-go/v18/arrow" +) + +// pathElem is one step of a VariantPath: an object field when isField is set, else an array index. +type pathElem struct { + name string + index int + isField bool +} + +// VariantPath is an ordered list of steps to navigate into a variant value. The +// zero value is the root path; extend it with Field and Index. +type VariantPath struct { + elems []pathElem +} + +// Field returns a copy of the path with an object-field step appended. +func (p VariantPath) Field(name string) VariantPath { + return VariantPath{elems: append(p.grow(), pathElem{name: name, isField: true})} +} + +// Index returns a copy of the path with an array-index step appended. +func (p VariantPath) Index(i int) VariantPath { + return VariantPath{elems: append(p.grow(), pathElem{index: i})} +} + +// Join returns a copy of the path with other's steps appended. +func (p VariantPath) Join(other VariantPath) VariantPath { + return VariantPath{elems: append(p.grow(), other.elems...)} +} + +func (p VariantPath) grow() []pathElem { + return append(make([]pathElem, 0, len(p.elems)+1), p.elems...) +} + +// Len returns the number of steps in the path. +func (p VariantPath) Len() int { return len(p.elems) } + +// StepAt returns the i-th step's name and index, with isField true for an object-field step. +func (p VariantPath) StepAt(i int) (name string, index int, isField bool) { + e := p.elems[i] + + return e.name, e.index, e.isField +} + +// GetByPath navigates path into v and returns the leaf value. found is false when +// the path is cleanly absent (a missing object field, or an out-of-range or +// non-array index). It returns an error for a type error (a field step into a +// non-object) or corrupt data (a field id not present in the metadata). +func (v Value) GetByPath(path VariantPath) (leaf Value, found bool, err error) { + cur := v + for _, e := range path.elems { + if e.isField { + obj, ok := cur.Value().(ObjectValue) + if !ok { + return Value{}, false, fmt.Errorf("%w: variant path field %q applied to non-object", arrow.ErrInvalid, e.name) + } + field, ferr := obj.ValueByKey(e.name) + if ferr != nil { + if errors.Is(ferr, arrow.ErrNotFound) { + return Value{}, false, nil + } + + return Value{}, false, ferr + } + cur = field.Value + + continue + } + + arr, ok := cur.Value().(ArrayValue) + if !ok { + return Value{}, false, nil + } + if e.index < 0 || uint64(e.index) >= uint64(arr.Len()) { + return Value{}, false, nil + } + el, aerr := arr.Value(uint32(e.index)) + if aerr != nil { + return Value{}, false, nil + } + cur = el + } + + return cur, true, nil +} diff --git a/parquet/variant/path_test.go b/parquet/variant/path_test.go new file mode 100644 index 000000000..9de8b84ad --- /dev/null +++ b/parquet/variant/path_test.go @@ -0,0 +1,116 @@ +// 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 variant_test + +import ( + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func buildVar(t *testing.T, v any) variant.Value { + t.Helper() + var b variant.Builder + require.NoError(t, b.Append(v)) + val, err := b.Build() + require.NoError(t, err) + + return val +} + +func TestGetByPathFieldAndIndex(t *testing.T) { + v := buildVar(t, map[string]any{"a": map[string]any{"b": int64(5)}, "arr": []any{int64(10), int64(20)}}) + + leaf, found, err := v.GetByPath(variant.VariantPath{}.Field("a").Field("b")) + require.NoError(t, err) + require.True(t, found) + assert.EqualValues(t, 5, leaf.Value()) + + leaf, found, err = v.GetByPath(variant.VariantPath{}.Field("arr").Index(1)) + require.NoError(t, err) + require.True(t, found) + assert.EqualValues(t, 20, leaf.Value()) +} + +func TestGetByPathAbsent(t *testing.T) { + v := buildVar(t, map[string]any{"a": int64(1), "arr": []any{int64(10)}}) + + for _, p := range []variant.VariantPath{ + variant.VariantPath{}.Field("missing"), // absent object field + variant.VariantPath{}.Field("arr").Index(9), // out-of-range index + variant.VariantPath{}.Field("a").Index(0), // index into a scalar + variant.VariantPath{}.Index(0), // index into an object + } { + _, found, err := v.GetByPath(p) + require.NoError(t, err) + assert.False(t, found) + } +} + +// TestGetByPathFieldOnScalarErrors: a field step into a non-object is a type error. +func TestGetByPathFieldOnScalarErrors(t *testing.T) { + v := buildVar(t, int64(1)) + _, _, err := v.GetByPath(variant.VariantPath{}.Field("a")) + require.ErrorIs(t, err, arrow.ErrInvalid) +} + +// TestGetByPathHugeIndex: a huge index must not wrap; it is simply absent. +func TestGetByPathHugeIndex(t *testing.T) { + v := buildVar(t, []any{int64(10), int64(20)}) + _, found, err := v.GetByPath(variant.VariantPath{}.Index(1 << 40)) + require.NoError(t, err) + assert.False(t, found) +} + +func TestVariantPathJoinAndStepAt(t *testing.T) { + p := variant.VariantPath{}.Field("a").Join(variant.VariantPath{}.Index(2).Field("b")) + require.Equal(t, 3, p.Len()) + + name, _, isField := p.StepAt(0) + assert.Equal(t, "a", name) + assert.True(t, isField) + name, idx, isField := p.StepAt(1) + assert.Equal(t, "", name) + assert.Equal(t, 2, idx) + assert.False(t, isField) + name, _, isField = p.StepAt(2) + assert.Equal(t, "b", name) + assert.True(t, isField) +} + +// TestGetByPathEmptyKey covers the empty-string object key: Field("") is a field +// step distinct from Index(0), so it must resolve the "" key rather than index 0. +func TestGetByPathEmptyKey(t *testing.T) { + v := buildVar(t, map[string]any{"": int64(42)}) + + _, _, isField := variant.VariantPath{}.Field("").StepAt(0) + assert.True(t, isField, `Field("") must be a field step, not an index step`) + + leaf, found, err := v.GetByPath(variant.VariantPath{}.Field("")) + require.NoError(t, err) + require.True(t, found) + assert.EqualValues(t, 42, leaf.Value()) + + // Index(0) on the object must not match the "" key. + _, found, err = v.GetByPath(variant.VariantPath{}.Index(0)) + require.NoError(t, err) + assert.False(t, found) +}