This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new a8f60003 feat(compute): add list_element (#1138)
a8f60003 is described below
commit a8f600039d59a2de7e7232bef1216896d3c970b2
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 31 18:25:11 2026 +0200
feat(compute): add list_element (#1138)
### Rationale for this change
Arrow-Go supports the list, large-list, list-view, and fixed-size-list
array types, but callers currently need to read offsets and child arrays
themselves to select one element from every row. A list_element kernel
provides the common nested-data operation directly through compute.
### What changes are included in this PR?
- Add list_element for List, LargeList, ListView, LargeListView, and
FixedSizeList.
- Accept signed or unsigned integer scalar indexes. A one-element
integer array is accepted for matching one-row input.
- Return an array of the list element type while preserving parent nulls
and child nulls.
- Return errors for null or negative indexes, empty lists, and indexes
outside the length of a valid list.
- Use the existing take kernels for primitive, binary, and fixed-width
child values to avoid one array allocation per input row. Keep the
generic concatenate path for nested child values.
- Handle list-view offsets and sizes without requiring the input to be
materialized first.
- Reject StringView, BinaryView, and dictionary-backed extension
children, including nested children, before execution.
- Fix null-array construction and run-end encoded span metadata needed
by the nested fallback.
### Are these changes tested?
- go test ./arrow/compute/... -count=1
- Added coverage for sliced List and FixedSizeList inputs,
non-contiguous ListView offsets, unsigned indexes, string and
nested-list children, parent and child nulls, and invalid indexes.
- Added regression coverage for dictionary-backed extension rejection
and null run-end encoded values with nested dictionaries.
- Added BenchmarkListElement cases for 1K, 100K, and 1M rows with
allocation reporting.
### Are there any user-facing changes?
Yes. This adds the public `compute.ListElement` function. It also fixes
`array.MakeArrayOfNull` for list views, unions, and run-end encoded
arrays, keeps run-end encoded `ArraySpan` buffer metadata consistent,
and preserves storage validity when extracting extension scalars.
---
arrow/array/util.go | 87 +-
arrow/array/util_test.go | 129 ++
arrow/compute/exec/span.go | 8 +-
arrow/compute/exec/span_test.go | 56 +
arrow/compute/internal/kernels/scalar_nested.go | 602 ++++++
arrow/compute/registry.go | 1 +
arrow/compute/scalar_nested.go | 263 +++
arrow/compute/scalar_nested_test.go | 2489 +++++++++++++++++++++++
arrow/scalar/scalar.go | 4 +-
9 files changed, 3604 insertions(+), 35 deletions(-)
diff --git a/arrow/array/util.go b/arrow/array/util.go
index afd0f39b..ca692911 100644
--- a/arrow/array/util.go
+++ b/arrow/array/util.go
@@ -463,12 +463,18 @@ func getMaxBufferLen(dt arrow.DataType, length int) int {
bufferLen = maxOf(getMaxBufferLen(f.Type, 1))
}
return bufferLen
+ case *arrow.ListViewType:
+ return maxOf(dt.OffsetTypeTraits().BytesRequired(length))
+ case *arrow.LargeListViewType:
+ return maxOf(dt.OffsetTypeTraits().BytesRequired(length))
case arrow.OffsetsDataType:
return maxOf(dt.OffsetTypeTraits().BytesRequired(length + 1))
case arrow.BinaryViewDataType:
return maxOf(arrow.ViewHeaderSizeBytes * length)
case *arrow.FixedSizeListType:
return maxOf(getMaxBufferLen(dt.Elem(), int(dt.Len())*length))
+ case *arrow.RunEndEncodedType:
+ return bufferLen
case arrow.ExtensionType:
return maxOf(getMaxBufferLen(dt.StorageType(), length))
default:
@@ -497,7 +503,11 @@ func (n *nullArrayFactory) create() *Data {
childData []arrow.ArrayData
dictData arrow.ArrayData
)
- defer bufs[0].Release()
+ defer func() {
+ if bufs[0] != nil {
+ bufs[0].Release()
+ }
+ }()
if ex, ok := dt.(arrow.ExtensionType); ok {
dt = ex.StorageType()
@@ -520,41 +530,12 @@ func (n *nullArrayFactory) create() *Data {
bufs = append(bufs, n.buf)
case arrow.BinaryDataType:
bufs = append(bufs, n.buf, n.buf)
- case arrow.OffsetsDataType:
- bufs = append(bufs, n.buf)
- childData[0] = n.createChild(dt, 0, 0)
- defer childData[0].Release()
- case *arrow.FixedSizeListType:
- childData[0] = n.createChild(dt, 0, n.len*int(dt.Len()))
- defer childData[0].Release()
- case *arrow.StructType:
- for i := range dt.Fields() {
- childData[i] = n.createChild(dt, i, n.len)
- defer childData[i].Release()
- }
- case *arrow.RunEndEncodedType:
- bldr := NewBuilder(n.mem, dt.RunEnds())
- defer bldr.Release()
-
- switch b := bldr.(type) {
- case *Int16Builder:
- b.Append(int16(n.len))
- case *Int32Builder:
- b.Append(int32(n.len))
- case *Int64Builder:
- b.Append(int64(n.len))
- }
-
- childData[0] = bldr.newData()
- defer childData[0].Release()
- childData[1] = n.createChild(dt.Encoded(), 1, 1)
- defer childData[1].Release()
case arrow.UnionType:
bufs[0].Release()
bufs[0] = nil
bufs = append(bufs, n.buf)
// buffer is zeroed, but 0 may not be a valid type code
- if dt.TypeCodes()[0] != 0 {
+ if n.len > 0 && dt.TypeCodes()[0] != 0 {
bufs[1] = memory.NewResizableBuffer(n.mem)
bufs[1].Resize(n.len)
defer bufs[1].Release()
@@ -573,9 +554,51 @@ func (n *nullArrayFactory) create() *Data {
childData[i] = n.createChild(dt, i, childLen)
defer childData[i].Release()
}
+ case *arrow.ListViewType:
+ bufs = append(bufs, n.buf, n.buf)
+ childData[0] = n.createChild(dt, 0, 0)
+ defer childData[0].Release()
+ case *arrow.LargeListViewType:
+ bufs = append(bufs, n.buf, n.buf)
+ childData[0] = n.createChild(dt, 0, 0)
+ defer childData[0].Release()
+ case arrow.OffsetsDataType:
+ bufs = append(bufs, n.buf)
+ childData[0] = n.createChild(dt, 0, 0)
+ defer childData[0].Release()
+ case *arrow.FixedSizeListType:
+ childData[0] = n.createChild(dt, 0, n.len*int(dt.Len()))
+ defer childData[0].Release()
+ case *arrow.StructType:
+ for i := range dt.Fields() {
+ childData[i] = n.createChild(dt, i, n.len)
+ defer childData[i].Release()
+ }
+ case *arrow.RunEndEncodedType:
+ // Build the run ends separately from encoded types without
builders.
+ bldr := NewRunEndEncodedBuilder(n.mem, dt.RunEnds(), arrow.Null)
+ defer bldr.Release()
+
+ if n.len > 0 {
+ bldr.AppendNull()
+ bldr.ContinueRun(uint64(n.len - 1))
+ }
+
+ arr := bldr.NewArray()
+ defer arr.Release()
+ values := MakeArrayOfNull(n.mem, dt.Encoded(),
arr.Data().Children()[1].Len())
+ defer values.Release()
+ childData[0] = arr.Data().Children()[0]
+ childData[1] = values.Data()
+ bufs[0].Release()
+ bufs[0] = nil
}
- out := NewData(n.dt, n.len, bufs, childData, n.len, 0)
+ nulls := n.len
+ if dt.ID() == arrow.RUN_END_ENCODED || arrow.IsUnion(dt.ID()) {
+ nulls = 0
+ }
+ out := NewData(n.dt, n.len, bufs, childData, nulls, 0)
if dictData != nil {
out.SetDictionary(dictData)
}
diff --git a/arrow/array/util_test.go b/arrow/array/util_test.go
index 6870cf8b..e8e496c8 100644
--- a/arrow/array/util_test.go
+++ b/arrow/array/util_test.go
@@ -63,6 +63,135 @@ func TestGetDictArrayDataNullInSuffix(t *testing.T) {
assert.True(t, dict.IsNull(1))
}
+func TestMakeArrayOfNullListViews(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for _, typ := range []arrow.DataType{
+ arrow.ListViewOf(arrow.PrimitiveTypes.Int32),
+ arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32),
+ } {
+ t.Run(typ.String(), func(t *testing.T) {
+ for _, length := range []int{0, 1} {
+ arr := array.MakeArrayOfNull(mem, typ, length)
+ require.Equal(t, length, arr.Len())
+ require.Equal(t, length, arr.NullN())
+ require.NoError(t, array.ValidateFull(arr))
+ arr.Release()
+ }
+ })
+ }
+}
+
+func TestMakeArrayOfNullRunEndEncoded(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ typ := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int32)
+ for _, length := range []int{0, 1, 4} {
+ t.Run(fmt.Sprintf("length-%d", length), func(t *testing.T) {
+ arr := array.MakeArrayOfNull(mem, typ, length)
+ defer arr.Release()
+
+ require.Equal(t, length, arr.Len())
+ require.Equal(t, 0, arr.NullN())
+ require.NoError(t, array.ValidateFull(arr))
+
+ rle := arr.(*array.RunEndEncoded)
+ if length == 0 {
+ require.Zero(t, rle.RunEndsArr().Len())
+ require.Zero(t, rle.Values().Len())
+ } else {
+ require.Equal(t, 1, rle.RunEndsArr().Len())
+ require.Equal(t, 1, rle.Values().Len())
+ require.True(t, rle.Values().IsNull(0))
+ }
+ })
+ }
+}
+
+func TestMakeArrayOfNullRunEndEncodedNestedDictionary(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.ListOf(arrow.PrimitiveTypes.Int32),
+ }
+ for _, runEndType := range []arrow.DataType{
+ arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int64,
+ } {
+ t.Run(runEndType.String(), func(t *testing.T) {
+ typ := arrow.RunEndEncodedOf(runEndType, dictType)
+ for _, length := range []int{0, 1, 4} {
+ t.Run(fmt.Sprintf("length-%d", length), func(t
*testing.T) {
+ arr := array.MakeArrayOfNull(mem, typ,
length)
+ defer arr.Release()
+ require.Equal(t, length, arr.Len())
+ require.NoError(t,
array.ValidateFull(arr))
+ encoded := arr.(*array.RunEndEncoded)
+ values :=
encoded.Values().(*array.Dictionary)
+ require.Empty(t,
values.Dictionary().(*array.List).ListValues().(*array.Int32).Int32Values())
+ if length == 0 {
+ require.Zero(t,
encoded.RunEndsArr().Len())
+ require.Zero(t, values.Len())
+ } else {
+ require.Equal(t, 1,
encoded.RunEndsArr().Len())
+ require.Equal(t, 1,
values.Len())
+ require.True(t,
values.IsNull(0))
+ require.Equal(t,
fmt.Sprint(length), encoded.RunEndsArr().ValueStr(0))
+ }
+ })
+ }
+ })
+ }
+}
+
+func TestMakeArrayOfNullEmptyUnions(t *testing.T) {
+ for _, typ := range []arrow.DataType{arrow.SparseUnionOf(nil, nil),
arrow.DenseUnionOf(nil, nil)} {
+ t.Run(typ.ID().String(), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ values := array.MakeArrayOfNull(mem, typ, 0)
+ defer values.Release()
+ require.Zero(t, values.Len())
+ require.NoError(t, array.ValidateFull(values))
+ })
+ }
+}
+
+func TestMakeArrayOfNullUnions(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ types := []arrow.DataType{
+ arrow.SparseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type:
arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "text", Type: arrow.BinaryTypes.String,
Nullable: true},
+ },
+ []arrow.UnionTypeCode{5, 42},
+ ),
+ arrow.DenseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type:
arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "text", Type: arrow.BinaryTypes.String,
Nullable: true},
+ },
+ []arrow.UnionTypeCode{5, 42},
+ ),
+ }
+
+ for _, typ := range types {
+ t.Run(typ.String(), func(t *testing.T) {
+ arr := array.MakeArrayOfNull(mem, typ, 1)
+ defer arr.Release()
+
+ require.Equal(t, 0, arr.NullN())
+ require.NoError(t, array.ValidateFull(arr))
+ })
+ }
+}
+
var typemap = map[arrow.DataType]reflect.Type{
arrow.PrimitiveTypes.Int8: reflect.TypeOf(int8(0)),
arrow.PrimitiveTypes.Uint8: reflect.TypeOf(uint8(0)),
diff --git a/arrow/compute/exec/span.go b/arrow/compute/exec/span.go
index c989cfe7..49350eeb 100644
--- a/arrow/compute/exec/span.go
+++ b/arrow/compute/exec/span.go
@@ -183,6 +183,9 @@ func (a *ArraySpan) MakeData() arrow.ArrayData {
} else if dt.ID() == arrow.DENSE_UNION || dt.ID() == arrow.SPARSE_UNION
{
bufs[0] = nil
nulls = 0
+ } else if dt.ID() == arrow.RUN_END_ENCODED {
+ bufs[0] = nil
+ nulls = 0
}
if len(a.Children) > 0 {
@@ -578,10 +581,11 @@ type ExecSpan struct {
func getNumBuffers(dt arrow.DataType) int {
switch dt.ID() {
case arrow.RUN_END_ENCODED:
- return 0
+ return 1
case arrow.NULL, arrow.STRUCT, arrow.FIXED_SIZE_LIST:
return 1
- case arrow.BINARY, arrow.LARGE_BINARY, arrow.STRING,
arrow.LARGE_STRING, arrow.DENSE_UNION:
+ case arrow.BINARY, arrow.LARGE_BINARY, arrow.STRING, arrow.LARGE_STRING,
+ arrow.DENSE_UNION:
return 3
case arrow.BINARY_VIEW, arrow.STRING_VIEW:
// bitmap + view-header buffer + a single overflow data buffer.
diff --git a/arrow/compute/exec/span_test.go b/arrow/compute/exec/span_test.go
index feff91e3..90d4c93e 100644
--- a/arrow/compute/exec/span_test.go
+++ b/arrow/compute/exec/span_test.go
@@ -34,6 +34,7 @@ import (
"github.com/apache/arrow-go/v18/arrow/scalar"
"github.com/apache/arrow-go/v18/internal/types"
"github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
)
func TestBufferSpan_SetBuffer(t *testing.T) {
@@ -205,6 +206,7 @@ func TestArraySpan_NumBuffers(t *testing.T) {
{"null", fields{Type: arrow.Null}, 1},
{"struct", fields{Type: arrow.StructOf()}, 1},
{"fixed size list", fields{Type: arrow.FixedSizeListOf(4,
arrow.PrimitiveTypes.Int32)}, 1},
+ {"run end encoded", fields{Type:
arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32)},
1},
{"binary", fields{Type: arrow.BinaryTypes.Binary}, 3},
{"large binary", fields{Type: arrow.BinaryTypes.LargeBinary},
3},
{"string", fields{Type: arrow.BinaryTypes.String}, 3},
@@ -232,6 +234,60 @@ func TestArraySpan_NumBuffers(t *testing.T) {
}
}
+func TestArraySpanFillZeroLengthRunEndEncoded(t *testing.T) {
+ typ := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int32)
+ var span exec.ArraySpan
+ exec.FillZeroLength(typ, &span)
+
+ actual := span.MakeArray()
+ defer actual.Release()
+ require.NoError(t, array.ValidateFull(actual))
+}
+
+func TestArraySpan_MakeArrayPreservesListViewBuffers(t *testing.T) {
+ tests := []struct {
+ name string
+ build func(memory.Allocator) arrow.Array
+ }{
+ {"list view", func(mem memory.Allocator) arrow.Array {
+ builder := array.NewListViewBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ values := builder.ValueBuilder().(*array.Int32Builder)
+ values.AppendValues([]int32{10, 11, 12}, nil)
+ builder.AppendDimensions(1, 2)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+ }},
+ {"large list view", func(mem memory.Allocator) arrow.Array {
+ builder := array.NewLargeListViewBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ values := builder.ValueBuilder().(*array.Int32Builder)
+ values.AppendValues([]int32{10, 11, 12}, nil)
+ builder.AppendDimensions(1, 2)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+ }},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ input := tt.build(mem)
+ defer input.Release()
+
+ var span exec.ArraySpan
+ span.SetMembers(input.Data())
+ actual := span.MakeArray()
+ defer actual.Release()
+
+ assert.NoError(t, array.ValidateFull(actual))
+ assert.True(t, array.Equal(input, actual), "expected:
%s\ngot: %s", input, actual)
+ })
+ }
+}
+
func TestArraySpan_MakeData(t *testing.T) {
type fields struct {
Type arrow.DataType
diff --git a/arrow/compute/internal/kernels/scalar_nested.go
b/arrow/compute/internal/kernels/scalar_nested.go
new file mode 100644
index 00000000..dbf09527
--- /dev/null
+++ b/arrow/compute/internal/kernels/scalar_nested.go
@@ -0,0 +1,602 @@
+// 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.
+
+//go:build go1.18
+
+package kernels
+
+import (
+ "fmt"
+ "math"
+
+ "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/compute/exec"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/arrow/scalar"
+)
+
+func listElementOutputType(_ *exec.KernelCtx, inputTypes []arrow.DataType)
(arrow.DataType, error) {
+ listType, ok := inputTypes[0].(arrow.ListLikeType)
+ if !ok {
+ return nil, fmt.Errorf("%w: list_element requires a list-like
input", arrow.ErrType)
+ }
+ return listType.Elem(), nil
+}
+
+func getListElementIndex(value *exec.ExecValue) (uint64, error) {
+ if value.IsScalar() {
+ return listElementScalarIndex(value.Scalar)
+ }
+
+ if value.Array.Len == 0 {
+ return 0, fmt.Errorf("%w: list_element index array is empty",
arrow.ErrInvalid)
+ }
+ if value.Array.Len > 1 {
+ return 0, fmt.Errorf("%w: list_element does not support arrays
of list indices", arrow.ErrNotImplemented)
+ }
+ if value.Array.UpdateNullCount() != 0 {
+ return 0, fmt.Errorf("%w: list_element index must not contain
nulls", arrow.ErrInvalid)
+ }
+
+ switch value.Array.Type.ID() {
+ case arrow.INT8:
+ return unsignedIndex(exec.GetSpanValues[int8](&value.Array,
1)[0])
+ case arrow.INT16:
+ return unsignedIndex(exec.GetSpanValues[int16](&value.Array,
1)[0])
+ case arrow.INT32:
+ return unsignedIndex(exec.GetSpanValues[int32](&value.Array,
1)[0])
+ case arrow.INT64:
+ return unsignedIndex(exec.GetSpanValues[int64](&value.Array,
1)[0])
+ case arrow.UINT8:
+ return uint64(exec.GetSpanValues[uint8](&value.Array, 1)[0]),
nil
+ case arrow.UINT16:
+ return uint64(exec.GetSpanValues[uint16](&value.Array, 1)[0]),
nil
+ case arrow.UINT32:
+ return uint64(exec.GetSpanValues[uint32](&value.Array, 1)[0]),
nil
+ case arrow.UINT64:
+ return exec.GetSpanValues[uint64](&value.Array, 1)[0], nil
+ default:
+ return 0, fmt.Errorf("%w: invalid list_element index type %s",
arrow.ErrType, value.Array.Type)
+ }
+}
+
+func ValidateListElementScalarIndex(value scalar.Scalar) error {
+ _, err := listElementScalarIndex(value)
+ return err
+}
+
+func ListElementScalarIndex(value scalar.Scalar) (uint64, error) {
+ return listElementScalarIndex(value)
+}
+
+func listElementScalarIndex(value scalar.Scalar) (uint64, error) {
+ if !value.IsValid() {
+ return 0, fmt.Errorf("%w: list_element index must not be null",
arrow.ErrInvalid)
+ }
+ return scalarIndex(value)
+}
+
+func scalarIndex(value scalar.Scalar) (uint64, error) {
+ switch value := value.(type) {
+ case *scalar.Int8:
+ return unsignedIndex(value.Value)
+ case *scalar.Int16:
+ return unsignedIndex(value.Value)
+ case *scalar.Int32:
+ return unsignedIndex(value.Value)
+ case *scalar.Int64:
+ return unsignedIndex(value.Value)
+ case *scalar.Uint8:
+ return uint64(value.Value), nil
+ case *scalar.Uint16:
+ return uint64(value.Value), nil
+ case *scalar.Uint32:
+ return uint64(value.Value), nil
+ case *scalar.Uint64:
+ return value.Value, nil
+ default:
+ return 0, fmt.Errorf("%w: invalid list_element index type %s",
arrow.ErrType, value.DataType())
+ }
+}
+
+func unsignedIndex[T arrow.IntType](value T) (uint64, error) {
+ if value < 0 {
+ return 0, fmt.Errorf("%w: list_element index %d is out of
bounds: should be greater than or equal to 0", arrow.ErrInvalid, value)
+ }
+ return uint64(value), nil
+}
+
+func listElementValueOffsets(list *exec.ArraySpan, i int64) (int64, int64,
error) {
+ check := func(start, end int64) (int64, int64, error) {
+ if start < 0 || end < start || end > list.Children[0].Len {
+ return 0, 0, fmt.Errorf("%w: list_element input has
invalid value offsets", arrow.ErrInvalid)
+ }
+ return start, end, nil
+ }
+
+ switch list.Type.ID() {
+ case arrow.LIST:
+ offsets := exec.GetSpanOffsets[int32](list, 1)
+ return check(int64(offsets[i]), int64(offsets[i+1]))
+ case arrow.LARGE_LIST:
+ offsets := exec.GetSpanOffsets[int64](list, 1)
+ return check(offsets[i], offsets[i+1])
+ case arrow.LIST_VIEW:
+ offsets := exec.GetSpanValues[int32](list, 1)
+ sizes := exec.GetSpanValues[int32](list, 2)
+ start := int64(offsets[i])
+ size := int64(sizes[i])
+ if size < 0 {
+ return 0, 0, fmt.Errorf("%w: list_element input has
invalid value offsets", arrow.ErrInvalid)
+ }
+ return check(start, start+size)
+ case arrow.LARGE_LIST_VIEW:
+ offsets := exec.GetSpanValues[int64](list, 1)
+ sizes := exec.GetSpanValues[int64](list, 2)
+ start := offsets[i]
+ size := sizes[i]
+ if start < 0 || size < 0 || size > math.MaxInt64-start {
+ return 0, 0, fmt.Errorf("%w: list_element input has
invalid value offsets", arrow.ErrInvalid)
+ }
+ return check(start, start+size)
+ case arrow.FIXED_SIZE_LIST:
+ size := int64(list.Type.(*arrow.FixedSizeListType).Len())
+ if list.Offset < 0 || i < 0 || i > math.MaxInt64-list.Offset {
+ return 0, 0, fmt.Errorf("%w: list_element input has
invalid value offsets", arrow.ErrInvalid)
+ }
+ position := list.Offset + i
+ if size < 0 || (size > 0 && position > math.MaxInt64/size) {
+ return 0, 0, fmt.Errorf("%w: list_element input has
invalid value offsets", arrow.ErrInvalid)
+ }
+ start := position * size
+ if size > math.MaxInt64-start {
+ return 0, 0, fmt.Errorf("%w: list_element input has
invalid value offsets", arrow.ErrInvalid)
+ }
+ return check(start, start+size)
+ default:
+ return 0, 0, fmt.Errorf("%w: unsupported list_element input
type %s", arrow.ErrType, list.Type)
+ }
+}
+
+func listElementExec(ctx *exec.KernelCtx, batch *exec.ExecSpan, out
*exec.ExecResult) error {
+ var listSpan exec.ArraySpan
+ if batch.Values[0].IsScalar() {
+ listSpan.FillFromScalar(batch.Values[0].Scalar)
+ } else {
+ listSpan = batch.Values[0].Array
+ }
+ list := &listSpan
+ if len(list.Children) == 0 {
+ return fmt.Errorf("%w: list_element input has no values child",
arrow.ErrInvalid)
+ }
+
+ index, err := getListElementIndex(&batch.Values[1])
+ if err != nil {
+ return err
+ }
+
+ elemType := list.Type.(arrow.ListLikeType).Elem()
+ if !ListElementOutputTypeSupported(elemType) {
+ return fmt.Errorf("%w: list_element output type %s is not
supported", arrow.ErrNotImplemented, elemType)
+ }
+ if list.Len == 0 {
+ values := list.Children[0].MakeArray()
+ defer values.Release()
+ empty := array.NewSlice(values, 0, 0)
+ defer empty.Release()
+ out.TakeOwnership(empty.Data())
+ return nil
+ }
+ if !listElementTakeSupported(elemType) {
+ return listElementConcat(ctx, list, index, elemType, out)
+ }
+
+ indexBuilder := array.NewInt64Builder(exec.GetAllocator(ctx.Ctx))
+ defer indexBuilder.Release()
+ indexBuilder.Reserve(int(list.Len))
+ for i := int64(0); i < list.Len; i++ {
+ if len(list.Buffers[0].Buf) != 0 &&
bitutil.BitIsNotSet(list.Buffers[0].Buf, int(list.Offset+i)) {
+ indexBuilder.AppendNull()
+ continue
+ }
+
+ start, end, err := listElementValueOffsets(list, i)
+ if err != nil {
+ return err
+ }
+ if end < start {
+ return fmt.Errorf("%w: list_element input has invalid
value offsets", arrow.ErrInvalid)
+ }
+ length := uint64(end - start)
+ if index >= length {
+ return fmt.Errorf("%w: list_element index %d is out of
bounds: should be in [0, %d)", arrow.ErrInvalid, index, length)
+ }
+ indexBuilder.Append(start + int64(index))
+ }
+
+ indices := indexBuilder.NewArray()
+ defer indices.Release()
+ return listElementTakeOrFallback(ctx, &list.Children[0], indices, out)
+}
+
+func ListElementOutputTypeSupported(typ arrow.DataType) bool {
+ switch typ.ID() {
+ case arrow.BINARY_VIEW, arrow.STRING_VIEW:
+ return false
+ case arrow.SPARSE_UNION, arrow.DENSE_UNION:
+ if typ.(arrow.UnionType).NumFields() == 0 {
+ return false
+ }
+ case arrow.EXTENSION:
+ storageType := typ.(arrow.ExtensionType).StorageType()
+ // ArraySpan does not preserve a dictionary stored under an
extension.
+ return storageType.ID() != arrow.DICTIONARY &&
ListElementOutputTypeSupported(storageType)
+ case arrow.DICTIONARY:
+ return
ListElementOutputTypeSupported(typ.(*arrow.DictionaryType).ValueType)
+ }
+
+ nested, ok := typ.(arrow.NestedType)
+ if !ok {
+ return true
+ }
+ for _, field := range nested.Fields() {
+ if !ListElementOutputTypeSupported(field.Type) {
+ return false
+ }
+ }
+ return true
+}
+
+func listElementTakeSupported(typ arrow.DataType) bool {
+ id := typ.ID()
+ if id == arrow.NULL || arrow.IsBinaryLike(id) ||
arrow.IsLargeBinaryLike(id) ||
+ arrow.IsFixedSizeBinary(id) || id == arrow.SPARSE_UNION || id
== arrow.DENSE_UNION ||
+ id == arrow.EXTENSION {
+ return true
+ }
+ if !arrow.IsPrimitive(id) {
+ return false
+ }
+
+ // PrimitiveTake has specialized implementations for these widths only.
+ // In particular, INTERVAL_MONTH_DAY_NANO is a primitive 128-bit type
and
+ // must use the generic concatenation fallback below.
+ fixed, ok := typ.(arrow.FixedWidthDataType)
+ if !ok {
+ return false
+ }
+ switch fixed.BitWidth() {
+ case 1, 8, 16, 32, 64:
+ return true
+ default:
+ return false
+ }
+}
+
+func listElementConcat(ctx *exec.KernelCtx, list *exec.ArraySpan, index
uint64, elemType arrow.DataType, out *exec.ExecResult) error {
+ values := list.Children[0].MakeArray()
+ defer values.Release()
+ pieces := make([]arrow.Array, 0, int(list.Len))
+ defer func() {
+ for _, piece := range pieces {
+ piece.Release()
+ }
+ }()
+
+ for i := int64(0); i < list.Len; i++ {
+ if len(list.Buffers[0].Buf) != 0 &&
bitutil.BitIsNotSet(list.Buffers[0].Buf, int(list.Offset+i)) {
+ pieces = append(pieces, listElementMakeNullLike(ctx,
values, elemType))
+ continue
+ }
+
+ start, end, err := listElementValueOffsets(list, i)
+ if err != nil {
+ return err
+ }
+ if end < start {
+ return fmt.Errorf("%w: list_element input has invalid
value offsets", arrow.ErrInvalid)
+ }
+ length := uint64(end - start)
+ if index >= length {
+ return fmt.Errorf("%w: list_element index %d is out of
bounds: should be in [0, %d)", arrow.ErrInvalid, index, length)
+ }
+ selected := start + int64(index)
+ pieces = append(pieces, array.NewSlice(values, selected,
selected+1))
+ }
+
+ result, err := array.Concatenate(pieces, exec.GetAllocator(ctx.Ctx))
+ if err != nil {
+ return err
+ }
+ defer result.Release()
+ out.TakeOwnership(result.Data())
+ return nil
+}
+
+func listElementMakeNullLike(ctx *exec.KernelCtx, values arrow.Array, elemType
arrow.DataType) arrow.Array {
+ mem := exec.GetAllocator(ctx.Ctx)
+ if extType, ok := elemType.(arrow.ExtensionType); ok {
+ storageValues := values
+ if extValues, ok := values.(array.ExtensionArray); ok {
+ storageValues = extValues.Storage()
+ }
+ storage := listElementMakeNullLike(ctx, storageValues,
extType.StorageType())
+ result := array.NewExtensionArrayWithStorage(extType, storage)
+ storage.Release()
+ return result
+ }
+ if elemType.ID() == arrow.RUN_END_ENCODED {
+ runEndType := elemType.(*arrow.RunEndEncodedType)
+ // Build only the run ends; nested encoded values may not have
a builder.
+ builder := array.NewRunEndEncodedBuilder(mem,
runEndType.RunEnds(), arrow.Null)
+ defer builder.Release()
+ builder.AppendNull()
+ runEnds := builder.NewRunEndEncodedArray()
+ defer runEnds.Release()
+ nulls := listElementMakeNullLike(ctx,
values.(*array.RunEndEncoded).Values(), runEndType.Encoded())
+ defer nulls.Release()
+ return array.NewRunEndEncodedArrayWithType(runEndType,
runEnds.RunEndsArr(), nulls, 1, 0)
+ }
+ if values.Len() == 0 || len(values.Data().Buffers()) == 0 ||
+ elemType.ID() == arrow.NULL || arrow.IsUnion(elemType.ID()) {
+ return array.MakeArrayOfNull(mem, elemType, 1)
+ }
+
+ source := array.NewSlice(values, 0, 1)
+ defer source.Release()
+
+ sourceData := source.Data()
+ validity := memory.NewResizableBuffer(mem)
+ validity.Resize(int(bitutil.BytesForBits(int64(sourceData.Offset() +
1))))
+ memory.Set(validity.Bytes(), 0)
+ defer validity.Release()
+
+ buffers := append([]*memory.Buffer(nil), sourceData.Buffers()...)
+ buffers[0] = validity
+ data := array.NewData(sourceData.DataType(), 1, buffers,
sourceData.Children(), 1, sourceData.Offset())
+ if dictionary := sourceData.Dictionary(); dictionary != nil {
+ data.SetDictionary(dictionary)
+ }
+ defer data.Release()
+ return array.MakeFromData(data)
+}
+
+func listElementTakeFallback(ctx *exec.KernelCtx, values *exec.ArraySpan,
indices arrow.Array, out *exec.ExecResult) error {
+ elemType := values.Type
+ valuesArray := values.MakeArray()
+ defer valuesArray.Release()
+ if indices.Len() == 0 {
+ empty := array.NewSlice(valuesArray, 0, 0)
+ defer empty.Release()
+ out.TakeOwnership(empty.Data())
+ return nil
+ }
+
+ pieces := make([]arrow.Array, 0, indices.Len())
+ defer func() {
+ for _, piece := range pieces {
+ piece.Release()
+ }
+ }()
+
+ for i := 0; i < indices.Len(); i++ {
+ if indices.IsNull(i) {
+ pieces = append(pieces, listElementMakeNullLike(ctx,
valuesArray, elemType))
+ continue
+ }
+ selected, err := listElementTakeIndex(indices, i)
+ if err != nil {
+ return err
+ }
+ pieces = append(pieces, array.NewSlice(valuesArray, selected,
selected+1))
+ }
+
+ result, err := array.Concatenate(pieces, exec.GetAllocator(ctx.Ctx))
+ if err != nil {
+ return err
+ }
+ defer result.Release()
+ out.TakeOwnership(result.Data())
+ return nil
+}
+
+func listElementTakeIndex(indices arrow.Array, i int) (int64, error) {
+ switch indexArray := indices.(type) {
+ case *array.Int32:
+ return int64(indexArray.Value(i)), nil
+ case *array.Int64:
+ return indexArray.Value(i), nil
+ default:
+ return 0, fmt.Errorf("%w: list_element fallback received
unsupported index type %s", arrow.ErrType, indices.DataType())
+ }
+}
+
+func listElementTakeOrFallback(ctx *exec.KernelCtx, values *exec.ArraySpan,
indices arrow.Array, out *exec.ExecResult) error {
+ if indices.Len() == 0 {
+ return listElementTakeFallback(ctx, values, indices, out)
+ }
+ if handled, err := listElementTake(ctx, values, indices, out); handled {
+ return err
+ }
+ return listElementTakeFallback(ctx, values, indices, out)
+}
+
+func listElementDenseUnionTake(ctx *exec.KernelCtx, values *exec.ArraySpan,
indices arrow.Array, out *exec.ExecResult) error {
+ var indexSpan exec.ArraySpan
+ indexSpan.SetMembers(indices.Data())
+ batch := &exec.ExecSpan{
+ Len: int64(indices.Len()),
+ Values: []exec.ExecValue{
+ {Array: *values},
+ {Array: indexSpan},
+ },
+ }
+ takeCtx := *ctx
+ takeCtx.State = TakeOptions{BoundsCheck: false}
+ if err := TakeExec(DenseUnionImpl)(&takeCtx, batch, out); err != nil {
+ return err
+ }
+
+ for i := range out.Children {
+ childIndices := out.Children[i].MakeArray()
+ out.Children[i] = exec.ArraySpan{}
+ childOut := &exec.ExecResult{Type: values.Children[i].Type}
+ err := listElementTakeOrFallback(ctx, &values.Children[i],
childIndices, childOut)
+ childIndices.Release()
+ if err != nil {
+ childOut.Release()
+ return err
+ }
+
+ childData := childOut.MakeData()
+ out.Children[i].TakeOwnership(childData)
+ childData.Release()
+ }
+ return nil
+}
+
+func listElementSparseUnionTake(ctx *exec.KernelCtx, values *exec.ArraySpan,
indices arrow.Array, out *exec.ExecResult) error {
+ valuesArray := values.MakeArray()
+ defer valuesArray.Release()
+
+ union := valuesArray.(*array.SparseUnion)
+ unionType := values.Type.(*arrow.SparseUnionType)
+ typeIDBuilder := array.NewInt8Builder(exec.GetAllocator(ctx.Ctx))
+ defer typeIDBuilder.Release()
+ typeIDBuilder.Reserve(indices.Len())
+
+ for i := 0; i < indices.Len(); i++ {
+ if indices.IsNull(i) {
+ typeIDBuilder.Append(unionType.TypeCodes()[0])
+ } else {
+ selected, err := listElementTakeIndex(indices, i)
+ if err != nil {
+ return err
+ }
+ typeIDBuilder.Append(union.TypeCode(int(selected)))
+ }
+ }
+
+ typeIDs := typeIDBuilder.NewArray()
+ defer typeIDs.Release()
+
+ children := make([]arrow.Array, len(values.Children))
+ defer func() {
+ for _, child := range children {
+ if child != nil {
+ child.Release()
+ }
+ }
+ }()
+ for i := range values.Children {
+ field := union.Field(i)
+ var childSpan exec.ArraySpan
+ childSpan.SetMembers(field.Data())
+ childOut := &exec.ExecResult{Type: childSpan.Type}
+ if err := listElementTakeOrFallback(ctx, &childSpan, indices,
childOut); err != nil {
+ childOut.Release()
+ return err
+ }
+ children[i] = childOut.MakeArray()
+ }
+
+ childData := make([]arrow.ArrayData, len(children))
+ for i, child := range children {
+ childData[i] = child.Data()
+ }
+ data := array.NewData(unionType, indices.Len(),
+ []*memory.Buffer{nil, typeIDs.Data().Buffers()[1]}, childData,
0, 0)
+ result := array.NewSparseUnionData(data)
+ data.Release()
+ defer result.Release()
+ out.TakeOwnership(result.Data())
+ return nil
+}
+
+func listElementTake(ctx *exec.KernelCtx, values *exec.ArraySpan, indices
arrow.Array, out *exec.ExecResult) (bool, error) {
+ if !listElementTakeSupported(values.Type) {
+ return false, nil
+ }
+
+ var indexSpan exec.ArraySpan
+ indexSpan.SetMembers(indices.Data())
+ batch := &exec.ExecSpan{
+ Len: int64(indices.Len()),
+ Values: []exec.ExecValue{
+ {Array: *values},
+ {Array: indexSpan},
+ },
+ }
+ takeCtx := *ctx
+ takeCtx.State = TakeOptions{BoundsCheck: false}
+
+ switch id := values.Type.ID(); {
+ case id == arrow.NULL:
+ return true, NullTake(&takeCtx, batch, out)
+ case arrow.IsPrimitive(id):
+ return true, PrimitiveTake(&takeCtx, batch, out)
+ case arrow.IsBinaryLike(id):
+ return true, TakeExec(VarBinaryImpl[int32])(&takeCtx, batch,
out)
+ case arrow.IsLargeBinaryLike(id):
+ return true, TakeExec(VarBinaryImpl[int64])(&takeCtx, batch,
out)
+ case arrow.IsFixedSizeBinary(id):
+ return true, TakeExec(FSBImpl)(&takeCtx, batch, out)
+ case id == arrow.SPARSE_UNION:
+ return true, listElementSparseUnionTake(ctx, values, indices,
out)
+ case id == arrow.DENSE_UNION:
+ return true, listElementDenseUnionTake(ctx, values, indices,
out)
+ case id == arrow.EXTENSION:
+ extType := values.Type.(arrow.ExtensionType)
+ storage := *values
+ storage.Type = extType.StorageType()
+ handled, err := listElementTake(ctx, &storage, indices, out)
+ if handled {
+ // The storage take produces the physical buffers and
children. Restore
+ // the logical extension type so ArraySpan.MakeData
reconstructs an
+ // ExtensionArray around them.
+ out.Type = values.Type
+ }
+ return handled, err
+ default:
+ return false, nil
+ }
+}
+
+func GetListElementKernels() []exec.ScalarKernel {
+ kernels := make([]exec.ScalarKernel, 0, 5)
+ for _, listID := range []arrow.Type{
+ arrow.LIST,
+ arrow.LARGE_LIST,
+ arrow.LIST_VIEW,
+ arrow.LARGE_LIST_VIEW,
+ arrow.FIXED_SIZE_LIST,
+ } {
+ kernel := exec.NewScalarKernel(
+ []exec.InputType{
+ exec.NewIDInput(listID),
+ exec.NewMatchedInput(exec.Integer()),
+ },
+ exec.NewComputedOutputType(listElementOutputType),
+ listElementExec,
+ nil)
+ kernel.NullHandling = exec.NullComputedNoPrealloc
+ kernel.MemAlloc = exec.MemNoPrealloc
+ kernels = append(kernels, kernel)
+ }
+ return kernels
+}
diff --git a/arrow/compute/registry.go b/arrow/compute/registry.go
index bea37025..cccf58ef 100644
--- a/arrow/compute/registry.go
+++ b/arrow/compute/registry.go
@@ -49,6 +49,7 @@ func GetFunctionRegistry() FunctionRegistry {
registry = NewRegistry()
RegisterScalarCast(registry)
RegisterVectorSelection(registry)
+ RegisterScalarNested(registry)
RegisterVectorSort(registry)
RegisterScalarBoolean(registry)
RegisterScalarArithmetic(registry)
diff --git a/arrow/compute/scalar_nested.go b/arrow/compute/scalar_nested.go
new file mode 100644
index 00000000..23d1c56d
--- /dev/null
+++ b/arrow/compute/scalar_nested.go
@@ -0,0 +1,263 @@
+// 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.
+
+//go:build go1.18
+
+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/compute/internal/kernels"
+ "github.com/apache/arrow-go/v18/arrow/scalar"
+)
+
+var listElementDoc = FunctionDoc{
+ Summary: "Compute elements using nested list values and an index",
+ Description: "For each list value, return the element at the requested
index.\n" +
+ "Use an integral scalar for broadcast selection. A one-element
array-like\n" +
+ "index is accepted only for a matching one-element execution
and is not\n" +
+ "broadcast over longer input. Supported inputs are List,
LargeList,\n" +
+ "ListView, LargeListView, and FixedSizeList; Map is not
supported.\n" +
+ "Null lists and null child values produce null results. A null
or negative\n" +
+ "index is invalid, as is an index outside any non-null list's
length\n" +
+ "(including an empty list). StringView, BinaryView, and
dictionary-backed\n" +
+ "extension children are not supported, including when nested
inside\n" +
+ "another child type.",
+ ArgNames: []string{"lists", "index"},
+}
+
+func validateListElementIndex(index Datum) error {
+ switch index.Kind() {
+ case KindScalar:
+ value, ok := index.(*ScalarDatum)
+ if !ok || value.Value == nil {
+ return fmt.Errorf("%w: list_element requires a scalar
index datum", arrow.ErrType)
+ }
+ return kernels.ValidateListElementScalarIndex(value.Value)
+ case KindArray, KindChunked:
+ if index.Len() == 0 {
+ return fmt.Errorf("%w: list_element index array is
empty", arrow.ErrInvalid)
+ }
+ if index.Len() > 1 {
+ return fmt.Errorf("%w: list_element does not support
arrays of list indices", arrow.ErrNotImplemented)
+ }
+ }
+ return nil
+}
+
+func validateListElementOutputType(lists Datum) error {
+ var typ arrow.DataType
+ switch lists.Kind() {
+ case KindScalar:
+ value, ok := lists.(*ScalarDatum)
+ if !ok || value.Value == nil {
+ return nil
+ }
+ typ = value.Value.DataType()
+ case KindArray, KindChunked:
+ value, ok := lists.(ArrayLikeDatum)
+ if !ok {
+ return nil
+ }
+ typ = value.Type()
+ default:
+ return nil
+ }
+
+ listType, ok := typ.(arrow.ListLikeType)
+ if !ok || kernels.ListElementOutputTypeSupported(listType.Elem()) {
+ return nil
+ }
+ return fmt.Errorf("%w: list_element output type %s is not supported",
arrow.ErrNotImplemented, listType.Elem())
+}
+
+type listElementFunction struct {
+ ScalarFunction
+}
+
+func listElementScalarResultSupported(typ arrow.DataType, nullResult bool)
bool {
+ switch typ.ID() {
+ case arrow.BINARY_VIEW, arrow.STRING_VIEW, arrow.LIST_VIEW,
arrow.LARGE_LIST_VIEW,
+ arrow.DECIMAL32, arrow.DECIMAL64:
+ return false
+ case arrow.EXTENSION:
+ return
listElementScalarResultSupported(typ.(arrow.ExtensionType).StorageType(),
nullResult)
+ case arrow.STRUCT:
+ for _, field := range typ.(*arrow.StructType).Fields() {
+ if !listElementScalarResultSupported(field.Type,
nullResult) {
+ return false
+ }
+ }
+ case arrow.SPARSE_UNION:
+ if typ.(arrow.UnionType).NumFields() == 0 {
+ return false
+ }
+ for _, field := range typ.(arrow.UnionType).Fields() {
+ if !listElementScalarResultSupported(field.Type,
nullResult) {
+ return false
+ }
+ }
+ case arrow.DENSE_UNION:
+ if !nullResult {
+ return true
+ }
+ fields := typ.(arrow.UnionType).Fields()
+ return len(fields) != 0 &&
listElementScalarResultSupported(fields[0].Type, true)
+ case arrow.RUN_END_ENCODED:
+ return
listElementScalarResultSupported(typ.(*arrow.RunEndEncodedType).Encoded(),
nullResult)
+ }
+ return true
+}
+
+func listElementScalarArrayValueSupported(values arrow.Array, index int) bool {
+ if values.IsNull(index) {
+ return listElementScalarResultSupported(values.DataType(), true)
+ }
+
+ switch values := values.(type) {
+ case *array.BinaryView, *array.StringView, *array.ListView,
*array.LargeListView,
+ *array.Decimal32, *array.Decimal64:
+ return false
+ case array.ExtensionArray:
+ return listElementScalarArrayValueSupported(values.Storage(),
index)
+ case *array.Struct:
+ for i := 0; i < values.NumField(); i++ {
+ if
!listElementScalarArrayValueSupported(values.Field(i), index) {
+ return false
+ }
+ }
+ case *array.SparseUnion:
+ for i := 0; i < values.NumFields(); i++ {
+ if
!listElementScalarArrayValueSupported(values.Field(i), index) {
+ return false
+ }
+ }
+ case *array.DenseUnion:
+ child := values.Field(values.ChildID(index))
+ if child == nil {
+ return false
+ }
+ offset := values.ValueOffset(index)
+ if offset < 0 || int64(offset) >= int64(child.Len()) {
+ return false
+ }
+ return listElementScalarArrayValueSupported(child, int(offset))
+ case *array.RunEndEncoded:
+ return listElementScalarArrayValueSupported(values.Values(),
values.GetPhysicalIndex(index))
+ }
+ return true
+}
+
+// Validate the index before scalar execution splits arguments into spans. The
+// index contract is defined by the original Datum, not by the length of an
+// execution span.
+func (fn *listElementFunction) Execute(ctx context.Context, opts
FunctionOptions, args ...Datum) (Datum, error) {
+ if err := fn.checkArity(len(args)); err != nil {
+ return nil, err
+ }
+ if err := checkOptions(fn, opts); err != nil {
+ return nil, err
+ }
+
+ if err := validateListElementIndex(args[1]); err != nil {
+ return nil, err
+ }
+ if err := validateListElementOutputType(args[0]); err != nil {
+ return nil, err
+ }
+ if args[0].Kind() == KindScalar && args[1].Kind() == KindScalar {
+ indexDatum, ok := args[1].(*ScalarDatum)
+ if !ok {
+ return nil, fmt.Errorf("%w: list_element requires a
scalar index datum", arrow.ErrType)
+ }
+ listDatum, ok := args[0].(*ScalarDatum)
+ if !ok {
+ return nil, fmt.Errorf("%w: list_element requires a
list-like input", arrow.ErrType)
+ }
+ if listDatum.Value == nil {
+ return nil, fmt.Errorf("%w: list_element requires a
list-like scalar input", arrow.ErrType)
+ }
+
+ listType, ok := listDatum.Type().(arrow.ListLikeType)
+ if !ok {
+ return nil, fmt.Errorf("%w: list_element requires a
list-like input", arrow.ErrType)
+ }
+ listValue, ok := listDatum.Value.(scalar.ListScalar)
+ if !ok {
+ return nil, fmt.Errorf("%w: list_element requires a
list-like scalar input", arrow.ErrType)
+ }
+ if !listElementScalarResultSupported(listType.Elem(),
!listValue.IsValid()) {
+ return nil, fmt.Errorf("%w: list_element scalar output
type %s is not supported", arrow.ErrNotImplemented, listType.Elem())
+ }
+ if !listValue.IsValid() {
+ if _, err := fn.DispatchExact(listDatum.Type(),
indexDatum.Type()); err != nil {
+ return nil, err
+ }
+ if err := context.Cause(ctx); err != nil {
+ return nil, err
+ }
+ // Null scalar construction does not need array
concatenation or
+ // unboxing, which do not support every nested scalar
type.
+ return &ScalarDatum{Value:
scalar.MakeNullScalar(listType.Elem())}, nil
+ }
+ if listValue.GetList() != nil {
+ index, err :=
kernels.ListElementScalarIndex(indexDatum.Value)
+ if err != nil {
+ return nil, err
+ }
+ values := listValue.GetList()
+ if index < uint64(values.Len()) {
+ if
!listElementScalarArrayValueSupported(values, int(index)) {
+ return nil, fmt.Errorf("%w:
list_element scalar output value %s is not supported", arrow.ErrNotImplemented,
values.DataType())
+ }
+ if values.IsNull(int(index)) {
+ if _, err :=
fn.DispatchExact(listDatum.Type(), indexDatum.Type()); err != nil {
+ return nil, err
+ }
+ if err := context.Cause(ctx); err !=
nil {
+ return nil, err
+ }
+ value, err := scalar.GetScalar(values,
int(index))
+ if err != nil {
+ return nil, err
+ }
+ return &ScalarDatum{Value: value}, nil
+ }
+ }
+ }
+ }
+
+ return fn.ScalarFunction.Execute(ctx, opts, args...)
+}
+
+func RegisterScalarNested(reg FunctionRegistry) {
+ fn := &listElementFunction{ScalarFunction:
*NewScalarFunction("list_element", Binary(), listElementDoc)}
+ for _, kernel := range kernels.GetListElementKernels() {
+ if err := fn.AddKernel(kernel); err != nil {
+ panic(err)
+ }
+ }
+
+ reg.AddFunction(fn, false)
+}
+
+func ListElement(ctx context.Context, lists, index Datum) (Datum, error) {
+ return CallFunction(ctx, "list_element", nil, lists, index)
+}
diff --git a/arrow/compute/scalar_nested_test.go
b/arrow/compute/scalar_nested_test.go
new file mode 100644
index 00000000..4030c188
--- /dev/null
+++ b/arrow/compute/scalar_nested_test.go
@@ -0,0 +1,2489 @@
+// 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.
+
+//go:build go1.18
+
+package compute_test
+
+import (
+ "context"
+ "fmt"
+ "reflect"
+ "strings"
+ "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/decimal"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/arrow/scalar"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+type denseUnionExtensionArray struct {
+ array.ExtensionArrayBase
+}
+
+func (a *denseUnionExtensionArray) ValueStr(i int) string {
+ if a.IsNull(i) {
+ return array.NullValueStr
+ }
+ return "dense_union"
+}
+
+type denseUnionExtensionType struct {
+ arrow.ExtensionBase
+}
+
+func (denseUnionExtensionType) ArrayType() reflect.Type {
+ return reflect.TypeOf(denseUnionExtensionArray{})
+}
+
+func (denseUnionExtensionType) ExtensionName() string {
+ return "compute-test.dense-union"
+}
+
+func (t *denseUnionExtensionType) ExtensionEquals(other arrow.ExtensionType)
bool {
+ rhs, ok := other.(*denseUnionExtensionType)
+ return ok && arrow.TypeEqual(t.StorageType(), rhs.StorageType())
+}
+
+func (denseUnionExtensionType) Serialize() string { return "" }
+
+func (t *denseUnionExtensionType) Deserialize(storage arrow.DataType, _
string) (arrow.ExtensionType, error) {
+ return &denseUnionExtensionType{ExtensionBase:
arrow.ExtensionBase{Storage: storage}}, nil
+}
+
+type runEndExtensionArray struct {
+ array.ExtensionArrayBase
+}
+
+func (a *runEndExtensionArray) ValueStr(i int) string {
+ return a.Storage().ValueStr(i)
+}
+
+type runEndExtensionType struct {
+ arrow.ExtensionBase
+}
+
+func (runEndExtensionType) ArrayType() reflect.Type {
+ return reflect.TypeOf(runEndExtensionArray{})
+}
+
+func (runEndExtensionType) ExtensionName() string {
+ return "compute-test.run-end"
+}
+
+func (t *runEndExtensionType) ExtensionEquals(other arrow.ExtensionType) bool {
+ rhs, ok := other.(*runEndExtensionType)
+ return ok && arrow.TypeEqual(t.StorageType(), rhs.StorageType())
+}
+
+func (runEndExtensionType) Serialize() string { return "" }
+
+func (t *runEndExtensionType) Deserialize(storage arrow.DataType, _ string)
(arrow.ExtensionType, error) {
+ return &runEndExtensionType{ExtensionBase: arrow.ExtensionBase{Storage:
storage}}, nil
+}
+
+func listElementInput(t *testing.T, mem memory.Allocator, typ arrow.DataType,
values string) arrow.Array {
+ arr, _, err := array.FromJSON(mem, typ, strings.NewReader(values))
+ require.NoError(t, err)
+ return arr
+}
+
+func TestListElement(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ types := []arrow.DataType{
+ arrow.ListOf(arrow.PrimitiveTypes.Int32),
+ arrow.LargeListOf(arrow.PrimitiveTypes.Int32),
+ arrow.ListViewOf(arrow.PrimitiveTypes.Int32),
+ arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32),
+ arrow.FixedSizeListOf(2, arrow.PrimitiveTypes.Int32),
+ }
+ for _, typ := range types {
+ t.Run(typ.String(), func(t *testing.T) {
+ input := listElementInput(t, mem, typ, `[[1, 2], [3,
4], null, [5, 6]]`)
+ defer input.Release()
+ expected := listElementInput(t, mem,
arrow.PrimitiveTypes.Int32, `[2, 4, null, 6]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual),
"expected: %s\ngot: %s", expected, actual)
+ })
+ }
+}
+
+func TestListElementAllIntegerIndexTypes(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ listTypes := []arrow.DataType{
+ arrow.ListOf(arrow.PrimitiveTypes.Int32),
+ arrow.LargeListOf(arrow.PrimitiveTypes.Int32),
+ arrow.ListViewOf(arrow.PrimitiveTypes.Int32),
+ arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32),
+ arrow.FixedSizeListOf(2, arrow.PrimitiveTypes.Int32),
+ }
+ indexes := []struct {
+ name string
+ typ arrow.DataType
+ scalar scalar.Scalar
+ }{
+ {name: "int8", typ: arrow.PrimitiveTypes.Int8, scalar:
scalar.NewInt8Scalar(1)},
+ {name: "int16", typ: arrow.PrimitiveTypes.Int16, scalar:
scalar.NewInt16Scalar(1)},
+ {name: "int32", typ: arrow.PrimitiveTypes.Int32, scalar:
scalar.NewInt32Scalar(1)},
+ {name: "int64", typ: arrow.PrimitiveTypes.Int64, scalar:
scalar.NewInt64Scalar(1)},
+ {name: "uint8", typ: arrow.PrimitiveTypes.Uint8, scalar:
scalar.NewUint8Scalar(1)},
+ {name: "uint16", typ: arrow.PrimitiveTypes.Uint16, scalar:
scalar.NewUint16Scalar(1)},
+ {name: "uint32", typ: arrow.PrimitiveTypes.Uint32, scalar:
scalar.NewUint32Scalar(1)},
+ {name: "uint64", typ: arrow.PrimitiveTypes.Uint64, scalar:
scalar.NewUint64Scalar(1)},
+ }
+
+ for _, listType := range listTypes {
+ t.Run(listType.String(), func(t *testing.T) {
+ input := listElementInput(t, mem, listType, `[[10, 20],
null]`)
+ defer input.Release()
+ expected := listElementInput(t, mem,
arrow.PrimitiveTypes.Int32, `[20, null]`)
+ defer expected.Release()
+
+ for _, tc := range indexes {
+ tc := tc
+ t.Run(tc.name+" scalar", func(t *testing.T) {
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value:
input.Data()},
+ &compute.ScalarDatum{Value:
tc.scalar},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual :=
result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.NoError(t,
array.ValidateFull(actual))
+ assert.True(t, array.Equal(expected,
actual), "expected: %s\ngot: %s", expected, actual)
+ })
+
+ t.Run(tc.name+" array", func(t *testing.T) {
+ singleInput := listElementInput(t, mem,
listType, `[[10, 20]]`)
+ defer singleInput.Release()
+ index := listElementInput(t, mem,
tc.typ, `[1]`)
+ defer index.Release()
+ expectedSingle := listElementInput(t,
mem, arrow.PrimitiveTypes.Int32, `[20]`)
+ defer expectedSingle.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value:
singleInput.Data()},
+ &compute.ArrayDatum{Value:
index.Data()},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual :=
result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.NoError(t,
array.ValidateFull(actual))
+ assert.True(t,
array.Equal(expectedSingle, actual), "expected: %s\ngot: %s", expectedSingle,
actual)
+ })
+ }
+ })
+ }
+}
+
+func TestListElementNumericChildren(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ childTypes := []arrow.DataType{
+ arrow.PrimitiveTypes.Int8,
+ arrow.PrimitiveTypes.Int16,
+ arrow.PrimitiveTypes.Int32,
+ arrow.PrimitiveTypes.Int64,
+ arrow.PrimitiveTypes.Uint8,
+ arrow.PrimitiveTypes.Uint16,
+ arrow.PrimitiveTypes.Uint32,
+ arrow.PrimitiveTypes.Uint64,
+ arrow.FixedWidthTypes.Float16,
+ arrow.PrimitiveTypes.Float32,
+ arrow.PrimitiveTypes.Float64,
+ }
+ for _, childType := range childTypes {
+ t.Run(childType.String(), func(t *testing.T) {
+ listTypes := []arrow.DataType{
+ arrow.ListOf(childType),
+ arrow.LargeListOf(childType),
+ arrow.ListViewOf(childType),
+ arrow.LargeListViewOf(childType),
+ arrow.FixedSizeListOf(2, childType),
+ }
+ for _, listType := range listTypes {
+ t.Run(listType.String(), func(t *testing.T) {
+ input := listElementInput(t, mem,
listType, `[[1, 2], [3, 4], null, [5, 6]]`)
+ defer input.Release()
+ expected := listElementInput(t, mem,
childType, `[2, 4, null, 6]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value:
input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual :=
result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.NoError(t,
array.ValidateFull(actual))
+ assert.True(t, array.Equal(expected,
actual), "expected: %s\ngot: %s", expected, actual)
+ })
+ }
+ })
+ }
+}
+
+func TestListElementPreservesChildNulls(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ input := listElementInput(t, mem,
arrow.ListOf(arrow.PrimitiveTypes.Int32), `[[1, null], [2, 3]]`)
+ defer input.Release()
+ expected := listElementInput(t, mem, arrow.PrimitiveTypes.Int32,
`[null, 3]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt8Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementSlicedInputs(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []arrow.DataType{
+ arrow.ListOf(arrow.PrimitiveTypes.Int32),
+ arrow.FixedSizeListOf(2, arrow.PrimitiveTypes.Int32),
+ }
+ for _, typ := range tests {
+ t.Run(typ.String(), func(t *testing.T) {
+ input := listElementInput(t, mem, typ, `[[0, 1], [2,
3], [4, 5], [6, 7]]`)
+ defer input.Release()
+ sliced := array.NewSlice(input, 1, 3)
+ defer sliced.Release()
+ expected := listElementInput(t, mem,
arrow.PrimitiveTypes.Int32, `[3, 5]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: sliced.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual),
"expected: %s\ngot: %s", expected, actual)
+
+ })
+ }
+}
+
+func TestListElementListViewUsesSizes(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ name string
+ input arrow.Array
+ }{
+ {name: "list view", input:
makeListViewWithOutOfOrderOffsets(mem)},
+ {name: "large list view", input:
makeLargeListViewWithOutOfOrderOffsets(mem)},
+ }
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ defer tc.input.Release()
+ expected := listElementInput(t, mem,
arrow.PrimitiveTypes.Int32, `[14, 10, 12]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: tc.input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual),
"expected: %s\ngot: %s", expected, actual)
+
+ sliced := array.NewSlice(tc.input, 1,
int64(tc.input.Len()))
+ defer sliced.Release()
+ slicedExpected := listElementInput(t, mem,
arrow.PrimitiveTypes.Int32, `[10, 12]`)
+ defer slicedExpected.Release()
+ slicedResult, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: sliced.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer slicedResult.Release()
+ slicedActual :=
slicedResult.(*compute.ArrayDatum).MakeArray()
+ defer slicedActual.Release()
+ assert.True(t, array.Equal(slicedExpected,
slicedActual), "expected: %s\ngot: %s", slicedExpected, slicedActual)
+ })
+ }
+}
+
+func TestListElementRejectsInvalidListViewOffsets(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ name string
+ build func(memory.Allocator) arrow.Array
+ }{
+ {"list view negative offset", func(mem memory.Allocator)
arrow.Array {
+ builder := array.NewListViewBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ builder.ValueBuilder().(*array.Int32Builder).Append(7)
+ builder.AppendDimensions(-1, 1)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+ }},
+ {"list view child overflow", func(mem memory.Allocator)
arrow.Array {
+ builder := array.NewListViewBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ builder.ValueBuilder().(*array.Int32Builder).Append(7)
+ builder.AppendDimensions(1, 1)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+ }},
+ {"large list view negative offset", func(mem memory.Allocator)
arrow.Array {
+ builder := array.NewLargeListViewBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ builder.ValueBuilder().(*array.Int32Builder).Append(7)
+ builder.AppendDimensions(-1, 1)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+ }},
+ {"large list view child overflow", func(mem memory.Allocator)
arrow.Array {
+ builder := array.NewLargeListViewBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ builder.ValueBuilder().(*array.Int32Builder).Append(7)
+ builder.AppendDimensions(1, 1)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+ }},
+ {"large list view offset overflow", func(mem memory.Allocator)
arrow.Array {
+ builder := array.NewLargeListViewBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ builder.AppendDimensions(int(^uint(0)>>1), 1)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+ }},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ input := tc.build(mem)
+ defer input.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+ })
+ }
+}
+
+func TestListElementRejectsFixedSizeListOffsetOverflow(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ valuesBuilder := array.NewInt32Builder(mem)
+ valuesBuilder.AppendValues([]int32{10, 20, 30, 40}, nil)
+ values := valuesBuilder.NewArray()
+ valuesBuilder.Release()
+ defer values.Release()
+
+ const offset = int64(6148914691236517205)
+ typ := arrow.FixedSizeListOf(3, arrow.PrimitiveTypes.Int32)
+ data := array.NewData(typ, 1, []*memory.Buffer{nil},
[]arrow.ArrayData{values.Data()}, 0, int(offset))
+ input := array.NewFixedSizeListData(data)
+ data.Release()
+ defer input.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+}
+
+func makeListViewWithOutOfOrderOffsets(mem memory.Allocator) arrow.Array {
+ builder := array.NewListViewBuilder(mem, arrow.PrimitiveTypes.Int32)
+ values := builder.ValueBuilder().(*array.Int32Builder)
+ values.AppendValues([]int32{10, 11, 12, 13, 14, 15}, nil)
+ builder.AppendDimensions(4, 2)
+ builder.AppendDimensions(0, 1)
+ builder.AppendDimensions(2, 2)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+}
+
+func makeLargeListViewWithOutOfOrderOffsets(mem memory.Allocator) arrow.Array {
+ builder := array.NewLargeListViewBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ values := builder.ValueBuilder().(*array.Int32Builder)
+ values.AppendValues([]int32{10, 11, 12, 13, 14, 15}, nil)
+ builder.AppendDimensions(4, 2)
+ builder.AppendDimensions(0, 1)
+ builder.AppendDimensions(2, 2)
+ result := builder.NewArray()
+ builder.Release()
+ return result
+}
+
+func TestListElementNestedListViewChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ name string
+ build func(memory.Allocator) arrow.Array
+ elemTyp arrow.DataType
+ }{
+ {"list view", func(mem memory.Allocator) arrow.Array {
+ outer := array.NewListBuilder(mem,
arrow.ListViewOf(arrow.PrimitiveTypes.Int32))
+ inner := outer.ValueBuilder().(*array.ListViewBuilder)
+ values := inner.ValueBuilder().(*array.Int32Builder)
+ outer.Append(true)
+ values.AppendValues([]int32{10, 11, 20, 21}, nil)
+ inner.AppendDimensions(0, 2)
+ inner.AppendDimensions(2, 2)
+ result := outer.NewArray()
+ outer.Release()
+ return result
+ }, arrow.ListViewOf(arrow.PrimitiveTypes.Int32)},
+ {"large list view", func(mem memory.Allocator) arrow.Array {
+ outer := array.NewListBuilder(mem,
arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32))
+ inner :=
outer.ValueBuilder().(*array.LargeListViewBuilder)
+ values := inner.ValueBuilder().(*array.Int32Builder)
+ outer.Append(true)
+ values.AppendValues([]int32{10, 11, 20, 21}, nil)
+ inner.AppendDimensions(0, 2)
+ inner.AppendDimensions(2, 2)
+ result := outer.NewArray()
+ outer.Release()
+ return result
+ }, arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32)},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ input := tc.build(mem)
+ defer input.Release()
+ expected := listElementInput(t, mem, tc.elemTyp, `[[20,
21]]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.NoError(t, array.ValidateFull(actual))
+ assert.True(t, array.Equal(expected, actual),
"expected: %s\ngot: %s", expected, actual)
+ })
+ }
+}
+
+func TestListElementEmptyNestedListViews(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for _, elemType := range []arrow.DataType{
+ arrow.ListViewOf(arrow.PrimitiveTypes.Int32),
+ arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32),
+ } {
+ t.Run(elemType.String(), func(t *testing.T) {
+ builder := array.NewListBuilder(mem, elemType)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.Equal(t, 0, actual.Len())
+ require.True(t, arrow.TypeEqual(elemType,
actual.DataType()))
+ require.NoError(t, array.ValidateFull(actual))
+ })
+ }
+}
+
+func TestListElementEmptyRunEndEncodedChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ elemType := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int32)
+ builder := array.NewListBuilder(mem, elemType)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.Equal(t, 0, actual.Len())
+ require.True(t, arrow.TypeEqual(elemType, actual.DataType()))
+ require.NoError(t, array.ValidateFull(actual))
+}
+
+func TestListElementNullParentEmptyListViews(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for _, elemType := range []arrow.DataType{
+ arrow.ListViewOf(arrow.PrimitiveTypes.Int32),
+ arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32),
+ } {
+ t.Run(elemType.String(), func(t *testing.T) {
+ builder := array.NewListBuilder(mem, elemType)
+ builder.AppendNull()
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.Equal(t, 1, actual.Len())
+ require.Equal(t, 1, actual.NullN())
+ require.True(t, actual.IsNull(0))
+ require.True(t, arrow.TypeEqual(elemType,
actual.DataType()))
+ require.NoError(t, array.ValidateFull(actual))
+ })
+ }
+}
+
+func TestListElementScalarListWithViewElement(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for _, values := range []arrow.Array{
+ makeListViewWithOutOfOrderOffsets(mem),
+ makeLargeListViewWithOutOfOrderOffsets(mem),
+ } {
+ t.Run(values.DataType().String(), func(t *testing.T) {
+ defer values.Release()
+ lists := scalar.NewListScalar(values)
+ defer lists.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: lists},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ })
+ }
+}
+
+func makeListElementViewValues(mem memory.Allocator, binary bool) arrow.Array {
+ values := []string{strings.Repeat("a", 32), strings.Repeat("b", 32)}
+ if binary {
+ builder := array.NewBinaryViewBuilder(mem)
+ builder.SetBlockSize(1)
+ for _, value := range values {
+ builder.Append([]byte(value))
+ }
+ result := builder.NewArray()
+ builder.Release()
+ return result
+ }
+
+ builder := array.NewStringViewBuilder(mem)
+ builder.SetBlockSize(1)
+ for _, value := range values {
+ builder.Append(value)
+ }
+ result := builder.NewArray()
+ builder.Release()
+ return result
+}
+
+func makeListElementArrayWithChild(mem memory.Allocator, elemType
arrow.DataType, child arrow.Array) arrow.Array {
+ offsetsBuilder := array.NewInt32Builder(mem)
+ offsetsBuilder.AppendValues([]int32{0, 1, 2}, nil)
+ offsets := offsetsBuilder.NewArray()
+ offsetsBuilder.Release()
+
+ data := array.NewData(
+ arrow.ListOf(elemType),
+ 2,
+ []*memory.Buffer{nil, offsets.Data().Buffers()[1]},
+ []arrow.ArrayData{child.Data()},
+ 0,
+ 0,
+ )
+ result := array.NewListData(data)
+ data.Release()
+ offsets.Release()
+ child.Release()
+ return result
+}
+
+func TestListElementRejectsNestedViewChildren(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ viewType := arrow.BinaryTypes.StringView
+ tests := []struct {
+ name string
+ build func() arrow.Array
+ }{
+ {
+ name: "string_view",
+ build: func() arrow.Array {
+ return makeListElementArrayWithChild(mem,
viewType, makeListElementViewValues(mem, false))
+ },
+ },
+ {
+ name: "binary_view",
+ build: func() arrow.Array {
+ return makeListElementArrayWithChild(mem,
arrow.BinaryTypes.BinaryView, makeListElementViewValues(mem, true))
+ },
+ },
+ {
+ name: "struct_string_view",
+ build: func() arrow.Array {
+ typ := arrow.StructOf(arrow.Field{Name:
"value", Type: viewType, Nullable: true})
+ builder := array.NewStructBuilder(mem, typ)
+ values :=
builder.FieldBuilder(0).(*array.StringViewBuilder)
+ values.SetBlockSize(1)
+ for _, value := range
[]string{strings.Repeat("a", 32), strings.Repeat("b", 32)} {
+ builder.Append(true)
+ values.Append(value)
+ }
+ child := builder.NewArray()
+ builder.Release()
+ return makeListElementArrayWithChild(mem, typ,
child)
+ },
+ },
+ {
+ name: "list_string_view",
+ build: func() arrow.Array {
+ typ := arrow.ListOf(viewType)
+ builder := array.NewListBuilder(mem, viewType)
+ values :=
builder.ValueBuilder().(*array.StringViewBuilder)
+ values.SetBlockSize(1)
+ for _, value := range
[]string{strings.Repeat("a", 32), strings.Repeat("b", 32)} {
+ builder.Append(true)
+ values.Append(value)
+ }
+ child := builder.NewArray()
+ builder.Release()
+ return makeListElementArrayWithChild(mem, typ,
child)
+ },
+ },
+ {
+ name: "fixed_size_list_string_view",
+ build: func() arrow.Array {
+ typ := arrow.FixedSizeListOf(2, viewType)
+ builder := array.NewFixedSizeListBuilder(mem,
2, viewType)
+ values :=
builder.ValueBuilder().(*array.StringViewBuilder)
+ values.SetBlockSize(1)
+ for _, value := range
[]string{strings.Repeat("a", 32), strings.Repeat("b", 32)} {
+ builder.Append(true)
+ values.Append(value)
+ values.Append(value)
+ }
+ child := builder.NewArray()
+ builder.Release()
+ return makeListElementArrayWithChild(mem, typ,
child)
+ },
+ },
+ {
+ name: "dictionary_string_view",
+ build: func() arrow.Array {
+ typ := &arrow.DictionaryType{IndexType:
arrow.PrimitiveTypes.Int8, ValueType: viewType}
+ values := makeListElementViewValues(mem, false)
+ indicesBuilder := array.NewInt8Builder(mem)
+ indicesBuilder.AppendValues([]int8{0, 1}, nil)
+ indices := indicesBuilder.NewArray()
+ indicesBuilder.Release()
+ child := array.NewDictionaryArray(typ, indices,
values)
+ indices.Release()
+ values.Release()
+ return makeListElementArrayWithChild(mem, typ,
child)
+ },
+ },
+ {
+ name: "run_end_encoded_string_view",
+ build: func() arrow.Array {
+ typ :=
arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32, viewType)
+ builder := array.NewRunEndEncodedBuilder(mem,
typ.RunEnds(), typ.Encoded())
+ values :=
builder.ValueBuilder().(*array.StringViewBuilder)
+ values.SetBlockSize(1)
+ for _, value := range
[]string{strings.Repeat("a", 32), strings.Repeat("b", 32)} {
+ builder.Append(1)
+ values.Append(value)
+ }
+ child := builder.NewArray()
+ builder.Release()
+ return makeListElementArrayWithChild(mem, typ,
child)
+ },
+ },
+ {
+ name: "extension_string_view",
+ build: func() arrow.Array {
+ typ := &denseUnionExtensionType{ExtensionBase:
arrow.ExtensionBase{Storage: viewType}}
+ storage := makeListElementViewValues(mem, false)
+ child :=
array.NewExtensionArrayWithStorage(typ, storage)
+ storage.Release()
+ return makeListElementArrayWithChild(mem, typ,
child)
+ },
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ input := tc.build()
+ defer input.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ })
+ }
+}
+
+func TestListElementRejectsDictionaryBackedExtensionChildren(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8,
ValueType: arrow.BinaryTypes.String}
+ extType := &denseUnionExtensionType{ExtensionBase:
arrow.ExtensionBase{Storage: dictType}}
+ valuesBuilder := array.NewStringBuilder(mem)
+ valuesBuilder.AppendValues([]string{"a", "b"}, nil)
+ values := valuesBuilder.NewArray()
+ valuesBuilder.Release()
+ defer values.Release()
+ indicesBuilder := array.NewInt8Builder(mem)
+ indicesBuilder.AppendValues([]int8{0, 1}, nil)
+ indices := indicesBuilder.NewArray()
+ indicesBuilder.Release()
+ defer indices.Release()
+
+ data := array.NewDataWithDictionary(extType, 2,
indices.Data().Buffers(), 0, 0, values.Data().(*array.Data))
+ child := array.NewExtensionData(data)
+ data.Release()
+ defer child.Release()
+ require.NoError(t, array.ValidateFull(child.Storage()))
+ nested, err := array.NewStructArray([]arrow.Array{child},
[]string{"value"})
+ require.NoError(t, err)
+ defer nested.Release()
+
+ for _, child := range []arrow.Array{child, nested} {
+ t.Run(child.DataType().String(), func(t *testing.T) {
+ child.Retain()
+ input := makeListElementArrayWithChild(mem,
child.DataType(), child)
+ defer input.Release()
+ empty := array.NewSlice(input, 0, 0)
+ defer empty.Release()
+ chunked := arrow.NewChunked(input.DataType(),
[]arrow.Array{input})
+ defer chunked.Release()
+ listScalar := scalar.NewListScalar(child)
+ defer listScalar.Release()
+ nullScalar :=
scalar.MakeNullScalar(input.DataType()).(scalar.ListScalar)
+ defer nullScalar.Release()
+
+ for _, input := range []compute.Datum{
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ArrayDatum{Value: empty.Data()},
+ &compute.ChunkedDatum{Value: chunked},
+ &compute.ScalarDatum{Value: listScalar},
+ &compute.ScalarDatum{Value: nullScalar},
+ } {
+ result, err := compute.ListElement(
+ context.Background(), input,
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ }
+ })
+ }
+}
+
+func TestListElementDecimalArrayChildren(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ name string
+ typ arrow.DataType
+ add func(array.Builder, int32)
+ }{
+ {
+ name: "decimal32",
+ typ: &arrow.Decimal32Type{Precision: 6, Scale: 2},
+ add: func(builder array.Builder, value int32) {
+
builder.(*array.Decimal32Builder).Append(decimal.Decimal32(value))
+ },
+ },
+ {
+ name: "decimal64",
+ typ: &arrow.Decimal64Type{Precision: 12, Scale: 2},
+ add: func(builder array.Builder, value int32) {
+
builder.(*array.Decimal64Builder).Append(decimal.Decimal64(value))
+ },
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ listBuilder := array.NewListBuilder(mem, tc.typ)
+ values := listBuilder.ValueBuilder()
+ for _, row := range [][]int32{{1, 2}, {3, 4}} {
+ listBuilder.Append(true)
+ for _, value := range row {
+ tc.add(values, value)
+ }
+ }
+ input := listBuilder.NewArray()
+ listBuilder.Release()
+ defer input.Release()
+
+ expectedBuilder := array.NewBuilder(mem, tc.typ)
+ tc.add(expectedBuilder, 2)
+ tc.add(expectedBuilder, 4)
+ expected := expectedBuilder.NewArray()
+ expectedBuilder.Release()
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.NoError(t, array.ValidateFull(actual))
+ assert.True(t, array.Equal(expected, actual),
"expected: %s\ngot: %s", expected, actual)
+ })
+ }
+}
+
+func TestListElementScalarListWithUnsupportedDecimalValues(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ name string
+ typ arrow.DataType
+ }{
+ {name: "decimal32", typ: &arrow.Decimal32Type{Precision: 6,
Scale: 2}},
+ {name: "decimal64", typ: &arrow.Decimal64Type{Precision: 12,
Scale: 2}},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ builder := array.NewBuilder(mem, tc.typ)
+ switch b := builder.(type) {
+ case *array.Decimal32Builder:
+ b.Append(decimal.Decimal32(123))
+ case *array.Decimal64Builder:
+ b.Append(decimal.Decimal64(123))
+ default:
+ t.Fatalf("unexpected builder type %T", builder)
+ }
+ values := builder.NewArray()
+ builder.Release()
+ defer values.Release()
+
+ list := scalar.NewListScalar(values)
+ defer list.Release()
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: list},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ })
+ }
+}
+
+func TestListElementSingleIndexArray(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ input := listElementInput(t, mem,
arrow.ListOf(arrow.PrimitiveTypes.Int32), `[[1, 2]]`)
+ defer input.Release()
+ index := listElementInput(t, mem, arrow.PrimitiveTypes.Int64, `[1]`)
+ defer index.Release()
+ expected := listElementInput(t, mem, arrow.PrimitiveTypes.Int32, `[2]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ArrayDatum{Value: index.Data()},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementSingleIndexArrayRejectsMismatchedLengths(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ input := listElementInput(t, mem,
arrow.ListOf(arrow.PrimitiveTypes.Int32), `[[1, 2], [3, 4]]`)
+ defer input.Release()
+ index := listElementInput(t, mem, arrow.PrimitiveTypes.Int64, `[1]`)
+ defer index.Release()
+ result, err := compute.CallFunction(
+ context.Background(),
+ "list_element",
+ nil,
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ArrayDatum{Value: index.Data()},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+}
+
+func TestListElementDispatchBest(t *testing.T) {
+ listType := arrow.ListOf(arrow.PrimitiveTypes.Int32)
+ CheckDispatchBest(t, "list_element",
+ []arrow.DataType{listType, arrow.PrimitiveTypes.Int64},
+ []arrow.DataType{listType, arrow.PrimitiveTypes.Int64})
+}
+
+func TestListElementFunctionDoc(t *testing.T) {
+ fn, ok := compute.GetFunctionRegistry().GetFunction("list_element")
+ require.True(t, ok)
+ require.NoError(t, fn.Validate())
+}
+
+func TestListElementRejectsMultipleIndicesIndependentOfExecutionSpans(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ input := listElementInput(t, mem,
arrow.ListOf(arrow.PrimitiveTypes.Int32), `[[10, 11], [20, 21]]`)
+ defer input.Release()
+ index := listElementInput(t, mem, arrow.PrimitiveTypes.Int64, `[0, 1]`)
+ defer index.Release()
+
+ chunk0 := array.NewSlice(input, 0, 1)
+ defer chunk0.Release()
+ chunk1 := array.NewSlice(input, 1, 2)
+ defer chunk1.Release()
+ chunkedLists := arrow.NewChunked(input.DataType(),
[]arrow.Array{chunk0, chunk1})
+ defer chunkedLists.Release()
+
+ execCtx := compute.DefaultExecCtx()
+ execCtx.ChunkSize = 1
+
+ tests := []struct {
+ name string
+ ctx context.Context
+ lists compute.Datum
+ }{
+ {name: "regular execution span", ctx: context.Background(),
lists: &compute.ArrayDatum{Value: input.Data()}},
+ {name: "chunk size one", ctx:
compute.SetExecCtx(context.Background(), execCtx), lists:
&compute.ArrayDatum{Value: input.Data()}},
+ {name: "chunked lists", ctx: context.Background(), lists:
&compute.ChunkedDatum{Value: chunkedLists}},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ result, err := compute.ListElement(
+ tc.ctx,
+ tc.lists,
+ &compute.ArrayDatum{Value: index.Data()},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ })
+ }
+}
+
+func TestListElementSingleIndexChunkedArray(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ input := listElementInput(t, mem,
arrow.ListOf(arrow.PrimitiveTypes.Int32), `[[1, 2]]`)
+ defer input.Release()
+ index := listElementInput(t, mem, arrow.PrimitiveTypes.Int64, `[1]`)
+ defer index.Release()
+ chunkedIndex := arrow.NewChunked(index.DataType(), []arrow.Array{index})
+ defer chunkedIndex.Release()
+ expected := listElementInput(t, mem, arrow.PrimitiveTypes.Int32, `[2]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ChunkedDatum{Value: chunkedIndex},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ chunkedResult, ok := result.(*compute.ChunkedDatum)
+ require.True(t, ok)
+ require.Len(t, chunkedResult.Value.Chunks(), 1)
+ assert.True(t, array.Equal(expected, chunkedResult.Value.Chunk(0)),
+ "expected: %s\ngot: %s", expected, chunkedResult.Value.Chunk(0))
+}
+
+func TestListElementScalarList(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ values := listElementInput(t, mem, arrow.PrimitiveTypes.Int32, `[10,
20]`)
+ defer values.Release()
+ list := scalar.NewListScalar(values)
+ defer list.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: list},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ScalarDatum).Value
+ assert.True(t, scalar.Equals(scalar.NewInt32Scalar(20), actual),
"expected: 20\ngot: %s", actual)
+}
+
+func TestListElementRejectsNonListScalar(t *testing.T) {
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: scalar.NewInt32Scalar(7)},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrType)
+}
+
+func TestListElementRejectsNilScalarIndex(t *testing.T) {
+ result, err := compute.ListElement(
+ context.Background(),
+ compute.EmptyDatum{},
+ &compute.ScalarDatum{},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrType)
+}
+
+func TestListElementScalarListWithArrayIndex(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ values := listElementInput(t, mem, arrow.PrimitiveTypes.Int32, `[10,
20]`)
+ defer values.Release()
+ list := scalar.NewListScalar(values)
+ defer list.Release()
+ index := listElementInput(t, mem, arrow.PrimitiveTypes.Int64, `[1]`)
+ defer index.Release()
+ expected := listElementInput(t, mem, arrow.PrimitiveTypes.Int32, `[20]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: list},
+ &compute.ArrayDatum{Value: index.Data()},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ arrayResult, ok := result.(*compute.ArrayDatum)
+ require.True(t, ok)
+ actual := arrayResult.MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementErrors(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ name string
+ input string
+ index scalar.Scalar
+ }{
+ {name: "out of bounds", input: `[[1], [2, 3]]`, index:
scalar.NewInt64Scalar(1)},
+ {name: "empty list", input: `[[], [1]]`, index:
scalar.NewInt64Scalar(0)},
+ {name: "negative index", input: `[[1]]`, index:
scalar.NewInt64Scalar(-1)},
+ {name: "large unsigned index", input: `[[1]]`, index:
scalar.NewUint64Scalar(^uint64(0))},
+ {name: "null index", input: `[[1]]`, index:
scalar.MakeNullScalar(arrow.PrimitiveTypes.Int64)},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ input := listElementInput(t, mem,
arrow.ListOf(arrow.PrimitiveTypes.Int32), tc.input)
+ defer input.Release()
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: tc.index},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ assert.ErrorIs(t, err, arrow.ErrInvalid)
+ })
+ }
+}
+
+func TestListElementComplexChildren(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ name string
+ typ arrow.DataType
+ input string
+ expected string
+ }{
+ {name: "string", typ: arrow.ListOf(arrow.BinaryTypes.String),
input: `[["a", "b"], ["c", "d"]]`, expected: `["b", "d"]`},
+ {name: "nested list", typ:
arrow.ListOf(arrow.ListOf(arrow.PrimitiveTypes.Int32)), input: `[[[1, 2], [3]],
[[4], [5, 6]]]`, expected: `[[3], [5, 6]]`},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ input := listElementInput(t, mem, tc.typ, tc.input)
+ defer input.Release()
+ expected := listElementInput(t, mem,
tc.typ.(arrow.ListLikeType).Elem(), tc.expected)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewUint8Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual),
"expected: %s\ngot: %s", expected, actual)
+ })
+ }
+}
+
+func TestListElementMonthDayNanoIntervalChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ typ := arrow.ListOf(arrow.FixedWidthTypes.MonthDayNanoInterval)
+ input := listElementInput(t, mem, typ,
+ `[[{"months": 1, "days": 2, "nanoseconds": 3}, {"months": 4,
"days": 5, "nanoseconds": 6}], [{"months": 7, "days": 8, "nanoseconds": 9},
{"months": 10, "days": 11, "nanoseconds": 12}]]`)
+ defer input.Release()
+ expected := listElementInput(t, mem,
arrow.FixedWidthTypes.MonthDayNanoInterval,
+ `[{"months": 4, "days": 5, "nanoseconds": 6}, {"months": 10,
"days": 11, "nanoseconds": 12}]`)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementDenseUnionChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ {Name: "text", Type: arrow.BinaryTypes.String,
Nullable: true},
+ },
+ []arrow.UnionTypeCode{0, 1},
+ )
+
+ builder := array.NewListBuilder(mem, unionType)
+ values := builder.ValueBuilder().(*array.DenseUnionBuilder)
+ builder.Append(true)
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).Append(10)
+ values.Append(1)
+ values.Child(1).(*array.StringBuilder).Append("a")
+ builder.Append(true)
+ values.Append(1)
+ values.Child(1).(*array.StringBuilder).Append("b")
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).Append(20)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ expectedBuilder := array.NewDenseUnionBuilder(mem, unionType)
+ expectedBuilder.Append(1)
+ expectedBuilder.Child(1).(*array.StringBuilder).Append("a")
+ expectedBuilder.Append(0)
+ expectedBuilder.Child(0).(*array.Int32Builder).Append(20)
+ expected := expectedBuilder.NewArray()
+ expectedBuilder.Release()
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementDenseUnionExtensionChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ {Name: "text", Type: arrow.BinaryTypes.String,
Nullable: true},
+ },
+ []arrow.UnionTypeCode{0, 1},
+ )
+ extType := &denseUnionExtensionType{ExtensionBase:
arrow.ExtensionBase{Storage: unionType}}
+
+ builder := array.NewListBuilder(mem, extType)
+ values :=
builder.ValueBuilder().(*array.ExtensionBuilder).StorageBuilder().(*array.DenseUnionBuilder)
+ builder.Append(true)
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).Append(10)
+ values.Append(1)
+ values.Child(1).(*array.StringBuilder).Append("a")
+ builder.Append(true)
+ values.Append(1)
+ values.Child(1).(*array.StringBuilder).Append("b")
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).Append(20)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ expectedBuilder := array.NewDenseUnionBuilder(mem, unionType)
+ expectedBuilder.Append(1)
+ expectedBuilder.Child(1).(*array.StringBuilder).Append("a")
+ expectedBuilder.Append(0)
+ expectedBuilder.Child(0).(*array.Int32Builder).Append(20)
+ expectedStorage := expectedBuilder.NewArray()
+ expectedBuilder.Release()
+ expected := array.NewExtensionArrayWithStorage(extType, expectedStorage)
+ expectedStorage.Release()
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.True(t, arrow.TypeEqual(extType, actual.DataType()))
+ storage := actual.(array.ExtensionArray).Storage()
+ require.NoError(t, array.ValidateFull(storage))
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementExtensionRunEndEncodedNullParent(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ storageType := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int32)
+ extType := &runEndExtensionType{ExtensionBase:
arrow.ExtensionBase{Storage: storageType}}
+
+ builder := array.NewListBuilder(mem, extType)
+ builder.AppendNull()
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.True(t, arrow.TypeEqual(extType, actual.DataType()))
+ require.NoError(t, array.ValidateFull(actual))
+ storage := actual.(array.ExtensionArray).Storage()
+ require.True(t, storage.(*array.RunEndEncoded).Values().IsNull(0))
+ require.NoError(t, array.ValidateFull(storage))
+}
+
+func TestListElementDenseUnionWithUnusedGenericChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ {Name: "values", Type:
arrow.ListOf(arrow.PrimitiveTypes.Int32), Nullable: true},
+ },
+ []arrow.UnionTypeCode{0, 1},
+ )
+
+ builder := array.NewListBuilder(mem, unionType)
+ values := builder.ValueBuilder().(*array.DenseUnionBuilder)
+ builder.Append(true)
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).Append(10)
+ builder.Append(true)
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).Append(20)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ expectedBuilder := array.NewDenseUnionBuilder(mem, unionType)
+ expectedBuilder.Append(0)
+ expectedBuilder.Child(0).(*array.Int32Builder).Append(10)
+ expectedBuilder.Append(0)
+ expectedBuilder.Child(0).(*array.Int32Builder).Append(20)
+ expected := expectedBuilder.NewArray()
+ expectedBuilder.Release()
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementDenseUnionScalarWithUnusedUnsupportedChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ {Name: "values", Type:
arrow.ListViewOf(arrow.PrimitiveTypes.Int32), Nullable: true},
+ },
+ []arrow.UnionTypeCode{0, 1},
+ )
+
+ builder := array.NewListBuilder(mem, unionType)
+ values := builder.ValueBuilder().(*array.DenseUnionBuilder)
+ builder.Append(true)
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).Append(42)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ listValue, err := scalar.GetScalar(input, 0)
+ require.NoError(t, err)
+ defer listValue.(scalar.Releasable).Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: listValue},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ScalarDatum).Value
+ require.Equal(t, arrow.DENSE_UNION, actual.DataType().ID())
+ assert.Equal(t, int32(42),
actual.(scalar.Union).ChildValue().(*scalar.Int32).Value)
+}
+
+func TestListElementDenseUnionScalarWithActiveUnsupportedChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{{Name: "values", Type:
arrow.ListViewOf(arrow.PrimitiveTypes.Int32), Nullable: true}},
+ []arrow.UnionTypeCode{0},
+ )
+ builder := array.NewListBuilder(mem, unionType)
+ values := builder.ValueBuilder().(*array.DenseUnionBuilder)
+ listBuilder := values.Child(0).(*array.ListViewBuilder)
+ builder.Append(true)
+ values.Append(0)
+ listBuilder.ValueBuilder().(*array.Int32Builder).Append(9)
+ listBuilder.AppendDimensions(0, 1)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ listValue := scalar.NewListScalar(input.(*array.List).ListValues())
+ defer listValue.Release()
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: listValue},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+}
+
+func TestListElementDenseUnionScalarWithActiveUnsupportedDecimalChild(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for _, tc := range []struct {
+ name string
+ typ arrow.DataType
+ add func(array.Builder)
+ }{
+ {
+ name: "decimal32",
+ typ: &arrow.Decimal32Type{Precision: 6, Scale: 2},
+ add: func(builder array.Builder) {
+
builder.(*array.Decimal32Builder).Append(decimal.Decimal32(123))
+ },
+ },
+ {
+ name: "decimal64",
+ typ: &arrow.Decimal64Type{Precision: 12, Scale: 2},
+ add: func(builder array.Builder) {
+
builder.(*array.Decimal64Builder).Append(decimal.Decimal64(123))
+ },
+ },
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{{Name: "value", Type: tc.typ,
Nullable: true}},
+ []arrow.UnionTypeCode{0},
+ )
+ builder := array.NewListBuilder(mem, unionType)
+ values :=
builder.ValueBuilder().(*array.DenseUnionBuilder)
+ builder.Append(true)
+ values.Append(0)
+ tc.add(values.Child(0))
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ listValue :=
scalar.NewListScalar(input.(*array.List).ListValues())
+ defer listValue.Release()
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: listValue},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ })
+ }
+}
+
+func TestListElementDenseUnionScalarWithActiveNullUnsupportedChild(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for _, tc := range []struct {
+ name string
+ typ arrow.DataType
+ }{
+ {
+ name: "binary_view",
+ typ: arrow.BinaryTypes.BinaryView,
+ },
+ {
+ name: "string_view",
+ typ: arrow.BinaryTypes.StringView,
+ },
+ {
+ name: "decimal32",
+ typ: &arrow.Decimal32Type{Precision: 6, Scale: 2},
+ },
+ {
+ name: "decimal64",
+ typ: &arrow.Decimal64Type{Precision: 12, Scale: 2},
+ },
+ {
+ name: "list_view",
+ typ: arrow.ListViewOf(arrow.PrimitiveTypes.Int32),
+ },
+ {
+ name: "large_list_view",
+ typ: arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32),
+ },
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{{Name: "value", Type: tc.typ,
Nullable: true}},
+ []arrow.UnionTypeCode{0},
+ )
+ builder := array.NewListBuilder(mem, unionType)
+ values :=
builder.ValueBuilder().(*array.DenseUnionBuilder)
+ builder.Append(true)
+ values.Append(0)
+ values.Child(0).AppendNull()
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ listValue :=
scalar.NewListScalar(input.(*array.List).ListValues())
+ defer listValue.Release()
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: listValue},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ })
+ }
+}
+
+func TestListElementDenseUnionExtensionScalarWithUnusedUnsupportedChild(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ {Name: "values", Type:
arrow.ListViewOf(arrow.PrimitiveTypes.Int32), Nullable: true},
+ },
+ []arrow.UnionTypeCode{0, 1},
+ )
+ extType := &denseUnionExtensionType{ExtensionBase:
arrow.ExtensionBase{Storage: unionType}}
+
+ builder := array.NewListBuilder(mem, extType)
+ values :=
builder.ValueBuilder().(*array.ExtensionBuilder).StorageBuilder().(*array.DenseUnionBuilder)
+ builder.Append(true)
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).Append(42)
+ values.Append(0)
+ values.Child(0).(*array.Int32Builder).AppendNull()
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ listValue, err := scalar.GetScalar(input, 0)
+ require.NoError(t, err)
+ defer listValue.(scalar.Releasable).Release()
+
+ for i := int64(0); i < 2; i++ {
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ScalarDatum{Value: listValue},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(i)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ScalarDatum).Value
+ require.True(t, arrow.TypeEqual(extType, actual.DataType()))
+ assert.Equal(t, i == 0, actual.IsValid())
+ assert.NoError(t, actual.ValidateFull())
+ if i == 0 {
+ assert.Equal(t, int32(42),
actual.(*scalar.Extension).Value.(scalar.Union).ChildValue().(*scalar.Int32).Value)
+ }
+ }
+}
+
+func TestListElementDenseUnionMonthDayNanoIntervalChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ unionType := arrow.DenseUnionOf(
+ []arrow.Field{
+ {Name: "interval", Type:
arrow.FixedWidthTypes.MonthDayNanoInterval, Nullable: true},
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ },
+ []arrow.UnionTypeCode{0, 1},
+ )
+
+ builder := array.NewListBuilder(mem, unionType)
+ values := builder.ValueBuilder().(*array.DenseUnionBuilder)
+ builder.Append(true)
+ values.Append(0)
+
values.Child(0).(*array.MonthDayNanoIntervalBuilder).Append(arrow.MonthDayNanoInterval{Months:
1, Days: 2, Nanoseconds: 3})
+ builder.Append(true)
+ values.Append(0)
+
values.Child(0).(*array.MonthDayNanoIntervalBuilder).Append(arrow.MonthDayNanoInterval{Months:
4, Days: 5, Nanoseconds: 6})
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ expectedBuilder := array.NewDenseUnionBuilder(mem, unionType)
+ expectedBuilder.Append(0)
+
expectedBuilder.Child(0).(*array.MonthDayNanoIntervalBuilder).Append(arrow.MonthDayNanoInterval{Months:
1, Days: 2, Nanoseconds: 3})
+ expectedBuilder.Append(0)
+
expectedBuilder.Child(0).(*array.MonthDayNanoIntervalBuilder).Append(arrow.MonthDayNanoInterval{Months:
4, Days: 5, Nanoseconds: 6})
+ expected := expectedBuilder.NewArray()
+ expectedBuilder.Release()
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementSparseUnionChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ unionType := arrow.SparseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ {Name: "text", Type: arrow.BinaryTypes.String,
Nullable: true},
+ },
+ []arrow.UnionTypeCode{0, 1},
+ )
+
+ builder := array.NewListBuilder(mem, unionType)
+ values := builder.ValueBuilder().(*array.SparseUnionBuilder)
+ appendValue := func(code arrow.UnionTypeCode, number int32, text
string) {
+ values.Append(code)
+ if code == 0 {
+ values.Child(0).(*array.Int32Builder).Append(number)
+ values.Child(1).(*array.StringBuilder).AppendNull()
+ } else {
+ values.Child(0).(*array.Int32Builder).AppendNull()
+ values.Child(1).(*array.StringBuilder).Append(text)
+ }
+ }
+ builder.Append(true)
+ appendValue(0, 10, "")
+ appendValue(1, 0, "a")
+ builder.Append(true)
+ appendValue(1, 0, "b")
+ appendValue(0, 20, "")
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ expectedBuilder := array.NewSparseUnionBuilder(mem, unionType)
+ expectedBuilder.Append(1)
+ expectedBuilder.Child(0).(*array.Int32Builder).AppendNull()
+ expectedBuilder.Child(1).(*array.StringBuilder).Append("a")
+ expectedBuilder.Append(0)
+ expectedBuilder.Child(0).(*array.Int32Builder).Append(20)
+ expectedBuilder.Child(1).(*array.StringBuilder).AppendNull()
+ expected := expectedBuilder.NewArray()
+ expectedBuilder.Release()
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+}
+
+func TestListElementSparseUnionSlicedValuesChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ typeIDsBuilder := array.NewInt8Builder(mem)
+ typeIDsBuilder.AppendValues([]int8{0, 1, 0}, nil)
+ typeIDs := typeIDsBuilder.NewArray()
+ typeIDsBuilder.Release()
+ defer typeIDs.Release()
+
+ numbersBuilder := array.NewInt32Builder(mem)
+ numbersBuilder.Append(10)
+ numbersBuilder.AppendNull()
+ numbersBuilder.Append(20)
+ numbers := numbersBuilder.NewArray()
+ numbersBuilder.Release()
+ defer numbers.Release()
+
+ textBuilder := array.NewStringBuilder(mem)
+ textBuilder.AppendNull()
+ textBuilder.Append("a")
+ textBuilder.AppendNull()
+ texts := textBuilder.NewArray()
+ textBuilder.Release()
+ defer texts.Release()
+
+ union, err := array.NewSparseUnionFromArraysWithFieldCodes(
+ typeIDs,
+ []arrow.Array{numbers, texts},
+ []string{"number", "text"},
+ []arrow.UnionTypeCode{0, 1},
+ )
+ require.NoError(t, err)
+ defer union.Release()
+
+ slicedUnion := array.NewSlice(union, 1, 3)
+ defer slicedUnion.Release()
+
+ offsetsBuilder := array.NewInt32Builder(mem)
+ offsetsBuilder.AppendValues([]int32{0, 2}, nil)
+ offsets := offsetsBuilder.NewArray()
+ offsetsBuilder.Release()
+ defer offsets.Release()
+
+ data := array.NewData(
+ arrow.ListOf(union.DataType()),
+ 1,
+ []*memory.Buffer{nil, offsets.Data().Buffers()[1]},
+ []arrow.ArrayData{slicedUnion.Data()},
+ 0,
+ 0,
+ )
+ input := array.NewListData(data)
+ data.Release()
+ defer input.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray().(*array.SparseUnion)
+ defer actual.Release()
+ require.NoError(t, array.ValidateFull(actual))
+ require.Equal(t, []arrow.UnionTypeCode{1}, actual.RawTypeCodes())
+ require.True(t, actual.Field(0).IsNull(0))
+ require.Equal(t, "a", actual.Field(1).(*array.String).Value(0))
+}
+
+func TestListElementSparseUnionPreservesUnusualChildren(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictionaryType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.BinaryTypes.String,
+ }
+ encodedType := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32,
arrow.BinaryTypes.String)
+ unionType := arrow.SparseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ {Name: "dictionary", Type: dictionaryType, Nullable:
true},
+ {Name: "encoded", Type: encodedType, Nullable: true},
+ },
+ []arrow.UnionTypeCode{5, 42, 100},
+ )
+
+ builder := array.NewListBuilder(mem, unionType)
+ values := builder.ValueBuilder().(*array.SparseUnionBuilder)
+ numbers := values.Child(0).(*array.Int32Builder)
+ dictionaries := values.Child(1).(*array.BinaryDictionaryBuilder)
+ encoded := values.Child(2).(*array.RunEndEncodedBuilder)
+ encodedValues := encoded.ValueBuilder().(*array.StringBuilder)
+ appendValue := func(code arrow.UnionTypeCode, number int32, dictionary,
encodedValue string) {
+ values.Append(code)
+ if code == 5 {
+ numbers.Append(number)
+ } else {
+ numbers.AppendNull()
+ }
+ if code == 42 {
+ require.NoError(t,
dictionaries.AppendString(dictionary))
+ } else {
+ dictionaries.AppendNull()
+ }
+ encoded.Append(1)
+ if encodedValue == "" {
+ encodedValues.AppendNull()
+ } else {
+ encodedValues.Append(encodedValue)
+ }
+ }
+
+ builder.Append(true)
+ appendValue(5, 10, "", "")
+ appendValue(42, 0, "a", "")
+ builder.Append(true)
+ appendValue(100, 0, "", "inactive")
+ appendValue(5, 20, "", "")
+ builder.Append(false)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ expectedBuilder := array.NewSparseUnionBuilder(mem, unionType)
+ expectedNumbers := expectedBuilder.Child(0).(*array.Int32Builder)
+ expectedDictionaries :=
expectedBuilder.Child(1).(*array.BinaryDictionaryBuilder)
+ expectedEncoded :=
expectedBuilder.Child(2).(*array.RunEndEncodedBuilder)
+ expectedEncodedValues :=
expectedEncoded.ValueBuilder().(*array.StringBuilder)
+ expectedBuilder.Append(42)
+ expectedNumbers.AppendNull()
+ require.NoError(t, expectedDictionaries.AppendString("a"))
+ expectedEncoded.Append(1)
+ expectedEncodedValues.AppendNull()
+ expectedBuilder.Append(5)
+ expectedNumbers.Append(20)
+ expectedDictionaries.AppendNull()
+ expectedEncoded.Append(1)
+ expectedEncodedValues.AppendNull()
+ expectedBuilder.Append(5)
+ expectedNumbers.AppendNull()
+ expectedDictionaries.AppendNull()
+ expectedEncoded.Append(1)
+ expectedEncodedValues.AppendNull()
+ expected := expectedBuilder.NewArray()
+ expectedBuilder.Release()
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.NoError(t, array.ValidateFull(actual))
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+
+ sliced := array.NewSlice(input, 1, 3)
+ defer sliced.Release()
+ slicedResult, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: sliced.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(1)},
+ )
+ require.NoError(t, err)
+ defer slicedResult.Release()
+ slicedActual := slicedResult.(*compute.ArrayDatum).MakeArray()
+ defer slicedActual.Release()
+ require.NoError(t, array.ValidateFull(slicedActual))
+ slicedUnion := slicedActual.(*array.SparseUnion)
+ assert.Equal(t, []arrow.UnionTypeCode{5, 5}, slicedUnion.RawTypeCodes())
+ assert.Equal(t, int32(20), slicedUnion.Field(0).(*array.Int32).Value(0))
+ assert.True(t, slicedUnion.Field(0).IsNull(1))
+}
+
+func TestListElementNestedDictionaryWithNullParent(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictionaryType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.ListOf(arrow.PrimitiveTypes.Int32),
+ }
+ dictionaryValuesBuilder := array.NewListBuilder(mem,
arrow.PrimitiveTypes.Int32)
+ dictionaryValuesBuilder.Append(true)
+ dictionaryValuesBuilder.ValueBuilder().(*array.Int32Builder).Append(7)
+ dictionaryValues := dictionaryValuesBuilder.NewArray()
+ dictionaryValuesBuilder.Release()
+ defer dictionaryValues.Release()
+
+ dictionaryIndicesBuilder := array.NewInt8Builder(mem)
+ dictionaryIndicesBuilder.Append(0)
+ dictionaryIndices := dictionaryIndicesBuilder.NewArray()
+ dictionaryIndicesBuilder.Release()
+ dictionary := array.NewDictionaryArray(dictionaryType,
dictionaryIndices, dictionaryValues)
+ dictionaryIndices.Release()
+ defer dictionary.Release()
+
+ offsetsBuilder := array.NewInt32Builder(mem)
+ offsetsBuilder.Append(0)
+ offsetsBuilder.Append(1)
+ offsetsBuilder.Append(1)
+ offsets := offsetsBuilder.NewArray()
+ offsetsBuilder.Release()
+ defer offsets.Release()
+
+ validity := memory.NewResizableBuffer(mem)
+ validity.Resize(1)
+ validity.Bytes()[0] = 0x01
+ data := array.NewData(
+ arrow.ListOf(dictionaryType),
+ 2,
+ []*memory.Buffer{validity, offsets.Data().Buffers()[1]},
+ []arrow.ArrayData{dictionary.Data()},
+ 1,
+ 0,
+ )
+ validity.Release()
+ input := array.NewListData(data)
+ data.Release()
+ defer input.Release()
+
+ expectedIndicesBuilder := array.NewInt8Builder(mem)
+ expectedIndicesBuilder.Append(0)
+ expectedIndicesBuilder.AppendNull()
+ expectedIndices := expectedIndicesBuilder.NewArray()
+ expectedIndicesBuilder.Release()
+ defer expectedIndices.Release()
+ expected := array.NewDictionaryArray(dictionaryType, expectedIndices,
dictionaryValues)
+ defer expected.Release()
+
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.NoError(t, array.ValidateFull(actual))
+ assert.True(t, array.Equal(expected, actual), "expected: %s\ngot: %s",
expected, actual)
+ assert.True(t, array.Equal(dictionaryValues,
actual.(*array.Dictionary).Dictionary()))
+}
+
+func TestListElementRunEndEncodedNestedDictionaryWithNullParent(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int8,
+ ValueType: arrow.ListOf(arrow.PrimitiveTypes.Int32),
+ }
+ valuesBuilder := array.NewListBuilder(mem, arrow.PrimitiveTypes.Int32)
+ valuesBuilder.Append(true)
+ valuesBuilder.ValueBuilder().(*array.Int32Builder).Append(7)
+ values := valuesBuilder.NewArray()
+ valuesBuilder.Release()
+ defer values.Release()
+ indicesBuilder := array.NewInt8Builder(mem)
+ indicesBuilder.Append(0)
+ indices := indicesBuilder.NewArray()
+ indicesBuilder.Release()
+ defer indices.Release()
+ dictionary := array.NewDictionaryArray(dictType, indices, values)
+ defer dictionary.Release()
+
+ for _, runEndType := range []arrow.DataType{
+ arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int64,
+ } {
+ t.Run(runEndType.String(), func(t *testing.T) {
+ endsBuilder := array.NewBuilder(mem, runEndType)
+ require.NoError(t,
endsBuilder.AppendValueFromString("1"))
+ ends := endsBuilder.NewArray()
+ endsBuilder.Release()
+ defer ends.Release()
+ child := array.NewRunEndEncodedArray(ends, dictionary,
1, 0)
+ defer child.Release()
+
+ offsetsBuilder := array.NewInt32Builder(mem)
+ offsetsBuilder.AppendValues([]int32{0, 1, 1}, nil)
+ offsets := offsetsBuilder.NewArray()
+ offsetsBuilder.Release()
+ defer offsets.Release()
+ validity := memory.NewResizableBuffer(mem)
+ validity.Resize(1)
+ validity.Bytes()[0] = 0x01
+ data := array.NewData(arrow.ListOf(child.DataType()), 2,
+ []*memory.Buffer{validity,
offsets.Data().Buffers()[1]},
+ []arrow.ArrayData{child.Data()}, 1, 0)
+ validity.Release()
+ input := array.NewListData(data)
+ data.Release()
+ defer input.Release()
+ require.NoError(t, array.ValidateFull(input))
+
+ for _, start := range []int64{0, 1} {
+ sliced := array.NewSlice(input, start, 2)
+ defer sliced.Release()
+ result, err := compute.ListElement(
+
compute.WithAllocator(context.Background(), mem),
+ &compute.ArrayDatum{Value:
sliced.Data()},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)},
+ )
+ require.NoError(t, err)
+ defer result.Release()
+ actual :=
result.(*compute.ArrayDatum).MakeArray()
+ defer actual.Release()
+ require.NoError(t, array.ValidateFull(actual))
+ require.Equal(t, 2-int(start), actual.Len())
+ encoded := actual.(*array.RunEndEncoded)
+ actualValues :=
encoded.Values().(*array.Dictionary)
+ assert.True(t, array.Equal(values,
actualValues.Dictionary()))
+ assert.True(t,
actualValues.IsNull(encoded.GetPhysicalIndex(actual.Len()-1)))
+ if start == 0 {
+ assert.False(t,
actualValues.IsNull(encoded.GetPhysicalIndex(0)))
+ assert.Equal(t, 0,
actualValues.GetValueIndex(encoded.GetPhysicalIndex(0)))
+ }
+ }
+ })
+ }
+}
+
+func TestListElementDenseUnionRecursiveErrorReleasesTemporaryIndices(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ sparseType := arrow.SparseUnionOf(
+ []arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int32,
Nullable: true},
+ {Name: "text", Type: arrow.BinaryTypes.String,
Nullable: true},
+ },
+ []arrow.UnionTypeCode{0, 1},
+ )
+ structType := arrow.StructOf(arrow.Field{Name: "union", Type:
sparseType, Nullable: true})
+ denseType := arrow.DenseUnionOf(
+ []arrow.Field{{Name: "struct", Type: structType, Nullable:
true}},
+ []arrow.UnionTypeCode{0},
+ )
+
+ builder := array.NewListBuilder(mem, denseType)
+ values := builder.ValueBuilder().(*array.DenseUnionBuilder)
+ structBuilder := values.Child(0).(*array.StructBuilder)
+ unionBuilder :=
structBuilder.FieldBuilder(0).(*array.SparseUnionBuilder)
+ for i := 0; i < 2; i++ {
+ builder.Append(true)
+ values.Append(0)
+ structBuilder.Append(true)
+ unionBuilder.Append(0)
+ unionBuilder.Child(0).(*array.Int32Builder).Append(int32(i))
+ unionBuilder.Child(1).(*array.StringBuilder).AppendNull()
+ }
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ for i := 0; i < 20; i++ {
+ result, err := compute.ListElement(
+ context.Background(),
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.Error(t, err)
+ }
+}
+
+func TestListElementValidatesScalarIndexForEmptyInputs(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ empty := listElementInput(t, mem,
arrow.ListOf(arrow.PrimitiveTypes.Int32), `[]`)
+ defer empty.Release()
+ chunked := arrow.NewChunked(empty.DataType(), []arrow.Array{empty})
+ defer chunked.Release()
+
+ tests := []struct {
+ name string
+ index scalar.Scalar
+ wantErr bool
+ }{
+ {name: "null scalar", index:
scalar.MakeNullScalar(arrow.PrimitiveTypes.Int64), wantErr: true},
+ {name: "negative scalar", index: scalar.NewInt64Scalar(-1),
wantErr: true},
+ {name: "zero scalar", index: scalar.NewInt64Scalar(0)},
+ {name: "large unsigned scalar", index:
scalar.NewUint64Scalar(^uint64(0))},
+ }
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ for _, lists := range []compute.Datum{
+ &compute.ArrayDatum{Value: empty.Data()},
+ &compute.ChunkedDatum{Value: chunked},
+ } {
+ result, err := compute.ListElement(
+ context.Background(),
+ lists,
+ &compute.ScalarDatum{Value: tc.index},
+ )
+ if err == nil && result != nil {
+ result.Release()
+ }
+ if tc.wantErr {
+ assert.ErrorIs(t, err, arrow.ErrInvalid)
+ } else {
+ assert.NoError(t, err)
+ }
+ }
+ })
+ }
+}
+
+func TestListElementNullScalarDenseUnionChildren(t *testing.T) {
+ for _, unsupported := range []arrow.DataType{
+ &arrow.Decimal32Type{Precision: 6, Scale: 2},
+ &arrow.Decimal64Type{Precision: 12, Scale: 2},
+ arrow.ListViewOf(arrow.PrimitiveTypes.Int32),
+ arrow.LargeListViewOf(arrow.PrimitiveTypes.Int32),
+ } {
+ for _, unsupportedFirst := range []bool{true, false} {
+ fields := []arrow.Field{
+ {Name: "unsupported", Type: unsupported,
Nullable: true},
+ {Name: "number", Type:
arrow.PrimitiveTypes.Int32, Nullable: true},
+ }
+ if !unsupportedFirst {
+ fields[0], fields[1] = fields[1], fields[0]
+ }
+ union := arrow.DenseUnionOf(fields,
[]arrow.UnionTypeCode{3, 7})
+ for _, elem := range []arrow.DataType{
+ union,
+ arrow.StructOf(arrow.Field{Name: "union", Type:
union, Nullable: true}),
+ &denseUnionExtensionType{ExtensionBase:
arrow.ExtensionBase{Storage: union}},
+ arrow.DenseUnionOf([]arrow.Field{{Name:
"nested", Type: union, Nullable: true}}, []arrow.UnionTypeCode{5}),
+ } {
+ for _, listType := range
[]arrow.DataType{arrow.ListOf(elem), arrow.LargeListOf(elem),
arrow.FixedSizeListOf(2, elem)} {
+
t.Run(fmt.Sprintf("%s/unsupported-first=%t", listType, unsupportedFirst),
func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx :=
compute.WithAllocator(context.Background(), mem)
+ list :=
scalar.MakeNullScalar(listType)
+ defer
list.(scalar.Releasable).Release()
+ result, err :=
compute.ListElement(ctx,
+
&compute.ScalarDatum{Value: list},
+
&compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)})
+ if err == nil && result != nil {
+ defer result.Release()
+ }
+ if unsupportedFirst {
+ require.ErrorIs(t, err,
arrow.ErrNotImplemented)
+ } else {
+ require.NoError(t, err)
+ actual :=
result.(*compute.ScalarDatum).Value
+ require.False(t,
actual.IsValid())
+ require.NoError(t,
actual.ValidateFull())
+ }
+ })
+ }
+ }
+ }
+ }
+}
+
+func TestListElementEmptyUnionChild(t *testing.T) {
+ for _, elem := range []arrow.DataType{arrow.DenseUnionOf(nil, nil),
arrow.SparseUnionOf(nil, nil)} {
+ t.Run(elem.ID().String(), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ builder := array.NewListBuilder(mem, elem)
+ builder.AppendNull()
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+ chunked := arrow.NewChunked(input.DataType(),
[]arrow.Array{input})
+ defer chunked.Release()
+ for _, lists := range []compute.Datum{
+ &compute.ArrayDatum{Value: input.Data()},
+ &compute.ChunkedDatum{Value: chunked},
+ } {
+ result, err :=
compute.ListElement(compute.WithAllocator(context.Background(), mem),
+ lists, &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)})
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ }
+ })
+ }
+}
+
+func TestListElementNullScalarValidation(t *testing.T) {
+ for _, elem := range []arrow.DataType{
+ arrow.DenseUnionOf(nil, nil),
+ arrow.SparseUnionOf(nil, nil),
+ } {
+ t.Run(elem.ID().String(), func(t *testing.T) {
+ list := scalar.MakeNullScalar(arrow.ListOf(elem))
+ defer list.(scalar.Releasable).Release()
+ result, err := compute.ListElement(context.Background(),
+ &compute.ScalarDatum{Value: list},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)})
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ })
+ }
+
+ list := scalar.MakeNullScalar(arrow.ListOf(arrow.PrimitiveTypes.Int32))
+ defer list.(scalar.Releasable).Release()
+ for _, index := range []scalar.Scalar{scalar.NewInt64Scalar(-1),
scalar.MakeNullScalar(arrow.PrimitiveTypes.Int64)} {
+ result, err := compute.ListElement(context.Background(),
+ &compute.ScalarDatum{Value: list},
&compute.ScalarDatum{Value: index})
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+ }
+
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+ result, err := compute.ListElement(ctx,
+ &compute.ScalarDatum{Value: list}, &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)})
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.ErrorIs(t, err, context.Canceled)
+
+ mapValue :=
scalar.MakeNullScalar(arrow.MapOf(arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int32))
+ defer mapValue.(scalar.Releasable).Release()
+ result, err = compute.ListElement(context.Background(),
+ &compute.ScalarDatum{Value: mapValue},
&compute.ScalarDatum{Value: scalar.NewInt64Scalar(0)})
+ if err == nil && result != nil {
+ result.Release()
+ }
+ require.Error(t, err)
+}
+
+func TestListElementNullStructScalarWithDenseUnionChild(t *testing.T) {
+ for _, unsupportedFirst := range []bool{true, false} {
+ t.Run(fmt.Sprintf("unsupported-first=%t", unsupportedFirst),
func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ fields := []arrow.Field{
+ {Name: "unsupported", Type:
&arrow.Decimal32Type{Precision: 6, Scale: 2}, Nullable: true},
+ {Name: "number", Type:
arrow.PrimitiveTypes.Int32, Nullable: true},
+ }
+ childID := 1
+ if !unsupportedFirst {
+ fields[0], fields[1] = fields[1], fields[0]
+ childID = 0
+ }
+ codes := []arrow.UnionTypeCode{3, 7}
+ unionType := arrow.DenseUnionOf(fields, codes)
+ builder := array.NewDenseUnionBuilder(mem, unionType)
+ builder.Append(codes[childID])
+ builder.Child(childID).(*array.Int32Builder).Append(42)
+ values := builder.NewArray()
+ builder.Release()
+ defer values.Release()
+
+ structType := arrow.StructOf(arrow.Field{Name: "union",
Type: unionType, Nullable: true})
+ validity := memory.NewBufferBytes([]byte{0})
+ data := array.NewData(structType, 1,
[]*memory.Buffer{validity}, []arrow.ArrayData{values.Data()}, 1, 0)
+ validity.Release()
+ child := array.NewStructData(data)
+ data.Release()
+ defer child.Release()
+ list := scalar.NewListScalar(child)
+ defer list.Release()
+
+ result, err :=
compute.ListElement(compute.WithAllocator(context.Background(), mem),
+ &compute.ScalarDatum{Value: list},
+ &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)})
+ if err == nil && result != nil {
+ defer result.Release()
+ }
+ if unsupportedFirst {
+ require.ErrorIs(t, err, arrow.ErrNotImplemented)
+ } else {
+ require.NoError(t, err)
+ actual := result.(*compute.ScalarDatum).Value
+ require.False(t, actual.IsValid())
+ require.NoError(t, actual.ValidateFull())
+ }
+ })
+ }
+}
+
+func BenchmarkListElement(b *testing.B) {
+ for _, size := range []int{1_000, 100_000, 1_000_000} {
+ b.Run(fmt.Sprintf("%d", size), func(b *testing.B) {
+ mem := memory.NewGoAllocator()
+ input := makeBenchmarkList(mem, size)
+ defer input.Release()
+ lists := &compute.ArrayDatum{Value: input.Data()}
+ index := &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)}
+
+ b.ReportAllocs()
+ b.SetBytes(int64(size * 4))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ result, err :=
compute.ListElement(context.Background(), lists, index)
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+ })
+ }
+}
+
+func BenchmarkListElementNested(b *testing.B) {
+ for _, size := range []int{1_000, 100_000} {
+ b.Run(fmt.Sprintf("%d", size), func(b *testing.B) {
+ mem := memory.NewGoAllocator()
+ input := makeBenchmarkNestedList(mem, size)
+ defer input.Release()
+ lists := &compute.ArrayDatum{Value: input.Data()}
+ index := &compute.ScalarDatum{Value:
scalar.NewInt64Scalar(0)}
+
+ b.ReportAllocs()
+ b.SetBytes(int64(size * 4))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ result, err :=
compute.ListElement(context.Background(), lists, index)
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+ })
+ }
+}
+
+func makeBenchmarkList(mem memory.Allocator, length int) arrow.Array {
+ builder := array.NewListBuilder(mem, arrow.PrimitiveTypes.Int32)
+ values := builder.ValueBuilder().(*array.Int32Builder)
+ builder.Reserve(length)
+ values.Reserve(length)
+ for i := 0; i < length; i++ {
+ builder.Append(true)
+ values.Append(int32(i))
+ }
+ result := builder.NewArray()
+ builder.Release()
+ return result
+}
+
+func makeBenchmarkNestedList(mem memory.Allocator, length int) arrow.Array {
+ outer := array.NewListBuilder(mem,
arrow.ListOf(arrow.PrimitiveTypes.Int32))
+ inner := outer.ValueBuilder().(*array.ListBuilder)
+ values := inner.ValueBuilder().(*array.Int32Builder)
+ outer.Reserve(length)
+ inner.Reserve(length)
+ values.Reserve(length)
+ for i := 0; i < length; i++ {
+ outer.Append(true)
+ inner.Append(true)
+ values.Append(int32(i))
+ }
+ result := outer.NewArray()
+ outer.Release()
+ return result
+}
diff --git a/arrow/scalar/scalar.go b/arrow/scalar/scalar.go
index 1d8dd13a..34166df6 100644
--- a/arrow/scalar/scalar.go
+++ b/arrow/scalar/scalar.go
@@ -615,7 +615,9 @@ func GetScalar(arr arrow.Array, idx int) (Scalar, error) {
if err != nil {
return nil, err
}
- return NewExtensionScalar(storage, arr.DataType()), nil
+ result := NewExtensionScalar(storage, arr.DataType())
+ result.Valid = storage.IsValid()
+ return result, nil
case *array.FixedSizeBinary:
width := arr.DataType().(*arrow.FixedSizeBinaryType).ByteWidth
buf := scalarValueBuffer(arr.Data().Buffers()[1],
(arr.Data().Offset()+idx)*width, width)
