zeroshade commented on code in PR #1206: URL: https://github.com/apache/arrow-go/pull/1206#discussion_r3797971966
########## arrow/extensions/variant_get.go: ########## @@ -0,0 +1,453 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package extensions + +import ( + "fmt" + + "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/memory" + "github.com/apache/arrow-go/v18/parquet/variant" +) + +// VariantPathElement is a single step of a variant path: either an object field +// name or an array index. +type VariantPathElement struct { + name string + index int + isIndex bool +} + +// VariantPathField returns a path element selecting the named object field. +func VariantPathField(name string) VariantPathElement { + return VariantPathElement{name: name} +} + +// VariantPathIndex returns a path element selecting the array element at index. +func VariantPathIndex(index int) VariantPathElement { + return VariantPathElement{index: index, isIndex: true} +} + +// VariantPath is an ordered list of path elements to extract from a variant value. +type VariantPath []VariantPathElement + +// GetOptions controls VariantGet. +type GetOptions struct { + // Path is the path to extract from each variant value. + Path VariantPath + // AsType, when nil, makes VariantGet return a VariantArray pointing at the path. + // When set, the extracted value is cast to this type. Nested (struct/list) types + // are not yet supported and yield arrow.ErrNotImplemented. + AsType arrow.DataType + // Strict makes a cast failure return an error. The default (false) mirrors + // arrow-rs: a cast failure produces null. + Strict bool + // Mem is the allocator for output arrays; nil uses memory.DefaultAllocator. + Mem memory.Allocator +} + +// VariantGet extracts opts.Path from each value of a VariantArray. It follows the +// shredded typed_value columns as far as the path allows, then falls back to a +// per-row walk of the residual value for the remainder. +func VariantGet(input arrow.Array, opts GetOptions) (arrow.Array, error) { + va, ok := input.(*VariantArray) + if !ok { + return nil, fmt.Errorf("%w: VariantGet input must be a VariantArray, got %T", arrow.ErrInvalid, input) + } + + if opts.Mem == nil { + opts.Mem = memory.DefaultAllocator + } + + return shreddedGetPath(va, 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 stateFromVariant(va *VariantArray) shreddingState { + vt := va.ExtensionType().(*VariantType) + st := va.Storage().(*array.Struct) + + var value arrow.TypedArray[[]byte] + if vt.valueFieldIdx != -1 { + value = st.Field(vt.valueFieldIdx).(arrow.TypedArray[[]byte]) + } + + var typed arrow.Array + if vt.typedValueFieldIdx != -1 { + typed = st.Field(vt.typedValueFieldIdx) + } + + return shreddingState{value: value, typedValue: typed, length: va.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 +} + +// missingStep decides whether an absent typed field means the value is provably +// missing (value column all-null) or merely not shredded (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} +} + +// followFieldElement takes one field step deeper into the shredded columns. +func followFieldElement(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 +} + +func shreddedGetPath(va *VariantArray, opts GetOptions) (arrow.Array, error) { + state := stateFromVariant(va) + nulls := newNullTracker(va.Len()) + nulls.apply(va.Storage()) + + // Peel the field prefix of the path through the shredded columns. Index steps + // and non-shredded fields stop the columnar walk and hand the rest to a per-row + // fallback over the fully reassembled value at the current node. + idx := 0 + for idx < len(opts.Path) { + elem := opts.Path[idx] + if elem.isIndex { + break + } + + step, err := followFieldElement(state, elem.name) + if err != nil { + return nil, err + } + + switch step.kind { + case stepSuccess: + nulls.apply(state.typedValue) Review Comment: A valid shredded layout can have `typed_value == null` for one row while `value` contains that row's complete residual object. `VariantArray.Value()` would correctly reconstruct that row, but this is going to mark it null instead. For instance a two-row case of `[1, 2]` would return `[1, null]` if the `2` was in the residual `value` instead of in the `typedValue`. You should add root-level and intermediate nested mixed-row tests. ########## arrow/extensions/variant_get.go: ########## @@ -0,0 +1,453 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package extensions + +import ( + "fmt" + + "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/memory" + "github.com/apache/arrow-go/v18/parquet/variant" +) + +// VariantPathElement is a single step of a variant path: either an object field +// name or an array index. +type VariantPathElement struct { + name string + index int + isIndex bool +} + +// VariantPathField returns a path element selecting the named object field. +func VariantPathField(name string) VariantPathElement { + return VariantPathElement{name: name} +} + +// VariantPathIndex returns a path element selecting the array element at index. +func VariantPathIndex(index int) VariantPathElement { + return VariantPathElement{index: index, isIndex: true} +} + +// VariantPath is an ordered list of path elements to extract from a variant value. +type VariantPath []VariantPathElement + +// GetOptions controls VariantGet. +type GetOptions struct { + // Path is the path to extract from each variant value. + Path VariantPath + // AsType, when nil, makes VariantGet return a VariantArray pointing at the path. + // When set, the extracted value is cast to this type. Nested (struct/list) types + // are not yet supported and yield arrow.ErrNotImplemented. + AsType arrow.DataType + // Strict makes a cast failure return an error. The default (false) mirrors + // arrow-rs: a cast failure produces null. + Strict bool + // Mem is the allocator for output arrays; nil uses memory.DefaultAllocator. + Mem memory.Allocator +} + +// VariantGet extracts opts.Path from each value of a VariantArray. It follows the +// shredded typed_value columns as far as the path allows, then falls back to a +// per-row walk of the residual value for the remainder. +func VariantGet(input arrow.Array, opts GetOptions) (arrow.Array, error) { + va, ok := input.(*VariantArray) + if !ok { + return nil, fmt.Errorf("%w: VariantGet input must be a VariantArray, got %T", arrow.ErrInvalid, input) + } + + if opts.Mem == nil { + opts.Mem = memory.DefaultAllocator + } + + return shreddedGetPath(va, 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 stateFromVariant(va *VariantArray) shreddingState { + vt := va.ExtensionType().(*VariantType) + st := va.Storage().(*array.Struct) + + var value arrow.TypedArray[[]byte] + if vt.valueFieldIdx != -1 { + value = st.Field(vt.valueFieldIdx).(arrow.TypedArray[[]byte]) + } + + var typed arrow.Array + if vt.typedValueFieldIdx != -1 { + typed = st.Field(vt.typedValueFieldIdx) + } + + return shreddingState{value: value, typedValue: typed, length: va.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 +} + +// missingStep decides whether an absent typed field means the value is provably +// missing (value column all-null) or merely not shredded (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} +} + +// followFieldElement takes one field step deeper into the shredded columns. +func followFieldElement(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 +} + +func shreddedGetPath(va *VariantArray, opts GetOptions) (arrow.Array, error) { + state := stateFromVariant(va) + nulls := newNullTracker(va.Len()) + nulls.apply(va.Storage()) + + // Peel the field prefix of the path through the shredded columns. Index steps + // and non-shredded fields stop the columnar walk and hand the rest to a per-row + // fallback over the fully reassembled value at the current node. + idx := 0 + for idx < len(opts.Path) { + elem := opts.Path[idx] + if elem.isIndex { + break + } + + step, err := followFieldElement(state, elem.name) + if err != nil { + return nil, err + } + + switch step.kind { + case stepSuccess: + nulls.apply(state.typedValue) + state = step.state + idx++ + + continue + case stepMissing: + return allNullResult(va, opts) + } + + break // stepNotShredded + } + + remaining := opts.Path[idx:] + target, err := buildTargetVariant(va, state, nulls, opts.Mem) + if err != nil { + return nil, err + } + defer target.Release() + + if len(remaining) == 0 { + if opts.AsType == nil { + target.Retain() + + return target, nil + } + + if shredded := tryPerfectShredding(state, nulls, opts.AsType); shredded != nil { + return shredded, nil + } + } + + return shredBasicVariant(target, remaining, opts) +} + +// shredBasicVariant walks the remaining path per row and produces either a +// VariantArray (AsType nil) or a typed array. +func shredBasicVariant(target *VariantArray, remaining VariantPath, opts GetOptions) (arrow.Array, error) { + if opts.AsType == nil { + bldr := NewVariantBuilder(opts.Mem, NewDefaultVariantType()) + defer bldr.Release() + bldr.Reserve(target.Len()) + + for i := 0; i < target.Len(); i++ { + leaf, ok, err := navigateRow(target, i, remaining) + if err != nil { + return nil, err + } + if !ok { + bldr.AppendNull() + + continue + } + bldr.Append(leaf) + } + + return bldr.NewArray(), nil + } + + if _, ok := opts.AsType.(arrow.NestedType); ok { + return nil, fmt.Errorf("%w: VariantGet cast to nested type %s", arrow.ErrNotImplemented, opts.AsType) + } + + bldr := array.NewBuilder(opts.Mem, opts.AsType) + defer bldr.Release() + bldr.Reserve(target.Len()) + + for i := 0; i < target.Len(); i++ { + leaf, ok, err := navigateRow(target, i, remaining) + if err != nil { + return nil, err + } + if !ok || leaf.Type() == variant.Null { + bldr.AppendNull() + + continue + } + + if appendVariantToTypedBuilder(bldr, leaf) { Review Comment: The `appendVariantToTypedBuilder` method is for the shredding writer where only lossless physical fits should succeed. As a result, ordinary conversions would silently produce null under the default mode such as: * `int64` -> `float64` * `int64` -> `decimal128` * `decimal(scale 2)` -> `decimal(scale 10)` * `string` -> `binary` Once this moves to `arrow/compute` use the existing cast machinery there ########## arrow/extensions/variant_get.go: ########## @@ -0,0 +1,453 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package extensions + +import ( + "fmt" + + "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/memory" + "github.com/apache/arrow-go/v18/parquet/variant" +) + +// VariantPathElement is a single step of a variant path: either an object field +// name or an array index. +type VariantPathElement struct { + name string + index int + isIndex bool +} + +// VariantPathField returns a path element selecting the named object field. +func VariantPathField(name string) VariantPathElement { + return VariantPathElement{name: name} +} + +// VariantPathIndex returns a path element selecting the array element at index. +func VariantPathIndex(index int) VariantPathElement { + return VariantPathElement{index: index, isIndex: true} +} + +// VariantPath is an ordered list of path elements to extract from a variant value. +type VariantPath []VariantPathElement + +// GetOptions controls VariantGet. +type GetOptions struct { + // Path is the path to extract from each variant value. + Path VariantPath + // AsType, when nil, makes VariantGet return a VariantArray pointing at the path. + // When set, the extracted value is cast to this type. Nested (struct/list) types + // are not yet supported and yield arrow.ErrNotImplemented. + AsType arrow.DataType + // Strict makes a cast failure return an error. The default (false) mirrors + // arrow-rs: a cast failure produces null. + Strict bool + // Mem is the allocator for output arrays; nil uses memory.DefaultAllocator. + Mem memory.Allocator +} + +// VariantGet extracts opts.Path from each value of a VariantArray. It follows the +// shredded typed_value columns as far as the path allows, then falls back to a +// per-row walk of the residual value for the remainder. +func VariantGet(input arrow.Array, opts GetOptions) (arrow.Array, error) { + va, ok := input.(*VariantArray) + if !ok { + return nil, fmt.Errorf("%w: VariantGet input must be a VariantArray, got %T", arrow.ErrInvalid, input) + } + + if opts.Mem == nil { + opts.Mem = memory.DefaultAllocator + } + + return shreddedGetPath(va, 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 stateFromVariant(va *VariantArray) shreddingState { + vt := va.ExtensionType().(*VariantType) + st := va.Storage().(*array.Struct) + + var value arrow.TypedArray[[]byte] + if vt.valueFieldIdx != -1 { + value = st.Field(vt.valueFieldIdx).(arrow.TypedArray[[]byte]) + } + + var typed arrow.Array + if vt.typedValueFieldIdx != -1 { + typed = st.Field(vt.typedValueFieldIdx) + } + + return shreddingState{value: value, typedValue: typed, length: va.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 +} + +// missingStep decides whether an absent typed field means the value is provably +// missing (value column all-null) or merely not shredded (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} +} + +// followFieldElement takes one field step deeper into the shredded columns. +func followFieldElement(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 +} + +func shreddedGetPath(va *VariantArray, opts GetOptions) (arrow.Array, error) { + state := stateFromVariant(va) + nulls := newNullTracker(va.Len()) + nulls.apply(va.Storage()) + + // Peel the field prefix of the path through the shredded columns. Index steps + // and non-shredded fields stop the columnar walk and hand the rest to a per-row + // fallback over the fully reassembled value at the current node. + idx := 0 + for idx < len(opts.Path) { + elem := opts.Path[idx] + if elem.isIndex { + break + } + + step, err := followFieldElement(state, elem.name) + if err != nil { + return nil, err + } + + switch step.kind { + case stepSuccess: + nulls.apply(state.typedValue) + state = step.state + idx++ + + continue + case stepMissing: + return allNullResult(va, opts) + } + + break // stepNotShredded + } + + remaining := opts.Path[idx:] + target, err := buildTargetVariant(va, state, nulls, opts.Mem) + if err != nil { + return nil, err + } + defer target.Release() + + if len(remaining) == 0 { + if opts.AsType == nil { + target.Retain() + + return target, nil + } + + if shredded := tryPerfectShredding(state, nulls, opts.AsType); shredded != nil { Review Comment: try this *before* constructing the target, otherwise we're building and discarding a struct array and possible bitmap. ########## arrow/extensions/variant_get.go: ########## @@ -0,0 +1,453 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package extensions + +import ( + "fmt" + + "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/memory" + "github.com/apache/arrow-go/v18/parquet/variant" +) + +// VariantPathElement is a single step of a variant path: either an object field +// name or an array index. +type VariantPathElement struct { + name string + index int + isIndex bool +} + +// VariantPathField returns a path element selecting the named object field. +func VariantPathField(name string) VariantPathElement { + return VariantPathElement{name: name} +} + +// VariantPathIndex returns a path element selecting the array element at index. +func VariantPathIndex(index int) VariantPathElement { + return VariantPathElement{index: index, isIndex: true} +} + +// VariantPath is an ordered list of path elements to extract from a variant value. +type VariantPath []VariantPathElement + +// GetOptions controls VariantGet. +type GetOptions struct { + // Path is the path to extract from each variant value. + Path VariantPath + // AsType, when nil, makes VariantGet return a VariantArray pointing at the path. + // When set, the extracted value is cast to this type. Nested (struct/list) types + // are not yet supported and yield arrow.ErrNotImplemented. + AsType arrow.DataType + // Strict makes a cast failure return an error. The default (false) mirrors + // arrow-rs: a cast failure produces null. + Strict bool + // Mem is the allocator for output arrays; nil uses memory.DefaultAllocator. + Mem memory.Allocator +} + +// VariantGet extracts opts.Path from each value of a VariantArray. It follows the +// shredded typed_value columns as far as the path allows, then falls back to a +// per-row walk of the residual value for the remainder. +func VariantGet(input arrow.Array, opts GetOptions) (arrow.Array, error) { + va, ok := input.(*VariantArray) + if !ok { + return nil, fmt.Errorf("%w: VariantGet input must be a VariantArray, got %T", arrow.ErrInvalid, input) + } + + if opts.Mem == nil { + opts.Mem = memory.DefaultAllocator + } + + return shreddedGetPath(va, 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 stateFromVariant(va *VariantArray) shreddingState { + vt := va.ExtensionType().(*VariantType) + st := va.Storage().(*array.Struct) + + var value arrow.TypedArray[[]byte] + if vt.valueFieldIdx != -1 { + value = st.Field(vt.valueFieldIdx).(arrow.TypedArray[[]byte]) + } + + var typed arrow.Array + if vt.typedValueFieldIdx != -1 { + typed = st.Field(vt.typedValueFieldIdx) + } + + return shreddingState{value: value, typedValue: typed, length: va.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 +} + +// missingStep decides whether an absent typed field means the value is provably +// missing (value column all-null) or merely not shredded (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} +} + +// followFieldElement takes one field step deeper into the shredded columns. +func followFieldElement(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 +} + +func shreddedGetPath(va *VariantArray, opts GetOptions) (arrow.Array, error) { + state := stateFromVariant(va) + nulls := newNullTracker(va.Len()) + nulls.apply(va.Storage()) + + // Peel the field prefix of the path through the shredded columns. Index steps + // and non-shredded fields stop the columnar walk and hand the rest to a per-row + // fallback over the fully reassembled value at the current node. + idx := 0 + for idx < len(opts.Path) { + elem := opts.Path[idx] + if elem.isIndex { + break + } + + step, err := followFieldElement(state, elem.name) + if err != nil { + return nil, err + } + + switch step.kind { + case stepSuccess: + nulls.apply(state.typedValue) + state = step.state + idx++ + + continue + case stepMissing: + return allNullResult(va, opts) + } + + break // stepNotShredded + } + + remaining := opts.Path[idx:] + target, err := buildTargetVariant(va, state, nulls, opts.Mem) + if err != nil { + return nil, err + } + defer target.Release() + + if len(remaining) == 0 { + if opts.AsType == nil { + target.Retain() + + return target, nil + } + + if shredded := tryPerfectShredding(state, nulls, opts.AsType); shredded != nil { + return shredded, nil + } + } + + return shredBasicVariant(target, remaining, opts) +} + +// shredBasicVariant walks the remaining path per row and produces either a +// VariantArray (AsType nil) or a typed array. +func shredBasicVariant(target *VariantArray, remaining VariantPath, opts GetOptions) (arrow.Array, error) { + if opts.AsType == nil { + bldr := NewVariantBuilder(opts.Mem, NewDefaultVariantType()) + defer bldr.Release() + bldr.Reserve(target.Len()) + + for i := 0; i < target.Len(); i++ { + leaf, ok, err := navigateRow(target, i, remaining) + if err != nil { + return nil, err + } + if !ok { + bldr.AppendNull() + + continue + } + bldr.Append(leaf) + } + + return bldr.NewArray(), nil + } + + if _, ok := opts.AsType.(arrow.NestedType); ok { + return nil, fmt.Errorf("%w: VariantGet cast to nested type %s", arrow.ErrNotImplemented, opts.AsType) + } + + bldr := array.NewBuilder(opts.Mem, opts.AsType) + defer bldr.Release() + bldr.Reserve(target.Len()) + + for i := 0; i < target.Len(); i++ { + leaf, ok, err := navigateRow(target, i, remaining) + if err != nil { + return nil, err + } + if !ok || leaf.Type() == variant.Null { + bldr.AppendNull() + + continue + } + + if appendVariantToTypedBuilder(bldr, leaf) { + continue + } + + if opts.Strict { + return nil, fmt.Errorf("%w: cannot cast variant %v to %s", arrow.ErrInvalid, leaf.Type(), opts.AsType) + } + + bldr.AppendNull() + } + + return bldr.NewArray(), nil +} + +// navigateRow reassembles row i of target and walks path into it. It returns +// (value, false) when the row is null or the path is absent. +func navigateRow(target *VariantArray, i int, path VariantPath) (variant.Value, bool, error) { + if target.IsNull(i) { + return variant.Value{}, false, nil + } + + v, err := target.Value(i) + if err != nil { + return variant.Value{}, false, fmt.Errorf("variant: reassembling row %d: %w", i, err) + } + + return navigateValue(v, path) +} + +// navigateValue walks path into a fully reassembled variant value. +func navigateValue(v variant.Value, path VariantPath) (variant.Value, bool, error) { + cur := v + for _, elem := range path { + if elem.isIndex { + arr, ok := cur.Value().(variant.ArrayValue) + if !ok || elem.index < 0 || uint32(elem.index) >= arr.Len() { + return variant.Value{}, false, nil + } + el, err := arr.Value(uint32(elem.index)) + if err != nil { + return variant.Value{}, false, nil + } + cur = el + + continue + } + + obj, ok := cur.Value().(variant.ObjectValue) + if !ok { + return variant.Value{}, false, nil + } + field, err := obj.ValueByKey(elem.name) + if err != nil { + return variant.Value{}, false, nil + } Review Comment: Two distinct failures become a missing/null result: * Field access on a scalar returns null, while the arrow-rs implementation specifies a cast/type error * Every `ObjectValue.ValueByKey` error gets swallowed instead of propagated, only `arrow.ErrNotFound` is meaning absence, malformed metadata produces a different error that should propagate. Probing for corruption currently returns a nil error and a null output. ########## arrow/extensions/variant_get.go: ########## @@ -0,0 +1,453 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package extensions + +import ( + "fmt" + + "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/memory" + "github.com/apache/arrow-go/v18/parquet/variant" +) + +// VariantPathElement is a single step of a variant path: either an object field +// name or an array index. +type VariantPathElement struct { + name string + index int + isIndex bool +} + +// VariantPathField returns a path element selecting the named object field. +func VariantPathField(name string) VariantPathElement { + return VariantPathElement{name: name} +} + +// VariantPathIndex returns a path element selecting the array element at index. +func VariantPathIndex(index int) VariantPathElement { + return VariantPathElement{index: index, isIndex: true} +} + +// VariantPath is an ordered list of path elements to extract from a variant value. +type VariantPath []VariantPathElement + +// GetOptions controls VariantGet. +type GetOptions struct { + // Path is the path to extract from each variant value. + Path VariantPath + // AsType, when nil, makes VariantGet return a VariantArray pointing at the path. + // When set, the extracted value is cast to this type. Nested (struct/list) types + // are not yet supported and yield arrow.ErrNotImplemented. + AsType arrow.DataType + // Strict makes a cast failure return an error. The default (false) mirrors + // arrow-rs: a cast failure produces null. + Strict bool + // Mem is the allocator for output arrays; nil uses memory.DefaultAllocator. + Mem memory.Allocator +} + +// VariantGet extracts opts.Path from each value of a VariantArray. It follows the +// shredded typed_value columns as far as the path allows, then falls back to a +// per-row walk of the residual value for the remainder. +func VariantGet(input arrow.Array, opts GetOptions) (arrow.Array, error) { + va, ok := input.(*VariantArray) + if !ok { + return nil, fmt.Errorf("%w: VariantGet input must be a VariantArray, got %T", arrow.ErrInvalid, input) + } + + if opts.Mem == nil { + opts.Mem = memory.DefaultAllocator + } + + return shreddedGetPath(va, 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 Review Comment: this is never used, remove it? ########## arrow/extensions/variant_get.go: ########## @@ -0,0 +1,453 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package extensions + +import ( + "fmt" + + "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/memory" + "github.com/apache/arrow-go/v18/parquet/variant" +) + +// VariantPathElement is a single step of a variant path: either an object field +// name or an array index. +type VariantPathElement struct { + name string + index int + isIndex bool +} + +// VariantPathField returns a path element selecting the named object field. +func VariantPathField(name string) VariantPathElement { + return VariantPathElement{name: name} +} + +// VariantPathIndex returns a path element selecting the array element at index. +func VariantPathIndex(index int) VariantPathElement { + return VariantPathElement{index: index, isIndex: true} +} + +// VariantPath is an ordered list of path elements to extract from a variant value. +type VariantPath []VariantPathElement + +// GetOptions controls VariantGet. +type GetOptions struct { Review Comment: use `VariantGetOptions` instead of `GetOptions` since this is exported at the package level. ########## arrow/extensions/variant_get.go: ########## @@ -0,0 +1,453 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package extensions + +import ( + "fmt" + + "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/memory" + "github.com/apache/arrow-go/v18/parquet/variant" +) + +// VariantPathElement is a single step of a variant path: either an object field +// name or an array index. +type VariantPathElement struct { + name string + index int + isIndex bool +} + +// VariantPathField returns a path element selecting the named object field. +func VariantPathField(name string) VariantPathElement { + return VariantPathElement{name: name} +} + +// VariantPathIndex returns a path element selecting the array element at index. +func VariantPathIndex(index int) VariantPathElement { + return VariantPathElement{index: index, isIndex: true} +} + +// VariantPath is an ordered list of path elements to extract from a variant value. +type VariantPath []VariantPathElement + +// GetOptions controls VariantGet. +type GetOptions struct { + // Path is the path to extract from each variant value. + Path VariantPath + // AsType, when nil, makes VariantGet return a VariantArray pointing at the path. + // When set, the extracted value is cast to this type. Nested (struct/list) types + // are not yet supported and yield arrow.ErrNotImplemented. + AsType arrow.DataType + // Strict makes a cast failure return an error. The default (false) mirrors + // arrow-rs: a cast failure produces null. + Strict bool + // Mem is the allocator for output arrays; nil uses memory.DefaultAllocator. + Mem memory.Allocator +} + +// VariantGet extracts opts.Path from each value of a VariantArray. It follows the +// shredded typed_value columns as far as the path allows, then falls back to a +// per-row walk of the residual value for the remainder. +func VariantGet(input arrow.Array, opts GetOptions) (arrow.Array, error) { + va, ok := input.(*VariantArray) + if !ok { + return nil, fmt.Errorf("%w: VariantGet input must be a VariantArray, got %T", arrow.ErrInvalid, input) + } + + if opts.Mem == nil { + opts.Mem = memory.DefaultAllocator + } + + return shreddedGetPath(va, 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 stateFromVariant(va *VariantArray) shreddingState { + vt := va.ExtensionType().(*VariantType) + st := va.Storage().(*array.Struct) + + var value arrow.TypedArray[[]byte] + if vt.valueFieldIdx != -1 { + value = st.Field(vt.valueFieldIdx).(arrow.TypedArray[[]byte]) + } + + var typed arrow.Array + if vt.typedValueFieldIdx != -1 { + typed = st.Field(vt.typedValueFieldIdx) + } + + return shreddingState{value: value, typedValue: typed, length: va.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 +} + +// missingStep decides whether an absent typed field means the value is provably +// missing (value column all-null) or merely not shredded (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} +} + +// followFieldElement takes one field step deeper into the shredded columns. +func followFieldElement(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 +} + +func shreddedGetPath(va *VariantArray, opts GetOptions) (arrow.Array, error) { + state := stateFromVariant(va) + nulls := newNullTracker(va.Len()) + nulls.apply(va.Storage()) + + // Peel the field prefix of the path through the shredded columns. Index steps + // and non-shredded fields stop the columnar walk and hand the rest to a per-row + // fallback over the fully reassembled value at the current node. + idx := 0 + for idx < len(opts.Path) { + elem := opts.Path[idx] + if elem.isIndex { + break + } + + step, err := followFieldElement(state, elem.name) + if err != nil { + return nil, err + } + + switch step.kind { + case stepSuccess: + nulls.apply(state.typedValue) + state = step.state + idx++ + + continue + case stepMissing: + return allNullResult(va, opts) + } + + break // stepNotShredded + } + + remaining := opts.Path[idx:] + target, err := buildTargetVariant(va, state, nulls, opts.Mem) + if err != nil { + return nil, err + } + defer target.Release() + + if len(remaining) == 0 { + if opts.AsType == nil { + target.Retain() + + return target, nil + } + + if shredded := tryPerfectShredding(state, nulls, opts.AsType); shredded != nil { + return shredded, nil + } + } + + return shredBasicVariant(target, remaining, opts) +} + +// shredBasicVariant walks the remaining path per row and produces either a +// VariantArray (AsType nil) or a typed array. +func shredBasicVariant(target *VariantArray, remaining VariantPath, opts GetOptions) (arrow.Array, error) { + if opts.AsType == nil { + bldr := NewVariantBuilder(opts.Mem, NewDefaultVariantType()) + defer bldr.Release() + bldr.Reserve(target.Len()) + + for i := 0; i < target.Len(); i++ { + leaf, ok, err := navigateRow(target, i, remaining) + if err != nil { + return nil, err + } + if !ok { + bldr.AppendNull() + + continue + } + bldr.Append(leaf) + } + + return bldr.NewArray(), nil + } + + if _, ok := opts.AsType.(arrow.NestedType); ok { + return nil, fmt.Errorf("%w: VariantGet cast to nested type %s", arrow.ErrNotImplemented, opts.AsType) + } + + bldr := array.NewBuilder(opts.Mem, opts.AsType) + defer bldr.Release() + bldr.Reserve(target.Len()) + + for i := 0; i < target.Len(); i++ { + leaf, ok, err := navigateRow(target, i, remaining) + if err != nil { + return nil, err + } + if !ok || leaf.Type() == variant.Null { + bldr.AppendNull() + + continue + } + + if appendVariantToTypedBuilder(bldr, leaf) { + continue + } + + if opts.Strict { + return nil, fmt.Errorf("%w: cannot cast variant %v to %s", arrow.ErrInvalid, leaf.Type(), opts.AsType) + } + + bldr.AppendNull() + } + + return bldr.NewArray(), nil +} + +// navigateRow reassembles row i of target and walks path into it. It returns +// (value, false) when the row is null or the path is absent. +func navigateRow(target *VariantArray, i int, path VariantPath) (variant.Value, bool, error) { + if target.IsNull(i) { + return variant.Value{}, false, nil + } + + v, err := target.Value(i) + if err != nil { + return variant.Value{}, false, fmt.Errorf("variant: reassembling row %d: %w", i, err) + } + + return navigateValue(v, path) +} + +// navigateValue walks path into a fully reassembled variant value. +func navigateValue(v variant.Value, path VariantPath) (variant.Value, bool, error) { + cur := v + for _, elem := range path { + if elem.isIndex { + arr, ok := cur.Value().(variant.ArrayValue) + if !ok || elem.index < 0 || uint32(elem.index) >= arr.Len() { + return variant.Value{}, false, nil + } + el, err := arr.Value(uint32(elem.index)) Review Comment: a large index would wrap here. on a 64-bit platform, `VariantPathIndex(math.MaxUint32+1)` converts to zero and selects element 0. Should we check before narrowing? -- 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]
