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 cc09645c perf(arrow/array): reuse TableReader scratch slices (#1200)
cc09645c is described below
commit cc09645cd46034d0bb9ab9a65cb2c68d0cf2788a
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 21 17:43:55 2026 +0200
perf(arrow/array): reuse TableReader scratch slices (#1200)
## What
`TableReader.Next` allocated two column slices for every record batch.
This keeps the batch slice on `TableReader` and reads the current chunks
again after finding the batch size. `NewRecordBatch` copies the input
slice, so the scratch slice can be safely reused on the next call.
## Benchmark
The benchmark reuses one reader across passes. It uses Int32 columns
with 256 rows per chunk.
256 columns x 256 chunks:
- 10.2 ms -> 3.0 ms
- 3.76 MB -> 1.27 MB
- 1024 -> 512 allocations
The benchmark covers 4, 32, and 256 columns with 32 and 256 chunks.
## Tests
- `go test ./arrow/array -count=1`
- `go test -race ./arrow/array -run TestTableReader -count=1`
---
arrow/array/table.go | 19 +++----
arrow/array/table_reader_bench_test.go | 91 ++++++++++++++++++++++++++++++++++
arrow/array/table_test.go | 34 +++++++++++++
3 files changed, 135 insertions(+), 9 deletions(-)
diff --git a/arrow/array/table.go b/arrow/array/table.go
index 9c8f7012..9e3e83c7 100644
--- a/arrow/array/table.go
+++ b/arrow/array/table.go
@@ -308,6 +308,8 @@ type TableReader struct {
chunks []*arrow.Chunked
slots []int // chunk indices
offsets []int64 // chunk offsets
+ // batch is reused as scratch input; NewRecordBatch copies the slice.
+ batch []arrow.Array
}
// NewTableReader returns a new TableReader to iterate over the (possibly
chunked) Table.
@@ -322,6 +324,7 @@ func NewTableReader(tbl arrow.Table, chunkSize int64)
*TableReader {
chunks: make([]*arrow.Chunked, ncols),
slots: make([]int, ncols),
offsets: make([]int64, ncols),
+ batch: make([]arrow.Array, ncols),
}
tr.refCount.Add(1)
tr.tbl.Retain()
@@ -355,8 +358,7 @@ func (tr *TableReader) Next() bool {
// determine the minimum contiguous slice across all columns
chunksz := imin64(tr.max-tr.cur, tr.chksz)
- chunks := make([]arrow.Array, len(tr.chunks))
- for i := range chunks {
+ for i := range tr.chunks {
j := tr.slots[i]
chunk := tr.chunks[i].Chunk(j)
for chunk.Len() == 0 && j+1 < len(tr.chunks[i].Chunks()) {
@@ -369,12 +371,10 @@ func (tr *TableReader) Next() bool {
chunksz = remain
}
- chunks[i] = chunk
}
-
// slice the chunks, advance each chunk slot as appropriate.
- batch := make([]arrow.Array, len(tr.chunks))
- for i, chunk := range chunks {
+ for i := range tr.chunks {
+ chunk := tr.chunks[i].Chunk(tr.slots[i])
var slice arrow.Array
offset := tr.offsets[i]
switch int64(chunk.Len()) - offset {
@@ -393,13 +393,13 @@ func (tr *TableReader) Next() bool {
tr.offsets[i] += chunksz
slice = NewSlice(chunk, offset, offset+chunksz)
}
- batch[i] = slice
+ tr.batch[i] = slice
}
tr.cur += chunksz
- tr.rec = NewRecordBatch(tr.tbl.Schema(), batch, chunksz)
+ tr.rec = NewRecordBatch(tr.tbl.Schema(), tr.batch, chunksz)
- for _, arr := range batch {
+ for _, arr := range tr.batch {
arr.Release()
}
@@ -430,6 +430,7 @@ func (tr *TableReader) Release() {
tr.chunks = nil
tr.slots = nil
tr.offsets = nil
+ tr.batch = nil
}
}
func (tr *TableReader) Err() error { return nil }
diff --git a/arrow/array/table_reader_bench_test.go
b/arrow/array/table_reader_bench_test.go
new file mode 100644
index 00000000..641181dc
--- /dev/null
+++ b/arrow/array/table_reader_bench_test.go
@@ -0,0 +1,91 @@
+// 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
+
+import (
+ "fmt"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+func BenchmarkTableReaderNext(b *testing.B) {
+ const rowsPerChunk = 256
+ for _, numCols := range []int{4, 32, 256} {
+ for _, numChunks := range []int{32, 256} {
+ b.Run(fmt.Sprintf("columns=%d/chunks=%d", numCols,
numChunks), func(b *testing.B) {
+ table :=
makeTableReaderBenchmarkTable(memory.DefaultAllocator, numCols, numChunks,
rowsPerChunk)
+ defer table.Release()
+
+ reader := NewTableReader(table, rowsPerChunk)
+ defer reader.Release()
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ var rows int64
+ for reader.Next() {
+ rows +=
reader.RecordBatch().NumRows()
+ }
+ if rows != table.NumRows() {
+ b.Fatalf("invalid row count:
got=%d, want=%d", rows, table.NumRows())
+ }
+
+ reader.cur = 0
+ clear(reader.slots)
+ clear(reader.offsets)
+ }
+ })
+ }
+ }
+}
+
+func makeTableReaderBenchmarkTable(mem memory.Allocator, numCols, numChunks,
rowsPerChunk int) arrow.Table {
+ fields := make([]arrow.Field, numCols)
+ for i := range fields {
+ fields[i] = arrow.Field{Name: fmt.Sprintf("col_%d", i), Type:
arrow.PrimitiveTypes.Int32}
+ }
+ schema := arrow.NewSchema(fields, nil)
+
+ chunks := make([]arrow.Array, numChunks)
+ for i := range chunks {
+ bldr := NewInt32Builder(mem)
+ bldr.Reserve(rowsPerChunk)
+ for j := 0; j < rowsPerChunk; j++ {
+ bldr.Append(int32(i*rowsPerChunk + j))
+ }
+ chunks[i] = bldr.NewInt32Array()
+ bldr.Release()
+ }
+
+ cols := make([]arrow.Column, numCols)
+ for i, field := range fields {
+ chunked := arrow.NewChunked(field.Type, chunks)
+ cols[i] = *arrow.NewColumn(field, chunked)
+ chunked.Release()
+ }
+ table := NewTable(schema, cols, -1)
+
+ for i := range cols {
+ cols[i].Release()
+ }
+ for _, chunk := range chunks {
+ chunk.Release()
+ }
+ return table
+}
diff --git a/arrow/array/table_test.go b/arrow/array/table_test.go
index b1da6f3e..d42c0d10 100644
--- a/arrow/array/table_test.go
+++ b/arrow/array/table_test.go
@@ -1019,6 +1019,40 @@ func TestTableReaderSkipsEmptyChunks(t *testing.T) {
}
}
+func TestTableReaderRetainedRecordBatchSurvivesNext(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ builder := array.NewInt32Builder(mem)
+ builder.AppendValues([]int32{1, 2, 3, 4}, nil)
+ values := builder.NewInt32Array()
+ builder.Release()
+ defer values.Release()
+
+ field := arrow.Field{Name: "values", Type: arrow.PrimitiveTypes.Int32}
+ column := arrow.NewColumnFromArr(field, values)
+ defer column.Release()
+ table := array.NewTable(arrow.NewSchema([]arrow.Field{field}, nil),
[]arrow.Column{column}, -1)
+ defer table.Release()
+
+ reader := array.NewTableReader(table, 2)
+ defer reader.Release()
+ if !reader.Next() {
+ t.Fatal("expected the first record batch")
+ }
+ first := reader.RecordBatch()
+ first.Retain()
+ defer first.Release()
+
+ if !reader.Next() {
+ t.Fatal("expected the second record batch")
+ }
+ got := first.Column(0).(*array.Int32).Value(0)
+ if got != 1 {
+ t.Fatalf("first record batch changed after Next: got=%d,
want=1", got)
+ }
+}
+
func TestTableToString(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer mem.AssertSize(t, 0)
