You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ze...@apache.org on 2026/08/19 16:28:11 UTC
(arrow-go) branch main updated: perf(arrow/compute): take chunked binary values without concatenation (#1199)
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 d3b223f1 perf(arrow/compute): take chunked binary values without concatenation (#1199)
d3b223f1 is described below
commit d3b223f1c613bcb7b28fd62c257dca4fd4b93959
Author: Minh Vu <vu...@gmail.com>
AuthorDate: Wed Aug 19 18:28:05 2026 +0200
perf(arrow/compute): take chunked binary values without concatenation (#1199)
## What
Chunked string and binary `take` currently concatenates all value chunks
before selecting rows.
This adds a chunk-aware path for:
- STRING
- BINARY
- LARGE_STRING
- LARGE_BINARY
The new path resolves global indices to source chunks and copies only
selected values. It keeps one output chunk per index chunk and preserves
null and bounds behavior.
## Benchmark
64 chunks x 4096 rows, 32-byte values, random indices, Apple M1 Pro:
- 1% selected: 0.87 ms / 9.6 MB -> 0.08 ms / 0.11 MB
- 10% selected: 2.36 ms / 10.4 MB -> 1.43 ms / 0.96 MB
The benchmark covers 8 and 64 chunks, 1% / 10% / 50% / 100% selection,
and all four types.
## Tests
- `go test ./arrow/compute/... -count=1`
---------
Co-authored-by: Matt Topol <ma...@columnar.tech>
---
.github/workflows/test.yml | 1 +
arrow/compute/chunked_take_bench_test.go | 193 ++++++++++++++
arrow/compute/chunked_take_test.go | 247 ++++++++++++++++++
arrow/compute/executor.go | 6 +-
arrow/compute/executor_test.go | 27 ++
.../compute/internal/kernels/chunked_take_test.go | 42 +++
arrow/compute/internal/kernels/vector_selection.go | 282 ++++++++++++++++++++-
7 files changed, 793 insertions(+), 5 deletions(-)
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 567b34bb..6b7e713b 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -406,6 +406,7 @@ jobs:
archery docker run \
-e ARCHERY_DEFAULT_BRANCH=${{ github.event.repository.default_branch }} \
-e ARCHERY_INTEGRATION_TARGET_IMPLEMENTATIONS=go \
+ -e GTest_SOURCE=SYSTEM \
-e ARCHERY_INTEGRATION_WITH_DOTNET=1 \
-e ARCHERY_INTEGRATION_WITH_GO=1 \
-e ARCHERY_INTEGRATION_WITH_JAVA=1 \
diff --git a/arrow/compute/chunked_take_bench_test.go b/arrow/compute/chunked_take_bench_test.go
new file mode 100644
index 00000000..4ec255e4
--- /dev/null
+++ b/arrow/compute/chunked_take_bench_test.go
@@ -0,0 +1,193 @@
+// 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"
+ "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/memory"
+ "github.com/stretchr/testify/require"
+)
+
+func TestChunkedBinaryTake(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ ctx := compute.WithAllocator(context.Background(), mem)
+ for _, typ := range []arrow.DataType{
+ arrow.BinaryTypes.String,
+ arrow.BinaryTypes.Binary,
+ arrow.BinaryTypes.LargeString,
+ arrow.BinaryTypes.LargeBinary,
+ } {
+ t.Run(typ.String(), func(t *testing.T) {
+ chunk0Full := newTestBinaryArray(mem, typ,
+ [][]byte{[]byte("prefix"), []byte("hello"), []byte("world"), []byte("suffix")},
+ []bool{true, true, true, true})
+ chunk0 := array.NewSlice(chunk0Full, 1, 3)
+ chunk1 := newTestBinaryArray(mem, typ,
+ [][]byte{[]byte("unused"), []byte("foo"), []byte("bar"), []byte("baz")},
+ []bool{false, true, true, true})
+ empty := newTestBinaryArray(mem, typ, nil, nil)
+ values := arrow.NewChunked(typ, []arrow.Array{empty, chunk0, empty, chunk1})
+ defer values.Release()
+ empty.Release()
+ chunk0.Release()
+ chunk0Full.Release()
+ chunk1.Release()
+
+ indices := newTestInt64Array(mem, []int64{4, 0, 2, 0, 5, 1}, []bool{true, true, true, false, true, true})
+ defer indices.Release()
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ChunkedDatum{Value: values}, &compute.ArrayDatum{Value: indices.Data()})
+ require.NoError(t, err)
+ actual := result.(*compute.ChunkedDatum).Value
+ expectedArray := newTestBinaryArray(mem, typ,
+ [][]byte{[]byte("bar"), []byte("hello"), nil, nil, []byte("baz"), []byte("world")},
+ []bool{true, true, false, false, true, true})
+ expected := arrow.NewChunked(typ, []arrow.Array{expectedArray})
+ require.True(t, array.ChunkedEqual(expected, actual))
+ result.Release()
+
+ invalid := newTestInt64Array(mem, []int64{6}, nil)
+ _, err = compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ChunkedDatum{Value: values}, &compute.ArrayDatum{Value: invalid.Data()})
+ require.ErrorIs(t, err, arrow.ErrIndex)
+ invalid.Release()
+
+ indices0 := newTestInt64Array(mem, []int64{4, 0, 2}, nil)
+ indices1 := newTestInt64Array(mem, []int64{0, 5, 1}, []bool{false, true, true})
+ chunkedIndices := arrow.NewChunked(arrow.PrimitiveTypes.Int64, []arrow.Array{indices0, indices1})
+ defer chunkedIndices.Release()
+ indices0.Release()
+ indices1.Release()
+
+ result, err = compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ChunkedDatum{Value: values}, &compute.ChunkedDatum{Value: chunkedIndices})
+ require.NoError(t, err)
+ actual = result.(*compute.ChunkedDatum).Value
+ require.True(t, array.ChunkedEqual(expected, actual))
+ result.Release()
+ expected.Release()
+ expectedArray.Release()
+ })
+ }
+ mem.AssertSize(t, 0)
+}
+
+func newTestBinaryArray(mem memory.Allocator, typ arrow.DataType, values [][]byte, valid []bool) arrow.Array {
+ bldr := array.NewBinaryBuilder(mem, typ.(arrow.BinaryDataType))
+ bldr.Reserve(len(values))
+ for i, value := range values {
+ if len(valid) != 0 && !valid[i] {
+ bldr.AppendNull()
+ } else {
+ bldr.Append(value)
+ }
+ }
+ result := bldr.NewArray()
+ bldr.Release()
+ return result
+}
+
+func newTestInt64Array(mem memory.Allocator, values []int64, valid []bool) arrow.Array {
+ bldr := array.NewInt64Builder(mem)
+ bldr.Reserve(len(values))
+ for i, value := range values {
+ if len(valid) != 0 && !valid[i] {
+ bldr.AppendNull()
+ } else {
+ bldr.Append(value)
+ }
+ }
+ result := bldr.NewArray()
+ bldr.Release()
+ return result
+}
+
+func BenchmarkTakeChunkedBinary(b *testing.B) {
+ for _, typ := range []arrow.DataType{
+ arrow.BinaryTypes.String,
+ arrow.BinaryTypes.Binary,
+ arrow.BinaryTypes.LargeString,
+ arrow.BinaryTypes.LargeBinary,
+ } {
+ for _, numChunks := range []int{8, 64} {
+ for _, selectivity := range []int{1, 10, 50, 100} {
+ name := fmt.Sprintf("%s/chunks=%d/selectivity=%d%%", typ, numChunks, selectivity)
+ b.Run(name, func(b *testing.B) {
+ mem := memory.DefaultAllocator
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ const rowsPerChunk = 4096
+ values, indices := makeChunkedBinaryTakeInputs(mem, typ, numChunks, rowsPerChunk, selectivity)
+ defer values.Release()
+ defer indices.Release()
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ChunkedDatum{Value: values},
+ &compute.ArrayDatum{Value: indices.Data()})
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+ })
+ }
+ }
+ }
+}
+
+func makeChunkedBinaryTakeInputs(mem memory.Allocator, typ arrow.DataType, numChunks, rowsPerChunk, selectivity int) (*arrow.Chunked, arrow.Array) {
+ chunks := make([]arrow.Array, numChunks)
+ value := []byte("0123456789abcdefghijklmnopqrstuv")
+ for i := range chunks {
+ bldr := array.NewBinaryBuilder(mem, typ.(arrow.BinaryDataType))
+ bldr.Reserve(rowsPerChunk)
+ bldr.ReserveData(rowsPerChunk * len(value))
+ for j := 0; j < rowsPerChunk; j++ {
+ bldr.Append(value)
+ }
+ chunks[i] = bldr.NewArray()
+ bldr.Release()
+ }
+
+ values := arrow.NewChunked(typ, chunks)
+ for _, chunk := range chunks {
+ chunk.Release()
+ }
+
+ totalRows := numChunks * rowsPerChunk
+ numIndices := totalRows * selectivity / 100
+ indicesBldr := array.NewInt64Builder(mem)
+ indicesBldr.Reserve(numIndices)
+ for i := 0; i < numIndices; i++ {
+ idx := (int64(i)*1103515245 + 12345) % int64(totalRows)
+ indicesBldr.Append(idx)
+ }
+ indices := indicesBldr.NewArray()
+ indicesBldr.Release()
+ return values, indices
+}
diff --git a/arrow/compute/chunked_take_test.go b/arrow/compute/chunked_take_test.go
new file mode 100644
index 00000000..e6e0784f
--- /dev/null
+++ b/arrow/compute/chunked_take_test.go
@@ -0,0 +1,247 @@
+// 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"
+ "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/memory"
+ "github.com/stretchr/testify/require"
+)
+
+var chunkedTakeIndexTypes = []arrow.DataType{
+ arrow.PrimitiveTypes.Int8,
+ arrow.PrimitiveTypes.Uint8,
+ arrow.PrimitiveTypes.Int16,
+ arrow.PrimitiveTypes.Uint16,
+ arrow.PrimitiveTypes.Int32,
+ arrow.PrimitiveTypes.Uint32,
+ arrow.PrimitiveTypes.Int64,
+ arrow.PrimitiveTypes.Uint64,
+}
+
+var chunkedTakeBinaryTypes = []arrow.DataType{
+ arrow.BinaryTypes.String,
+ arrow.BinaryTypes.Binary,
+ arrow.BinaryTypes.LargeString,
+ arrow.BinaryTypes.LargeBinary,
+}
+
+func TestChunkedBinaryTakeIndexTypes(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ for _, typ := range chunkedTakeBinaryTypes {
+ t.Run(typ.String(), func(t *testing.T) {
+ values := makeChunkedTakeValues(mem, typ)
+ defer values.Release()
+
+ for _, indexType := range chunkedTakeIndexTypes {
+ t.Run(indexType.String(), func(t *testing.T) {
+ indices0Full := makeChunkedTakeIndexArray(t, mem, indexType,
+ []string{"77", "0", "5", "99", "3", "77"},
+ []bool{true, true, true, false, true, true})
+ indices0 := array.NewSlice(indices0Full, 1, 5)
+ indices1Full := makeChunkedTakeIndexArray(t, mem, indexType,
+ []string{"77", "2", "1", "77"}, nil)
+ indices1 := array.NewSlice(indices1Full, 1, 3)
+
+ indices := arrow.NewChunked(indexType, []arrow.Array{indices0, indices1})
+ indices0Full.Release()
+ indices0.Release()
+ indices1Full.Release()
+ indices1.Release()
+ defer indices.Release()
+
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ChunkedDatum{Value: values},
+ &compute.ChunkedDatum{Value: indices})
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ChunkedDatum).Value
+ require.Equal(t, 6, actual.Len())
+ require.Len(t, actual.Chunks(), 2)
+ require.Equal(t, 4, actual.Chunk(0).Len())
+ require.Equal(t, 2, actual.Chunk(1).Len())
+
+ expected0 := newTestBinaryArray(mem, typ,
+ [][]byte{[]byte("zero"), []byte("five"), nil, []byte("three")},
+ []bool{true, true, false, true})
+ expected1 := newTestBinaryArray(mem, typ,
+ [][]byte{[]byte("two"), []byte("one")}, nil)
+ expected := arrow.NewChunked(typ, []arrow.Array{expected0, expected1})
+ expected0.Release()
+ expected1.Release()
+ defer expected.Release()
+ require.True(t, array.ChunkedEqual(expected, actual))
+ })
+ }
+ })
+ }
+
+ mem.AssertSize(t, 0)
+}
+
+func TestChunkedBinaryTakeBoundsChecks(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ for _, typ := range chunkedTakeBinaryTypes {
+ t.Run(typ.String(), func(t *testing.T) {
+ values := makeChunkedTakeValues(mem, typ)
+ defer values.Release()
+
+ for _, indexType := range chunkedTakeIndexTypes {
+ cases := []struct {
+ name string
+ value string
+ }{
+ {name: "upper_bound", value: "6"},
+ }
+ if arrow.IsSignedInteger(indexType.ID()) {
+ cases = append(cases, struct {
+ name string
+ value string
+ }{name: "negative", value: "-1"})
+ }
+
+ for _, tc := range cases {
+ t.Run(indexType.String()+"/"+tc.name, func(t *testing.T) {
+ indices := makeChunkedTakeIndexArray(t, mem, indexType, []string{tc.value}, nil)
+ defer indices.Release()
+
+ _, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ChunkedDatum{Value: values},
+ &compute.ArrayDatum{Value: indices.Data()})
+ require.ErrorIs(t, err, arrow.ErrIndex)
+ })
+ }
+ }
+ })
+ }
+
+ mem.AssertSize(t, 0)
+}
+
+func TestChunkedBinaryTakeEmptyIndexChunks(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ for _, typ := range chunkedTakeBinaryTypes {
+ t.Run(typ.String(), func(t *testing.T) {
+ values := makeChunkedTakeValues(mem, typ)
+ defer values.Release()
+
+ for _, indexType := range chunkedTakeIndexTypes {
+ t.Run(indexType.String(), func(t *testing.T) {
+ t.Run("zero_chunks", func(t *testing.T) {
+ indices := arrow.NewChunked(indexType, nil)
+ defer indices.Release()
+
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ChunkedDatum{Value: values},
+ &compute.ChunkedDatum{Value: indices})
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ChunkedDatum).Value
+ require.Equal(t, 0, actual.Len())
+ require.Empty(t, actual.Chunks())
+ })
+
+ for _, tc := range []struct {
+ name string
+ chunkedValues [][]string
+ wantLen int
+ wantChunks int
+ }{
+ {name: "empty_chunk", chunkedValues: [][]string{nil}, wantLen: 0, wantChunks: 0},
+ {name: "leading_empty", chunkedValues: [][]string{nil, {"0", "5"}}, wantLen: 2, wantChunks: 1},
+ {name: "trailing_empty", chunkedValues: [][]string{{"0", "5"}, nil}, wantLen: 2, wantChunks: 1},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ chunks := make([]arrow.Array, len(tc.chunkedValues))
+ for i, chunkValues := range tc.chunkedValues {
+ chunks[i] = makeChunkedTakeIndexArray(t, mem, indexType, chunkValues, nil)
+ }
+ indices := arrow.NewChunked(indexType, chunks)
+ for _, chunk := range chunks {
+ chunk.Release()
+ }
+ defer indices.Release()
+
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ChunkedDatum{Value: values},
+ &compute.ChunkedDatum{Value: indices})
+ require.NoError(t, err)
+ defer result.Release()
+
+ actual := result.(*compute.ChunkedDatum).Value
+ require.Equal(t, tc.wantLen, actual.Len())
+ require.Len(t, actual.Chunks(), tc.wantChunks)
+ if tc.wantChunks == 1 {
+ require.Equal(t, 2, actual.Chunk(0).Len())
+ }
+ })
+ }
+ })
+ }
+ })
+ }
+
+ mem.AssertSize(t, 0)
+}
+
+func makeChunkedTakeValues(mem memory.Allocator, typ arrow.DataType) *arrow.Chunked {
+ full := newTestBinaryArray(mem, typ,
+ [][]byte{[]byte("prefix"), []byte("zero"), []byte("one"), []byte("two"), []byte("three"), []byte("suffix")}, nil)
+ middle := array.NewSlice(full, 1, 5)
+ tail := newTestBinaryArray(mem, typ, [][]byte{[]byte("four"), []byte("five")}, nil)
+ empty0 := newTestBinaryArray(mem, typ, nil, nil)
+ empty1 := newTestBinaryArray(mem, typ, nil, nil)
+
+ values := arrow.NewChunked(typ, []arrow.Array{empty0, middle, tail, empty1})
+ full.Release()
+ middle.Release()
+ tail.Release()
+ empty0.Release()
+ empty1.Release()
+ return values
+}
+
+func makeChunkedTakeIndexArray(t testing.TB, mem memory.Allocator, typ arrow.DataType, values []string, valid []bool) arrow.Array {
+ t.Helper()
+ bldr := array.NewBuilder(mem, typ)
+ defer bldr.Release()
+ bldr.Reserve(len(values))
+ for i, value := range values {
+ if len(valid) != 0 && !valid[i] {
+ bldr.AppendNull()
+ continue
+ }
+ require.NoError(t, bldr.AppendValueFromString(value))
+ }
+ result := bldr.NewArray()
+ return result
+}
diff --git a/arrow/compute/executor.go b/arrow/compute/executor.go
index 1728e49e..3468d640 100644
--- a/arrow/compute/executor.go
+++ b/arrow/compute/executor.go
@@ -1137,9 +1137,11 @@ func (v *vectorExecutor) execChunked(batch *ExecBatch, out chan<- Datum) error {
}
if len(result) == 0 {
- empty := output.MakeArray()
+ outType := output.Type
+ empty := array.MakeArrayOfNull(exec.GetAllocator(v.ctx.Ctx), outType, 0)
defer empty.Release()
- out <- &ChunkedDatum{Value: arrow.NewChunked(output.Type, []arrow.Array{empty})}
+ output.Release()
+ out <- &ChunkedDatum{Value: arrow.NewChunked(outType, []arrow.Array{empty})}
return nil
}
diff --git a/arrow/compute/executor_test.go b/arrow/compute/executor_test.go
index a8194108..dbc2312b 100644
--- a/arrow/compute/executor_test.go
+++ b/arrow/compute/executor_test.go
@@ -39,6 +39,33 @@ func (d *signalChunkedDatum) Chunks() []arrow.Array {
return d.Value.Chunks()
}
+func TestVectorExecutorWrapResultsReleasesEmptyArrayOutput(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ builder := array.NewBinaryBuilder(mem, arrow.BinaryTypes.String)
+ empty := builder.NewArray()
+ builder.Release()
+
+ output := make(chan Datum, 1)
+ output <- NewDatum(empty)
+ close(output)
+
+ executor := &vectorExecutor{
+ nonAggExecImpl: nonAggExecImpl{
+ kernel: &exec.VectorKernel{OutputChunked: true},
+ outType: arrow.BinaryTypes.String,
+ },
+ }
+
+ result := executor.WrapResults(context.Background(), output, true)
+ require.NotNil(t, result)
+ require.Equal(t, KindChunked, result.Kind())
+ require.Empty(t, result.(*ChunkedDatum).Value.Chunks())
+ result.Release()
+ empty.Release()
+}
+
func TestVectorExecutorWrapResultsReleasesEmptyChunkedOutput(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
diff --git a/arrow/compute/internal/kernels/chunked_take_test.go b/arrow/compute/internal/kernels/chunked_take_test.go
new file mode 100644
index 00000000..64b5ee30
--- /dev/null
+++ b/arrow/compute/internal/kernels/chunked_take_test.go
@@ -0,0 +1,42 @@
+// 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 (
+ "math"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/stretchr/testify/require"
+)
+
+func TestTakeChunkedBinaryOffsetLimits(t *testing.T) {
+ maxOffset, offsetBytes := binaryTakeOffsetLimits[int32]()
+ require.Equal(t, int64(math.MaxInt32), maxOffset)
+ require.Equal(t, 4, offsetBytes)
+ require.NoError(t, checkBinaryTakeOffset[int32](0, math.MaxInt32))
+ require.ErrorIs(t,
+ checkBinaryTakeOffset[int32](0, int64(math.MaxInt32)+1),
+ arrow.ErrInvalid)
+
+ maxOffset, offsetBytes = binaryTakeOffsetLimits[int64]()
+ require.Equal(t, int64(math.MaxInt64), maxOffset)
+ require.Equal(t, 8, offsetBytes)
+ require.NoError(t, checkBinaryTakeOffset[int64](0, int64(math.MaxInt32)+1))
+}
diff --git a/arrow/compute/internal/kernels/vector_selection.go b/arrow/compute/internal/kernels/vector_selection.go
index f08b53f7..00bc141a 100644
--- a/arrow/compute/internal/kernels/vector_selection.go
+++ b/arrow/compute/internal/kernels/vector_selection.go
@@ -662,6 +662,72 @@ func (c *chunkedPrimitiveGetter[T]) GetValue(i int64) T {
func (c *chunkedPrimitiveGetter[T]) NullCount() int64 { return c.nulls }
func (c *chunkedPrimitiveGetter[T]) Len() int64 { return c.len }
+type binaryGetter interface {
+ IsValid(int64) bool
+ GetValue(int64) []byte
+ NullCount() int64
+ Len() int64
+ DataLen() int64
+}
+
+type chunkedBinaryGetter[OffsetT int32 | int64] struct {
+ resolver *exec.ChunkResolver
+ offsets [][]OffsetT
+ values [][]byte
+ valuesIsValid [][]byte
+ valuesOffset []int64
+ nulls int64
+ length int64
+ dataLen int64
+}
+
+func newChunkedBinaryGetter[OffsetT int32 | int64](arr *arrow.Chunked) *chunkedBinaryGetter[OffsetT] {
+ chunks := make([]arrow.Array, 0, len(arr.Chunks()))
+ for _, chunk := range arr.Chunks() {
+ if chunk.Len() > 0 {
+ chunks = append(chunks, chunk)
+ }
+ }
+ getter := &chunkedBinaryGetter[OffsetT]{
+ resolver: exec.NewChunkResolver(chunks),
+ offsets: make([][]OffsetT, len(chunks)),
+ values: make([][]byte, len(chunks)),
+ valuesIsValid: make([][]byte, len(chunks)),
+ valuesOffset: make([]int64, len(chunks)),
+ nulls: int64(arr.NullN()),
+ length: int64(arr.Len()),
+ }
+
+ var span exec.ArraySpan
+ for i, chunk := range chunks {
+ span.SetMembers(chunk.Data())
+ getter.values[i] = span.Buffers[2].Buf
+ getter.valuesIsValid[i] = span.Buffers[0].Buf
+ getter.valuesOffset[i] = span.Offset
+ if span.Len > 0 {
+ getter.offsets[i] = exec.GetSpanOffsets[OffsetT](&span, 1)
+ getter.dataLen += int64(getter.offsets[i][span.Len] - getter.offsets[i][0])
+ }
+ }
+ return getter
+}
+
+func (c *chunkedBinaryGetter[OffsetT]) IsValid(i int64) bool {
+ chunk, index := c.resolver.Resolve(i)
+ bitmap := c.valuesIsValid[chunk]
+ return bitmap == nil || bitutil.BitIsSet(bitmap, int(c.valuesOffset[chunk]+index))
+}
+
+func (c *chunkedBinaryGetter[OffsetT]) GetValue(i int64) []byte {
+ chunk, index := c.resolver.Resolve(i)
+ offsets := c.offsets[chunk]
+ return c.values[chunk][offsets[index]:offsets[index+1]]
+}
+
+func (c *chunkedBinaryGetter[OffsetT]) NullCount() int64 { return c.nulls }
+func (c *chunkedBinaryGetter[OffsetT]) Len() int64 { return c.length }
+func (c *chunkedBinaryGetter[OffsetT]) DataLen() int64 { return c.dataLen }
+
// isSorted checks if indices are monotonically increasing (sorted)
// Returns true if sorted, false otherwise
// Uses sampling for large arrays to avoid full scan
@@ -1161,6 +1227,216 @@ func ChunkedPrimitiveTake(ctx *exec.KernelCtx, batch []*arrow.Chunked, out *exec
}
}
+func binaryTakeOffsetLimits[OffsetT int32 | int64]() (int64, int) {
+ var zero OffsetT
+ switch any(zero).(type) {
+ case int32:
+ return math.MaxInt32, 4
+ case int64:
+ return math.MaxInt64, 8
+ default:
+ panic("unsupported binary offset type")
+ }
+}
+
+func checkBinaryTakeOffset[OffsetT int32 | int64](offset OffsetT, valueLen int64) error {
+ maxOffset, _ := binaryTakeOffsetLimits[OffsetT]()
+ if offset < 0 || valueLen < 0 || valueLen > maxOffset-int64(offset) {
+ return fmt.Errorf("%w: binary output offset overflow", arrow.ErrInvalid)
+ }
+ return nil
+}
+
+func takeChunkedBinaryImpl[IdxT arrow.UintType, OffsetT int32 | int64](ctx *exec.KernelCtx, indices *exec.ArraySpan, values binaryGetter, out *exec.ExecResult) error {
+ var (
+ indicesValues = exec.GetSpanValues[IdxT](indices, 1)
+ indicesIsValid = bitutil.OptionalBitIndexer{Bitmap: indices.Buffers[0].Buf, Offset: int(indices.Offset)}
+ bitCounter = bitutils.NewOptionalBitBlockCounter(indices.Buffers[0].Buf, indices.Offset, indices.Len)
+ validityBuilder = validityBuilder{mem: exec.GetAllocator(ctx.Ctx)}
+ offsetBuilder = newBufferBuilder[OffsetT](exec.GetAllocator(ctx.Ctx))
+ dataBuilder = newBufferBuilder[uint8](exec.GetAllocator(ctx.Ctx))
+ valuesHaveNulls = values.NullCount() != 0
+ pos int64
+ offset OffsetT
+ )
+
+ maxInt := int(^uint(0) >> 1)
+ if indices.Len >= int64(maxInt) || indices.Len > math.MaxInt64-7 {
+ return fmt.Errorf("%w: binary take input length exceeds capacity", arrow.ErrInvalid)
+ }
+ offsetElements := int(indices.Len) + 1
+ _, offsetSize := binaryTakeOffsetLimits[OffsetT]()
+ if offsetElements > maxInt/offsetSize {
+ return fmt.Errorf("%w: binary take offset buffer exceeds capacity", arrow.ErrInvalid)
+ }
+
+ defer func() {
+ if validityBuilder.buffer != nil {
+ validityBuilder.buffer.Release()
+ }
+ if offsetBuilder.buffer != nil {
+ offsetBuilder.buffer.Release()
+ }
+ if dataBuilder.buffer != nil {
+ dataBuilder.buffer.Release()
+ }
+ }()
+
+ validityBuilder.Reserve(indices.Len)
+ offsetBuilder.reserve(offsetElements)
+ if values.Len() > 0 && values.DataLen() > 0 && indices.Len > 0 {
+ const maxPrealloc = int64(16777216)
+ meanValueLen := values.DataLen() / values.Len()
+ estimatedTotalSize := int64(0)
+ if meanValueLen > 0 {
+ if meanValueLen >= maxPrealloc || indices.Len > maxPrealloc/meanValueLen {
+ estimatedTotalSize = maxPrealloc
+ } else {
+ estimatedTotalSize = meanValueLen * indices.Len
+ }
+ }
+ dataBuilder.reserve(int(estimatedTotalSize))
+ }
+
+ spaceAvail := dataBuilder.cap()
+ appendValue := func(idx int64) error {
+ value := values.GetValue(idx)
+ valueLen := int64(len(value))
+ if err := checkBinaryTakeOffset(offset, valueLen); err != nil {
+ return err
+ }
+
+ dataLen := dataBuilder.len()
+ if len(value) > maxInt-dataLen {
+ return fmt.Errorf("%w: binary output size exceeds capacity", arrow.ErrInvalid)
+ }
+
+ offsetBuilder.unsafeAppend(offset)
+ if len(value) > spaceAvail {
+ needed := dataLen + len(value)
+ newCap := dataBuilder.cap()
+ if newCap < dataLen {
+ newCap = dataLen
+ }
+ if newCap == 0 {
+ newCap = len(value)
+ }
+ for newCap < needed {
+ if newCap > maxInt/2 {
+ newCap = needed
+ break
+ }
+ newCap *= 2
+ }
+ dataBuilder.reserve(newCap - dataLen)
+ spaceAvail = dataBuilder.cap() - dataBuilder.len()
+ }
+ dataBuilder.unsafeAppendSlice(value)
+ spaceAvail -= len(value)
+ offset += OffsetT(valueLen)
+ return nil
+ }
+ appendNull := func() {
+ offsetBuilder.unsafeAppend(offset)
+ }
+
+ for pos < indices.Len {
+ block := bitCounter.NextBlock()
+ indicesHaveNulls := block.Popcnt < block.Len
+ switch {
+ case !indicesHaveNulls && !valuesHaveNulls:
+ validityBuilder.UnsafeAppendN(int64(block.Len), true)
+ for i := 0; i < int(block.Len); i++ {
+ if err := appendValue(int64(indicesValues[pos])); err != nil {
+ return err
+ }
+ pos++
+ }
+ case block.Popcnt > 0:
+ for i := 0; i < int(block.Len); i++ {
+ idxValid := !indicesHaveNulls || indicesIsValid.GetBit(int(pos))
+ if idxValid && (!valuesHaveNulls || values.IsValid(int64(indicesValues[pos]))) {
+ validityBuilder.UnsafeAppend(true)
+ if err := appendValue(int64(indicesValues[pos])); err != nil {
+ return err
+ }
+ } else {
+ validityBuilder.UnsafeAppend(false)
+ appendNull()
+ }
+ pos++
+ }
+ default:
+ validityBuilder.UnsafeAppendN(int64(block.Len), false)
+ for i := 0; i < int(block.Len); i++ {
+ appendNull()
+ }
+ pos += int64(block.Len)
+ }
+ }
+
+ offsetBuilder.unsafeAppend(offset)
+ out.Len = indices.Len
+ out.Nulls = int64(validityBuilder.falseCount)
+ out.Buffers[0].WrapBuffer(validityBuilder.Finish())
+ out.Buffers[1].WrapBuffer(offsetBuilder.finish())
+ out.Buffers[2].WrapBuffer(dataBuilder.finish())
+ return nil
+}
+
+func takeChunkedBinaryDispatch[IdxT arrow.UintType, OffsetT int32 | int64](ctx *exec.KernelCtx, values binaryGetter, indices *arrow.Chunked, out []*exec.ExecResult) error {
+ var span exec.ArraySpan
+ for i, chunk := range indices.Chunks() {
+ span.SetMembers(chunk.Data())
+ if err := takeChunkedBinaryImpl[IdxT, OffsetT](ctx, &span, values, out[i]); err != nil {
+ return err
+ }
+ }
+ return nil
+}
+
+func ChunkedVarBinaryTake[OffsetT int32 | int64](ctx *exec.KernelCtx, batch []*arrow.Chunked, out *exec.ExecResult) ([]*exec.ExecResult, error) {
+ values, indices := batch[0], batch[1]
+ if ctx.State.(TakeState).BoundsCheck {
+ if err := checkIndexBoundsChunked(indices, uint64(values.Len())); err != nil {
+ return nil, err
+ }
+ }
+
+ outData := make([]*exec.ExecResult, len(indices.Chunks()))
+ for i := range outData {
+ outData[i] = &exec.ExecResult{Type: out.Type}
+ }
+
+ cleanup := func() {
+ for _, result := range outData {
+ if result != nil {
+ result.Release()
+ }
+ }
+ }
+
+ valuesGetter := newChunkedBinaryGetter[OffsetT](values)
+ var err error
+ switch indices.DataType().(arrow.FixedWidthDataType).Bytes() {
+ case 1:
+ err = takeChunkedBinaryDispatch[uint8, OffsetT](ctx, valuesGetter, indices, outData)
+ case 2:
+ err = takeChunkedBinaryDispatch[uint16, OffsetT](ctx, valuesGetter, indices, outData)
+ case 4:
+ err = takeChunkedBinaryDispatch[uint32, OffsetT](ctx, valuesGetter, indices, outData)
+ case 8:
+ err = takeChunkedBinaryDispatch[uint64, OffsetT](ctx, valuesGetter, indices, outData)
+ default:
+ err = fmt.Errorf("%w: invalid byte width for indices", arrow.ErrIndex)
+ }
+ if err != nil {
+ cleanup()
+ return nil, err
+ }
+ return outData, nil
+}
+
func NullTake(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error {
if ctx.State.(TakeState).BoundsCheck {
if err := checkIndexBounds(&batch.Values[1].Array, uint64(batch.Values[0].Array.Len)); err != nil {
@@ -2058,7 +2334,7 @@ type SelectionKernelData struct {
}
func ChunkedTakeSupported(dt arrow.DataType) bool {
- return arrow.IsPrimitive(dt.ID())
+ return arrow.IsPrimitive(dt.ID()) || arrow.IsBaseBinary(dt.ID())
}
func GetVectorSelectionKernels() (filterkernels, takeKernels []SelectionKernelData) {
@@ -2078,8 +2354,8 @@ func GetVectorSelectionKernels() (filterkernels, takeKernels []SelectionKernelDa
{In: exec.NewIDInput(arrow.DECIMAL128), Exec: TakeExec(FSBImpl)},
{In: exec.NewIDInput(arrow.DECIMAL256), Exec: TakeExec(FSBImpl)},
{In: exec.NewIDInput(arrow.FIXED_SIZE_BINARY), Exec: TakeExec(FSBImpl)},
- {In: exec.NewMatchedInput(exec.BinaryLike()), Exec: TakeExec(VarBinaryImpl[int32])},
- {In: exec.NewMatchedInput(exec.LargeBinaryLike()), Exec: TakeExec(VarBinaryImpl[int64])},
+ {In: exec.NewMatchedInput(exec.BinaryLike()), Exec: TakeExec(VarBinaryImpl[int32]), Chunked: ChunkedVarBinaryTake[int32]},
+ {In: exec.NewMatchedInput(exec.LargeBinaryLike()), Exec: TakeExec(VarBinaryImpl[int64]), Chunked: ChunkedVarBinaryTake[int64]},
}
return
}