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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
