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 372e1eab perf(arrow/ipc): release dictionary payloads eagerly (#1259)
372e1eab is described below
commit 372e1eabeb9d8db3962e75219049431677528083
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 31 22:38:59 2026 +0200
perf(arrow/ipc): release dictionary payloads eagerly (#1259)
## What
- Release each dictionary payload as soon as `WritePayload` returns.
- Keep the optional dictionary delta slice in the same short-lived
scope.
- This keeps encoded dictionary payloads from accumulating until
`writeDictionaryPayloads` returns.
## Benchmark
The benchmark uses a tracking allocator and reports peak bytes held by
dictionary payloads beyond the input record.
Apple M1 Pro, Go 1.26.3, 1 CPU, 150 ms, 7 runs:
| Dictionary columns | Before | After |
| ---: | ---: | ---: |
| 8 | 1,536 B | 192 B |
| 64 | 12,288 B | 192 B |
| 256 | 49,152 B | 192 B |
## Tests
- `go test ./arrow/ipc`
- `PARQUET_TEST_DATA="$(pwd)/parquet-testing/data" go test ./...`
- `go vet ./arrow/ipc`
---
arrow/ipc/dictionary_payload_lifetime_test.go | 217 ++++++++++++++++++++++++++
arrow/ipc/writer.go | 26 +--
2 files changed, 231 insertions(+), 12 deletions(-)
diff --git a/arrow/ipc/dictionary_payload_lifetime_test.go
b/arrow/ipc/dictionary_payload_lifetime_test.go
new file mode 100644
index 00000000..20e19c8a
--- /dev/null
+++ b/arrow/ipc/dictionary_payload_lifetime_test.go
@@ -0,0 +1,217 @@
+// 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 (
+ "errors"
+ "fmt"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/internal/dictutils"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/require"
+)
+
+type dictionaryTrackingAllocator struct {
+ memory.Allocator
+ current int64
+}
+
+func (a *dictionaryTrackingAllocator) Allocate(size int) []byte {
+ out := a.Allocator.Allocate(size)
+ a.current += int64(size)
+ return out
+}
+
+func (a *dictionaryTrackingAllocator) Reallocate(size int, b []byte) []byte {
+ out := a.Allocator.Reallocate(size, b)
+ a.current += int64(size - len(b))
+ return out
+}
+
+func (a *dictionaryTrackingAllocator) Free(b []byte) {
+ a.current -= int64(len(b))
+ a.Allocator.Free(b)
+}
+
+type dictionaryPayloadTrackingWriter struct {
+ allocator *dictionaryTrackingAllocator
+ peak int64
+ payloads []int64
+}
+
+func (w *dictionaryPayloadTrackingWriter) Start() error { return nil }
+func (w *dictionaryPayloadTrackingWriter) Close() error { return nil }
+
+func (w *dictionaryPayloadTrackingWriter) WritePayload(p Payload) error {
+ if p.msg != MessageDictionaryBatch {
+ return nil
+ }
+
+ current := w.allocator.current
+ w.payloads = append(w.payloads, current)
+ if current > w.peak {
+ w.peak = current
+ }
+ return nil
+}
+
+func newDictionaryRecordBatch(mem memory.Allocator, numColumns, numValues int)
arrow.RecordBatch {
+ dictType := &arrow.DictionaryType{
+ IndexType: arrow.PrimitiveTypes.Int32,
+ ValueType: arrow.PrimitiveTypes.Int64,
+ }
+ fields := make([]arrow.Field, numColumns)
+ for i := range fields {
+ fields[i] = arrow.Field{Name: fmt.Sprintf("dict_%d", i), Type:
dictType}
+ }
+
+ schema := arrow.NewSchema(fields, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ for i := 0; i < numColumns; i++ {
+ column := builder.Field(i).(*array.Int64DictionaryBuilder)
+ column.Reserve(numValues)
+ for j := 0; j < numValues; j++ {
+ column.UnsafeAppend(int64(j))
+ }
+ }
+
+ record := builder.NewRecordBatch()
+ builder.Release()
+ return record
+}
+
+func TestWriteDictionaryPayloadsReleasesPayloadPerDictionary(t *testing.T) {
+ const (
+ numColumns = 8
+ numValues = 4096
+ )
+
+ mem := &dictionaryTrackingAllocator{Allocator: memory.DefaultAllocator}
+ record := newDictionaryRecordBatch(mem, numColumns, numValues)
+ defer record.Release()
+
+ mapper := &dictutils.Mapper{}
+ mapper.ImportSchema(record.Schema())
+ encoder := newRecordEncoder(mem, 0, kMaxNestingDepth, true, -1, 1, 0,
nil)
+ lastWrittenDicts := make(map[int64]arrow.Array, numColumns)
+ defer func() {
+ for _, dict := range lastWrittenDicts {
+ dict.Release()
+ }
+ }()
+
+ writer := &dictionaryPayloadTrackingWriter{allocator: mem}
+ before := mem.current
+ require.NoError(t, writeDictionaryPayloads(mem, record, false, false,
mapper, lastWrittenDicts, writer, encoder))
+ require.Len(t, writer.payloads, numColumns)
+
+ payloadBytes := writer.payloads[0] - before
+ require.Positive(t, payloadBytes)
+ lastPayloadGrowth := writer.payloads[len(writer.payloads)-1] -
writer.payloads[0]
+ require.Less(t, lastPayloadGrowth, payloadBytes)
+}
+
+func TestWriteDictionaryPayloadsReleasesPayloadOnFailure(t *testing.T) {
+ const numColumns = 3
+ for _, delta := range []bool{false, true} {
+ for failAfter := 1; failAfter <= numColumns; failAfter++ {
+ t.Run(fmt.Sprintf("delta=%t/failAfter=%d", delta,
failAfter), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ record := newDictionaryRecordBatch(mem,
numColumns, 64)
+ defer record.Release()
+
+ mapper := &dictutils.Mapper{}
+ mapper.ImportSchema(record.Schema())
+ encoder := newRecordEncoder(mem, 0,
kMaxNestingDepth, true, -1, 1, 0, nil)
+ lastWrittenDicts := make(map[int64]arrow.Array)
+ defer func() {
+ for _, dict := range lastWrittenDicts {
+ dict.Release()
+ }
+ }()
+
+ if delta {
+ previous :=
newDictionaryRecordBatch(mem, numColumns, 32)
+ defer previous.Release()
+ require.NoError(t,
writeDictionaryPayloads(mem, previous, false, true,
+ mapper, lastWrittenDicts,
&failingPayloadWriter{}, encoder))
+ }
+
+ before := mem.CurrentAlloc()
+ want := errors.New("dictionary payload write
failed")
+ writer := &failingPayloadWriter{err: want,
failAfter: failAfter}
+ err := writeDictionaryPayloads(mem, record,
false, delta,
+ mapper, lastWrittenDicts, writer,
encoder)
+ require.ErrorIs(t, err, want)
+ require.Equal(t, failAfter, writer.payloads)
+ require.Equal(t, before, mem.CurrentAlloc())
+ })
+ }
+ }
+}
+
+func BenchmarkWriteDictionaryPayloadsPeak(b *testing.B) {
+ for _, tc := range []struct {
+ name string
+ numColumns int
+ numValues int
+ }{
+ {name: "8cols", numColumns: 8, numValues: 4096},
+ {name: "64cols", numColumns: 64, numValues: 4096},
+ {name: "256cols", numColumns: 256, numValues: 4096},
+ } {
+ tc := tc
+ b.Run(tc.name, func(b *testing.B) {
+ mem := &dictionaryTrackingAllocator{Allocator:
memory.DefaultAllocator}
+ record := newDictionaryRecordBatch(mem, tc.numColumns,
tc.numValues)
+ defer record.Release()
+
+ mapper := &dictutils.Mapper{}
+ mapper.ImportSchema(record.Schema())
+ encoder := newRecordEncoder(mem, 0, kMaxNestingDepth,
true, -1, 1, 0, nil)
+ base := mem.current
+ var peakPayloadBytes int64
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ writer :=
&dictionaryPayloadTrackingWriter{allocator: mem}
+ lastWrittenDicts := make(map[int64]arrow.Array,
tc.numColumns)
+ if err := writeDictionaryPayloads(mem, record,
false, false, mapper, lastWrittenDicts, writer, encoder); err != nil {
+ b.Fatal(err)
+ }
+ if payloadBytes := writer.peak - base;
payloadBytes > peakPayloadBytes {
+ peakPayloadBytes = payloadBytes
+ }
+ for _, dict := range lastWrittenDicts {
+ dict.Release()
+ }
+ if mem.current != base {
+ b.Fatalf("allocator bytes changed from
%d to %d", base, mem.current)
+ }
+ }
+ b.StopTimer()
+ b.ReportMetric(float64(peakPayloadBytes),
"peak-payload-B")
+ })
+ }
+}
+
+var _ memory.Allocator = (*dictionaryTrackingAllocator)(nil)
diff --git a/arrow/ipc/writer.go b/arrow/ipc/writer.go
index a1c768ba..ed8e60a9 100644
--- a/arrow/ipc/writer.go
+++ b/arrow/ipc/writer.go
@@ -291,19 +291,21 @@ func writeDictionaryPayloads(mem memory.Allocator, batch
arrow.RecordBatch, isFi
}
}
- var data = Payload{msg: MessageDictionaryBatch}
- defer data.Release()
-
- dict := pair.Dict
- if deltaStart > 0 {
- dict = array.NewSlice(dict, deltaStart,
int64(dict.Len()))
- defer dict.Release()
- }
- if err := enc.Encode(&data, pair.ID, deltaStart > 0, dict); err
!= nil {
- return err
- }
+ if err := func() error {
+ data := Payload{msg: MessageDictionaryBatch}
+ defer data.Release()
+
+ dict := pair.Dict
+ if deltaStart > 0 {
+ dict = array.NewSlice(dict, deltaStart,
int64(dict.Len()))
+ defer dict.Release()
+ }
+ if err := enc.Encode(&data, pair.ID, deltaStart > 0,
dict); err != nil {
+ return err
+ }
- if err := pw.WritePayload(data); err != nil {
+ return pw.WritePayload(data)
+ }(); err != nil {
return err
}
