You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "fallintoplace (via GitHub)" <gi...@apache.org> on 2026/08/12 17:13:28 UTC
[PR] perf(arrow/compute): take chunked binary values without concatenation [arrow-go]
fallintoplace opened a new pull request, #1199:
URL: https://github.com/apache/arrow-go/pull/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`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(arrow/compute): take chunked binary values without concatenation [arrow-go]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade merged PR #1199:
URL: https://github.com/apache/arrow-go/pull/1199
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(arrow/compute): take chunked binary values without concatenation [arrow-go]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #1199:
URL: https://github.com/apache/arrow-go/pull/1199#discussion_r3807806694
##########
arrow/compute/internal/kernels/vector_selection.go:
##########
@@ -1161,6 +1227,144 @@ func ChunkedPrimitiveTake(ctx *exec.KernelCtx, batch []*arrow.Chunked, out *exec
}
}
+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
+ )
+
+ validityBuilder.Reserve(indices.Len)
+ offsetBuilder.reserve(int(indices.Len) + 1)
+ if values.Len() > 0 {
+ meanValueLen := float64(values.DataLen()) / float64(values.Len())
+ estimatedTotalSize := min(int(meanValueLen*float64(indices.Len)), 16777216)
+ dataBuilder.reserve(estimatedTotalSize)
+ }
+
+ spaceAvail := dataBuilder.cap()
+ appendValue := func(idx int64) {
+ offsetBuilder.unsafeAppend(offset)
+ value := values.GetValue(idx)
+ if len(value) > spaceAvail {
+ needed := dataBuilder.len() + len(value)
+ newCap := dataBuilder.cap()
+ if newCap == 0 {
+ newCap = len(value)
+ }
+ for newCap < needed {
+ newCap *= 2
+ }
+ dataBuilder.reserve(newCap - dataBuilder.len())
+ spaceAvail = dataBuilder.cap() - dataBuilder.len()
+ }
+ dataBuilder.unsafeAppendSlice(value)
+ spaceAvail -= len(value)
+ offset += OffsetT(len(value))
+ }
+ 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++ {
+ appendValue(int64(indicesValues[pos]))
+ 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)
+ appendValue(int64(indicesValues[pos]))
+ } 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)
Review Comment:
**Blocking:** a zero-length index chunk leaks its output offsets buffer.
This terminating offset correctly creates a real offsets buffer for an empty binary array. However, `vectorExecutor.WrapResults.toChunked` calls `ArrayDatum.Chunks()`, which retains the resulting array, and then silently drops zero-length chunks without releasing them.
With multi-chunk string values and an empty index array:
- merge base: 0 bytes remain allocated;
- this PR: 64 bytes leak per call.
A leading empty chunk in chunked indices leaks similarly.
Please fix the shared executor path to release zero-length arrays that it drops, mirroring the release behavior in the later accumulation branch. Add checked-allocator tests for empty, leading-empty, and trailing-empty index chunks.
##########
arrow/compute/internal/kernels/vector_selection.go:
##########
@@ -1161,6 +1227,144 @@ func ChunkedPrimitiveTake(ctx *exec.KernelCtx, batch []*arrow.Chunked, out *exec
}
}
+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
+ )
+
+ validityBuilder.Reserve(indices.Len)
+ offsetBuilder.reserve(int(indices.Len) + 1)
+ if values.Len() > 0 {
+ meanValueLen := float64(values.DataLen()) / float64(values.Len())
+ estimatedTotalSize := min(int(meanValueLen*float64(indices.Len)), 16777216)
+ dataBuilder.reserve(estimatedTotalSize)
+ }
+
+ spaceAvail := dataBuilder.cap()
+ appendValue := func(idx int64) {
+ offsetBuilder.unsafeAppend(offset)
+ value := values.GetValue(idx)
+ if len(value) > spaceAvail {
+ needed := dataBuilder.len() + len(value)
+ newCap := dataBuilder.cap()
+ if newCap == 0 {
+ newCap = len(value)
+ }
+ for newCap < needed {
+ newCap *= 2
+ }
+ dataBuilder.reserve(newCap - dataBuilder.len())
+ spaceAvail = dataBuilder.cap() - dataBuilder.len()
+ }
+ dataBuilder.unsafeAppendSlice(value)
+ spaceAvail -= len(value)
+ offset += OffsetT(len(value))
+ }
+ 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++ {
+ appendValue(int64(indicesValues[pos]))
+ 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)
+ appendValue(int64(indicesValues[pos]))
+ } 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()))
Review Comment:
**Blocking:** when `indices.Chunks()` is empty, this returns an empty result slice. `execChunked` then calls `output.MakeArray()` on the uninitialized output span, causing a panic such as:
```text
arrow/array: string offset buffer must have at least 4 values
```
The merge-base concatenate path returns a normal empty chunked result for the same input.
Please preserve that behavior by constructing a valid empty variable-binary result, or by fixing the shared empty-result path so it can create the required offsets buffer. Add a zero-chunk-index regression test for all four supported types.
##########
arrow/compute/internal/kernels/vector_selection.go:
##########
@@ -1161,6 +1227,144 @@ func ChunkedPrimitiveTake(ctx *exec.KernelCtx, batch []*arrow.Chunked, out *exec
}
}
+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
+ )
+
+ validityBuilder.Reserve(indices.Len)
+ offsetBuilder.reserve(int(indices.Len) + 1)
+ if values.Len() > 0 {
+ meanValueLen := float64(values.DataLen()) / float64(values.Len())
+ estimatedTotalSize := min(int(meanValueLen*float64(indices.Len)), 16777216)
+ dataBuilder.reserve(estimatedTotalSize)
+ }
+
+ spaceAvail := dataBuilder.cap()
+ appendValue := func(idx int64) {
+ offsetBuilder.unsafeAppend(offset)
+ value := values.GetValue(idx)
+ if len(value) > spaceAvail {
+ needed := dataBuilder.len() + len(value)
+ newCap := dataBuilder.cap()
+ if newCap == 0 {
+ newCap = len(value)
+ }
+ for newCap < needed {
+ newCap *= 2
+ }
+ dataBuilder.reserve(newCap - dataBuilder.len())
+ spaceAvail = dataBuilder.cap() - dataBuilder.len()
+ }
+ dataBuilder.unsafeAppendSlice(value)
+ spaceAvail -= len(value)
+ offset += OffsetT(len(value))
Review Comment:
For the `int32` instantiation used by STRING/BINARY, this silently wraps once the selected data in an output index chunk exceeds `math.MaxInt32`.
The previous concatenate path rejected chunked Binary/String data that overflowed 32-bit offsets. Bypassing concatenation now allows large chunked sources, but the resulting individual output chunk must still satisfy the Binary offset limit.
Please check before incrementing and return an overflow/capacity error rather than emitting corrupt offsets. The capacity-growth arithmetic should likewise avoid integer overflow.
##########
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})
Review Comment:
The correctness test uses only `Int64` indices, leaving the 8-, 16-, and 32-bit dispatch branches and unsigned index types untested.
Please table-test all supported signed and unsigned index widths, including:
- negative and upper-bound failures;
- out-of-range garbage under null index slots;
- sliced index arrays;
- zero-length and zero-chunk index inputs;
- explicit output chunk-count assertions for chunked indices.
The logical `ChunkedEqual` assertion currently cannot detect an incorrect output chunk layout.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org