perf(arrow/array): reuse full chunks in NewChunkedSlice

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
+}

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++
+				}
+			}
+		})
+	}
+}

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++