From 095720747ccc76826d57ce29bdfe909790d50a6d Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sat, 29 Aug 2026 13:27:36 +0200 Subject: [PATCH 1/2] perf(arrow/array): reuse full chunks in NewChunkedSlice --- arrow/array/chunked_slice_benchmark_test.go | 87 +++++++++++++++++++++ arrow/array/table.go | 7 +- 2 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 arrow/array/chunked_slice_benchmark_test.go diff --git a/arrow/array/chunked_slice_benchmark_test.go b/arrow/array/chunked_slice_benchmark_test.go new file mode 100644 index 000000000..1713b2a34 --- /dev/null +++ b/arrow/array/chunked_slice_benchmark_test.go @@ -0,0 +1,87 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package array_test + +import ( + "fmt" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +var benchmarkNewChunkedSliceSink int + +func BenchmarkNewChunkedSlice(b *testing.B) { + const rowsPerChunk = 64 + + for _, numChunks := range []int{64, 1024, 8192} { + for _, withNulls := range []bool{false, true} { + for _, partial := range []bool{false, true} { + name := fmt.Sprintf("chunks=%d/nulls=%t/partial=%t", numChunks, withNulls, partial) + b.Run(name, func(b *testing.B) { + input := makeChunkedSliceBenchmarkInput(numChunks, rowsPerChunk, withNulls) + defer input.Release() + + var start, end int64 + if partial { + start = 1 + end = int64(input.Len() - 1) + } else { + end = int64(input.Len()) + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + result := array.NewChunkedSlice(input, start, end) + benchmarkNewChunkedSliceSink = result.Len() + result.Release() + } + }) + } + } + } +} + +func makeChunkedSliceBenchmarkInput(numChunks, rowsPerChunk int, withNulls bool) *arrow.Chunked { + chunks := make([]arrow.Array, numChunks) + values := make([]int64, rowsPerChunk) + validity := make([]bool, rowsPerChunk) + for i := range values { + values[i] = int64(i) + validity[i] = i%10 != 0 + } + + for i := range chunks { + builder := array.NewInt64Builder(memory.DefaultAllocator) + if withNulls { + builder.AppendValues(values, validity) + } else { + builder.AppendValues(values, nil) + } + chunks[i] = builder.NewArray() + builder.Release() + } + + result := arrow.NewChunked(arrow.PrimitiveTypes.Int64, chunks) + for _, chunk := range chunks { + chunk.Release() + } + return result +} diff --git a/arrow/array/table.go b/arrow/array/table.go index 9e3e83c7f..04b83f5d4 100644 --- a/arrow/array/table.go +++ b/arrow/array/table.go @@ -69,7 +69,12 @@ func NewChunkedSlice(a *arrow.Chunked, i, j int64) *arrow.Chunked { if end > int64(arr.Len()) { end = int64(arr.Len()) } - chunks = append(chunks, NewSlice(arr, beg, end)) + if beg == 0 && end == int64(arr.Len()) { + arr.Retain() + chunks = append(chunks, arr) + } else { + chunks = append(chunks, NewSlice(arr, beg, end)) + } sz -= int64(arr.Len()) - beg beg = 0 cur++ From f3a10b67eb9beb49e8e629f83680237f1fe030b9 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sun, 30 Aug 2026 22:38:02 +0200 Subject: [PATCH 2/2] test(arrow/array): cover reused chunk ownership --- arrow/array/chunked_slice_test.go | 81 +++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 arrow/array/chunked_slice_test.go diff --git a/arrow/array/chunked_slice_test.go b/arrow/array/chunked_slice_test.go new file mode 100644 index 000000000..4251ed692 --- /dev/null +++ b/arrow/array/chunked_slice_test.go @@ -0,0 +1,81 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package array_test + +import ( + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/stretchr/testify/require" +) + +func TestNewChunkedSliceRetainsFullChunks(t *testing.T) { + for _, tc := range []struct { + name string + begin, end int64 + want []int32 + valid []bool + fullChunks []int + }{ + {"all", 0, 6, []int32{1, 2, 3, 4, 5, 6}, []bool{true, false, true, true, false, true}, []int{0, 1, 2}}, + {"middle", 2, 4, []int32{3, 4}, []bool{true, true}, []int{1}}, + {"partial ends", 1, 5, []int32{2, 3, 4, 5}, []bool{false, true, true, false}, []int{-1, 1, -1}}, + {"empty", 6, 6, nil, nil, nil}, + } { + t.Run(tc.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + builder := array.NewInt32Builder(mem) + builder.AppendValues([]int32{0, 1, 2, 3, 4, 5, 6, 7}, []bool{true, true, false, true, true, false, true, true}) + backing := builder.NewInt32Array() + builder.Release() + chunks := []arrow.Array{ + array.NewSlice(backing, 1, 3), + array.NewSlice(backing, 3, 5), + array.NewSlice(backing, 5, 7), + } + backing.Release() + input := arrow.NewChunked(arrow.PrimitiveTypes.Int32, chunks) + for _, chunk := range chunks { + chunk.Release() + } + result := array.NewChunkedSlice(input, tc.begin, tc.end) + defer result.Release() + input.Release() + + require.Equal(t, len(tc.want), result.Len()) + require.Len(t, result.Chunks(), len(tc.fullChunks)) + position := 0 + for i, chunk := range result.Chunks() { + if full := tc.fullChunks[i]; full >= 0 { + require.Same(t, chunks[full], chunk) + } + values := chunk.(*array.Int32) + for j := range values.Len() { + require.Equal(t, tc.valid[position], values.IsValid(j)) + if tc.valid[position] { + require.Equal(t, tc.want[position], values.Value(j)) + } + position++ + } + } + }) + } +}