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 ad7a9338 perf(arrow/ipc): reuse record encoder scratch (#1221)
ad7a9338 is described below
commit ad7a933828725d03a889199696ab9876cb93f50d
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 28 23:25:21 2026 +0200
perf(arrow/ipc): reuse record encoder scratch (#1221)
## Summary
- Keep one `recordEncoder` on `Writer` and `FileWriter`.
- Reuse the `fields`, `meta`, and `variadicCounts` backing arrays.
- Restore the encoder recursion depth when resetting.
- Add a benchmark for repeated writes with 1, 16, and 64 columns.
- Add a regression test for resetting after a nested encode error.
## Benchmark
**Command**
```text
go test ./arrow/ipc -run '^$' -bench '^BenchmarkWriterRecordEncoderReuse$'
-benchmem -benchtime=100ms -count=3
```
**Machine:** Apple M1 Pro, arm64
**Go:** 1.26.3
The benchmark writes one warm-up batch before timing repeated writes.
Each width was tested with 16, 256, and 4096 rows. The allocation
results were the same across those row counts.
| Columns | Before | After |
| --- | ---: | ---: |
| 1 | 1,624 B/op, 21 allocs/op | 1,408 B/op, 18 allocs/op |
| 16 | 8,072 B/op, 33 allocs/op | 6,656 B/op, 26 allocs/op |
| 64 | 31,240 B/op, 41 allocs/op | 25,984 B/op, 32 allocs/op |
Timing was noisy on this machine, so the allocation reduction is the
main result.
## Tests
- `go test ./arrow/ipc`
- `go test ./arrow/...`
---
arrow/ipc/file_writer.go | 28 ++++++--
arrow/ipc/writer.go | 49 +++++++++-----
arrow/ipc/writer_benchmark_test.go | 81 +++++++++++++++++++++++
arrow/ipc/writer_test.go | 131 +++++++++++++++++++++++++++++++++++++
4 files changed, 265 insertions(+), 24 deletions(-)
diff --git a/arrow/ipc/file_writer.go b/arrow/ipc/file_writer.go
index 89c98999..ffba57c1 100644
--- a/arrow/ipc/file_writer.go
+++ b/arrow/ipc/file_writer.go
@@ -256,6 +256,7 @@ type FileWriter struct {
codec flatbuf.CompressionType
compressNP int
compressors []compressor
+ encoder *recordEncoder
minSpaceSavings float64
// map of the last written dictionaries by id
@@ -286,7 +287,25 @@ func NewFileWriter(w io.Writer, opts ...Option)
(*FileWriter, error) {
return &f, err
}
+func (f *FileWriter) getRecordEncoder() *recordEncoder {
+ if f.encoder == nil {
+ f.encoder = newRecordEncoder(
+ f.mem,
+ 0,
+ kMaxNestingDepth,
+ true,
+ f.codec,
+ f.compressNP,
+ f.minSpaceSavings,
+ f.compressors,
+ )
+ }
+ return f.encoder
+}
+
func (f *FileWriter) Close() error {
+ defer func() { f.encoder = nil }()
+
if f.closed {
return f.closeErr
}
@@ -334,13 +353,8 @@ func (f *FileWriter) Write(rec arrow.RecordBatch) error {
return fmt.Errorf("arrow/ipc: could not write header: %w", err)
}
- const allow64b = true
- var (
- data = Payload{msg: MessageRecordBatch}
- enc = newRecordEncoder(
- f.mem, 0, kMaxNestingDepth, allow64b, f.codec,
f.compressNP, f.minSpaceSavings, f.compressors,
- )
- )
+ data := Payload{msg: MessageRecordBatch}
+ enc := f.getRecordEncoder()
defer data.Release()
err := writeDictionaryPayloads(f.mem, rec, true, false, &f.mapper,
f.lastWrittenDicts, f.pw, enc)
diff --git a/arrow/ipc/writer.go b/arrow/ipc/writer.go
index a771849f..a1c768ba 100644
--- a/arrow/ipc/writer.go
+++ b/arrow/ipc/writer.go
@@ -96,6 +96,7 @@ type Writer struct {
codec flatbuf.CompressionType
compressNP int
compressors []compressor
+ encoder *recordEncoder
minSpaceSavings float64
// map of the last written dictionaries by id
@@ -137,7 +138,25 @@ func NewWriter(w io.Writer, opts ...Option) *Writer {
}
}
+func (w *Writer) getRecordEncoder() *recordEncoder {
+ if w.encoder == nil {
+ w.encoder = newRecordEncoder(
+ w.mem,
+ 0,
+ kMaxNestingDepth,
+ true,
+ w.codec,
+ w.compressNP,
+ w.minSpaceSavings,
+ w.compressors,
+ )
+ }
+ return w.encoder
+}
+
func (w *Writer) Close() error {
+ defer func() { w.encoder = nil }()
+
if w.err != nil {
return w.closeAfterFailure()
}
@@ -211,20 +230,8 @@ func (w *Writer) Write(rec arrow.RecordBatch) (err error) {
return errInconsistentSchema
}
- const allow64b = true
- var (
- data = Payload{msg: MessageRecordBatch}
- enc = newRecordEncoder(
- w.mem,
- 0,
- kMaxNestingDepth,
- allow64b,
- w.codec,
- w.compressNP,
- w.minSpaceSavings,
- w.compressors,
- )
- )
+ data := Payload{msg: MessageRecordBatch}
+ enc := w.getRecordEncoder()
defer data.Release()
err = writeDictionaryPayloads(w.mem, rec, false, w.emitDictDeltas,
&w.mapper, w.lastWrittenDicts, w.pw, enc)
@@ -368,6 +375,7 @@ type recordEncoder struct {
variadicCounts []int64
depth int64
+ maxDepth int64
start int64
allow64b bool
codec flatbuf.CompressionType
@@ -390,6 +398,7 @@ func newRecordEncoder(
mem: mem,
start: startOffset,
depth: maxDepth,
+ maxDepth: maxDepth,
allow64b: allow64b,
codec: codec,
compressNP: compressNP,
@@ -409,9 +418,11 @@ func (w *recordEncoder) shouldCompress(uncompressed,
compressed int) bool {
}
func (w *recordEncoder) reset() {
+ w.depth = w.maxDepth
w.start = 0
- w.fields = make([]fieldMetadata, 0)
- w.variadicCounts = nil
+ w.fields = w.fields[:0]
+ w.meta = w.meta[:0]
+ w.variadicCounts = w.variadicCounts[:0]
}
func (w *recordEncoder) getCompressor(id int) compressor {
@@ -550,7 +561,11 @@ func (w *recordEncoder) encode(p *Payload, rec
arrow.RecordBatch) error {
// position for the start of a buffer relative to the passed frame of
reference.
// may be 0 or some other position in an address space.
offset := w.start
- w.meta = make([]bufferMetadata, len(p.body))
+ if cap(w.meta) < len(p.body) {
+ w.meta = make([]bufferMetadata, len(p.body))
+ } else {
+ w.meta = w.meta[:len(p.body)]
+ }
// construct the metadata for the record batch header
for i, buf := range p.body {
diff --git a/arrow/ipc/writer_benchmark_test.go
b/arrow/ipc/writer_benchmark_test.go
new file mode 100644
index 00000000..3567ec4f
--- /dev/null
+++ b/arrow/ipc/writer_benchmark_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 ipc
+
+import (
+ "fmt"
+ "io"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+func benchmarkRecordBatch(numColumns, numRows int) arrow.RecordBatch {
+ fields := make([]arrow.Field, numColumns)
+ columns := make([]arrow.Array, numColumns)
+ values := make([]int32, numRows)
+ for i := range values {
+ values[i] = int32(i)
+ }
+
+ for i := range fields {
+ fields[i] = arrow.Field{Name: fmt.Sprintf("col%d", i), Type:
arrow.PrimitiveTypes.Int32}
+
+ builder := array.NewInt32Builder(memory.DefaultAllocator)
+ builder.AppendValues(values, nil)
+ columns[i] = builder.NewArray()
+ builder.Release()
+ }
+
+ schema := arrow.NewSchema(fields, nil)
+ record := array.NewRecordBatch(schema, columns, int64(numRows))
+ for _, column := range columns {
+ column.Release()
+ }
+ return record
+}
+
+func BenchmarkWriterRecordEncoderReuse(b *testing.B) {
+ for _, numColumns := range []int{1, 16, 64} {
+ for _, numRows := range []int{16, 256, 4096} {
+ b.Run(fmt.Sprintf("%dcols/%drows", numColumns,
numRows), func(b *testing.B) {
+ record := benchmarkRecordBatch(numColumns,
numRows)
+ defer record.Release()
+
+ writer := NewWriter(io.Discard,
WithSchema(record.Schema()))
+ if err := writer.Write(record); err != nil {
+ b.Fatal(err)
+ }
+
+ b.ReportAllocs()
+ b.SetBytes(int64(numColumns * numRows *
arrow.Int32SizeBytes))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ if err := writer.Write(record); err !=
nil {
+ b.Fatal(err)
+ }
+ }
+ b.StopTimer()
+ if err := writer.Close(); err != nil {
+ b.Fatal(err)
+ }
+ })
+ }
+ }
+}
diff --git a/arrow/ipc/writer_test.go b/arrow/ipc/writer_test.go
index e49c4f54..88b2dc6d 100644
--- a/arrow/ipc/writer_test.go
+++ b/arrow/ipc/writer_test.go
@@ -69,6 +69,20 @@ type failingWriter struct {
err error
}
+type closeFailWriter struct {
+ bytes.Buffer
+ err error
+ fail bool
+}
+
+func (w *closeFailWriter) Write(p []byte) (int, error) {
+ if w.fail {
+ return 0, w.err
+ }
+
+ return w.Buffer.Write(p)
+}
+
func (shortWriteWriter) Write(p []byte) (int, error) {
return len(p) - 1, io.ErrShortWrite
}
@@ -118,6 +132,39 @@ func TestWriterCloseFailureIsTerminal(t *testing.T) {
require.Equal(t, 1, payloadWriter.closeCall)
}
+func TestWriterCloseReleasesRecordEncoder(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{{Name: "col", Type:
arrow.PrimitiveTypes.Int32}}, nil)
+ builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+ defer builder.Release()
+ record := builder.NewRecordBatch()
+ defer record.Release()
+
+ var output bytes.Buffer
+ writer := NewWriter(&output, WithSchema(schema))
+ require.NoError(t, writer.Write(record))
+ require.NotNil(t, writer.encoder)
+
+ require.NoError(t, writer.Close())
+ require.Nil(t, writer.encoder)
+}
+
+func TestWriterCloseAfterFailureReleasesRecordEncoder(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{{Name: "col", Type:
arrow.PrimitiveTypes.Int32}}, nil)
+ builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+ defer builder.Release()
+ record := builder.NewRecordBatch()
+ defer record.Release()
+
+ want := errors.New("payload failed")
+ payloadWriter := &failingPayloadWriter{err: want, failAfter: 2}
+ writer := NewWriterWithPayloadWriter(payloadWriter, WithSchema(schema))
+ require.ErrorIs(t, writer.Write(record), want)
+ require.NotNil(t, writer.encoder)
+
+ require.ErrorIs(t, writer.Close(), want)
+ require.Nil(t, writer.encoder)
+}
+
func TestFileWriterCloseFailureIsTerminal(t *testing.T) {
schema := arrow.NewSchema([]arrow.Field{{Name: "col", Type:
arrow.PrimitiveTypes.Int32}}, nil)
want := errors.New("write failed")
@@ -128,6 +175,42 @@ func TestFileWriterCloseFailureIsTerminal(t *testing.T) {
require.ErrorIs(t, writer.Close(), want)
}
+func TestFileWriterCloseReleasesRecordEncoder(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{{Name: "col", Type:
arrow.PrimitiveTypes.Int32}}, nil)
+ builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+ defer builder.Release()
+ record := builder.NewRecordBatch()
+ defer record.Release()
+
+ var output bytes.Buffer
+ writer, err := NewFileWriter(&output, WithSchema(schema))
+ require.NoError(t, err)
+ require.NoError(t, writer.Write(record))
+ require.NotNil(t, writer.encoder)
+
+ require.NoError(t, writer.Close())
+ require.Nil(t, writer.encoder)
+}
+
+func TestFileWriterCloseFailureReleasesRecordEncoder(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{{Name: "col", Type:
arrow.PrimitiveTypes.Int32}}, nil)
+ builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+ defer builder.Release()
+ record := builder.NewRecordBatch()
+ defer record.Release()
+
+ want := errors.New("close failed")
+ output := &closeFailWriter{err: want}
+ writer, err := NewFileWriter(output, WithSchema(schema))
+ require.NoError(t, err)
+ require.NoError(t, writer.Write(record))
+ require.NotNil(t, writer.encoder)
+
+ output.fail = true
+ require.ErrorIs(t, writer.Close(), want)
+ require.Nil(t, writer.encoder)
+}
+
func TestWriterSchemaFailureIsTerminal(t *testing.T) {
schema := arrow.NewSchema([]arrow.Field{{Name: "col", Type:
arrow.PrimitiveTypes.Int32}}, nil)
builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
@@ -827,3 +910,51 @@ func TestVariadicCountsNotAccumulatedAcrossEncode(t
*testing.T) {
p.Release()
}
}
+
+func TestRecordEncoderResetRestoresDepth(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ childBuilder := array.NewInt32Builder(mem)
+ childBuilder.Append(1)
+ child := childBuilder.NewArray()
+ childBuilder.Release()
+ defer child.Release()
+
+ structArr, err := array.NewStructArrayWithFields(
+ []arrow.Array{child},
+ []arrow.Field{{Name: "value", Type:
arrow.PrimitiveTypes.Int32}},
+ )
+ require.NoError(t, err)
+ defer structArr.Release()
+
+ structSchema := arrow.NewSchema([]arrow.Field{{
+ Name: "struct",
+ Type: structArr.DataType(),
+ }}, nil)
+ structRecord := array.NewRecordBatch(structSchema,
[]arrow.Array{structArr}, 1)
+ defer structRecord.Release()
+
+ enc := newRecordEncoder(mem, 0, 1, true, -1, 1, 0, nil)
+ var nestedPayload Payload
+ require.ErrorIs(t, enc.encode(&nestedPayload, structRecord),
errMaxRecursion)
+ nestedPayload.Release()
+
+ enc.reset()
+
+ simpleSchema := arrow.NewSchema([]arrow.Field{{
+ Name: "value",
+ Type: arrow.PrimitiveTypes.Int32,
+ }}, nil)
+ simpleBuilder := array.NewInt32Builder(mem)
+ simpleBuilder.Append(1)
+ simple := simpleBuilder.NewArray()
+ simpleBuilder.Release()
+ defer simple.Release()
+ simpleRecord := array.NewRecordBatch(simpleSchema,
[]arrow.Array{simple}, 1)
+ defer simpleRecord.Release()
+
+ var simplePayload Payload
+ require.NoError(t, enc.encode(&simplePayload, simpleRecord))
+ simplePayload.Release()
+}
