nssalian commented on code in PR #1206: URL: https://github.com/apache/arrow-go/pull/1206#discussion_r3884740654
########## 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 Review Comment: VariantGet now rejects a nested AsType up front with arrow.ErrNotImplemented (restoring the guard that was dropped when the code moved to arrow/compute), before any leaf work — struct and list targets return the error instead of a silent all-null array. Covered by TestVariantGetNestedTypeNotImplemented (struct + list). Full nested materialization remains a follow-up. -- 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]
