zeroshade commented on code in PR #1206: URL: https://github.com/apache/arrow-go/pull/1206#discussion_r3855092605
########## parquet/variant/path.go: ########## @@ -0,0 +1,108 @@ +// 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 name != "", else an +// array index. +type pathElem struct { + name string + index int +} + +// 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})} Review Comment: **Blocking:** `Field("")` produces the same zero-valued `pathElem` as `Index(0)`, and `StepAt`/`GetByPath` distinguish steps using `name != ""`. Empty-string object keys are valid, so extracting `{"":42}` with `Field("")` is interpreted as array index 0 and returns null. Please retain an explicit field/index discriminator through `StepAt` and path rebuilding, with an empty-key regression test. ########## arrow/compute/variant_get.go: ########## @@ -0,0 +1,658 @@ +// 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; the default allows overflow and truncation via + // the cast kernels. Unlike arrow-rs safe mode there is no null-on-failure: an + // impossible cast always errors, since arrow-go's cast kernels have no safe flag. + 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) + } + + // 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 + stepNotShredded +) + +type pathStep struct { + kind pathStepKind + state shreddingState + owned []arrow.Array // intermediate take results the caller must release +} + +// missingStep reports whether an absent typed field is provably missing (the value +// column is all-null) or merely not shredded (a residual may hold it). +func (s shreddingState) missingStep() pathStep { + if s.value == nil || s.value.NullN() == s.value.Len() { + return pathStep{kind: stepMissing} + } + + return pathStep{kind: stepNotShredded} +} + +func fieldStep(s shreddingState, name string) (pathStep, error) { + if s.typedValue == nil { + return s.missingStep(), nil + } + st, ok := s.typedValue.(*array.Struct) + if !ok { + return s.missingStep(), nil + } + 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() { + name, index := opts.Path.StepAt(idx) + var ( + step pathStep + err error + ) + if name != "" { + 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 + } + if step.kind == stepMissing { + return allNullResult(mem, input.Len(), opts.AsType), nil + } + + break // stepNotShredded + } + + 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 + } + + src := buildNaturalArray(mem, leaves) + if src == nil { + return allNullResult(mem, len(leaves), opts.AsType), nil + } + defer src.Release() + + return CastArray(ctx, src, NewCastOptions(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 _, ok := asType.(arrow.NestedType); ok { + return nil + } + 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 := p.StepAt(i); name != "" { + 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() +} + +// buildNaturalArray materializes the leaves as an array of the first present leaf's +// natural Arrow type so the cast kernels can convert it. Rows whose value does not +// match that natural type become null. Returns nil when no leaf is present. +func buildNaturalArray(mem memory.Allocator, leaves []variantLeaf) arrow.Array { + var natural arrow.DataType + for _, l := range leaves { + if l.present && l.value.Type() != variant.Null { + natural = naturalArrowType(l.value) Review Comment: **Blocking:** Choosing the source type from the first present leaf silently discards later values of another natural width/type at lines 424–425. Variant integers commonly encode as different widths: values `1`, `1000`, and `5000000000` become `int8`, `int16`, and `int64`; requesting `AsType: int64` currently returns `[1,null,null]`. Results also become row-order-dependent and `Strict` behaves differently depending on the fast path. Please materialize against the requested target type or otherwise unify leaf types, with mixed-width and mixed-type regression coverage. ########## arrow/compute/variant_get.go: ########## @@ -0,0 +1,658 @@ +// 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; the default allows overflow and truncation via + // the cast kernels. Unlike arrow-rs safe mode there is no null-on-failure: an + // impossible cast always errors, since arrow-go's cast kernels have no safe flag. + 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) + } + + // 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 + stepNotShredded +) + +type pathStep struct { + kind pathStepKind + state shreddingState + owned []arrow.Array // intermediate take results the caller must release +} + +// missingStep reports whether an absent typed field is provably missing (the value +// column is all-null) or merely not shredded (a residual may hold it). +func (s shreddingState) missingStep() pathStep { + if s.value == nil || s.value.NullN() == s.value.Len() { + return pathStep{kind: stepMissing} + } + + return pathStep{kind: stepNotShredded} +} + +func fieldStep(s shreddingState, name string) (pathStep, error) { + if s.typedValue == nil { + return s.missingStep(), nil + } + st, ok := s.typedValue.(*array.Struct) + if !ok { + return s.missingStep(), nil + } + 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() { + name, index := opts.Path.StepAt(idx) + var ( + step pathStep + err error + ) + if name != "" { + 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) Review Comment: **Blocking:** This still drops residual-backed rows whenever the path is non-empty. A successful schema-level step merges the parent `typed_value` validity into the output mask, so a row with null `typed_value` but a complete value in the residual column is forced null. The new mixed-row test uses an empty path and never executes this line. For `$.a`, `[{"a":1},{"a":2}]` with row 2 residual-backed returns `[1,null]`. Please implement per-row fallback/partitioning and cover root field, nested field, and list-index paths. -- 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]
