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 56815b70 perf(arrow/ipc): reuse zstd encoders across writers (#1429)
56815b70 is described below
commit 56815b707612e05482ea0dc45c3b4bf85e52994c
Author: Arnaud <[email protected]>
AuthorDate: Thu Oct 29 17:44:33 2026 +0100
perf(arrow/ipc): reuse zstd encoders across writers (#1429)
### Rationale for this change
`getCompressor` builds a new `zstd.Encoder` for every writer. A new
encoder allocates its window and match tables, about 2.3 MB at the
default level, before it compresses anything. A service that writes each
table or batch with its own writer pays that on every call. In our case,
one writer per series in each Grafana query, this was about 63% of all
allocations on that path, and the GC load that followed was a third of
the CPU.
#1331 pooled the zstd decoders for the same reason, and
`parquet/compress` already pools its zstd encoders (#717). This does the
same for the IPC writer.
### What changes are included in this PR?
- zstd encoders come from a `sync.Pool`, one per encoder level: the
library default, plus the four `zstd.EncoderLevel` presets that
`WithZstdLevel` maps to. An encoder is only reused for the level it was
built with. The encoder options are unchanged.
- `Writer.Close`, `FileWriter.Close` and `GetRecordBatchPayload` return
the encoders. A writer keeps its encoder for its whole lifetime, so a
long-lived stream sees no change.
- Before going back to the pool, the encoder is `Reset(nil)`. That waits
for any in-flight block goroutines (also after a failed `Write`) and
drops the reference to the last output buffer.
- lz4 is unchanged.
### Are these changes tested?
New `TestZstdEncoderPool`:
- `Close` returns every encoder, for `Writer` and `FileWriter` with
`WithCompressConcurrency(4)`. This fails without the change.
- The output of a reused encoder is byte-identical to its first use,
after a larger stream went through it in between.
- 16 concurrent writers with mixed sizes and compression concurrency
round-trip correctly under `-race`.
`go test -race ./arrow/ipc/... ./arrow/flight/...` passes, and
`golangci-lint run ./arrow/ipc/` is clean. Run on macOS only.
New `BenchmarkWriterZstd`. Old and new test binaries were run
alternately, 8 runs each, on an Apple M2 Pro:
```
sec/op B/op
allocs/op
writer_per_record/256rows 246.4µ → 4.5µ (-98%) 2279Ki → 7.1Ki
(-99.7%) 72 → 44
writer_per_record/4096rows 280.7µ → 21.8µ (-92%) 2296Ki → 24.7Ki
(-98.9%) 72 → 44
one_writer/256rows 2.98µ → 2.98µ (~) 2.755Ki → 2.755Ki (~)
23 → 23
one_writer/4096rows 19.4µ → 19.6µ (~) 19.54Ki → 19.54Ki (~)
23 → 23
```
### Are there any user-facing changes?
No API change. Writers that use zstd allocate less. Idle encoders can
stay in the pool until GC clears it, at most one per level and per P.
---
arrow/ipc/compression.go | 36 ++++++++++++++++-
arrow/ipc/file_writer.go | 1 +
arrow/ipc/writer.go | 2 +
arrow/ipc/writer_benchmark_test.go | 36 +++++++++++++++++
arrow/ipc/writer_test.go | 79 ++++++++++++++++++++++++++++++++++++++
5 files changed, 153 insertions(+), 1 deletion(-)
diff --git a/arrow/ipc/compression.go b/arrow/ipc/compression.go
index 0807e661..7027d448 100644
--- a/arrow/ipc/compression.go
+++ b/arrow/ipc/compression.go
@@ -50,6 +50,36 @@ func (lz4Compressor) Type() flatbuf.CompressionType {
type zstdCompressor struct {
*zstd.Encoder
+ pool interface{ Put(any) }
+}
+
+// zstdEncoderPools holds reusable encoders: index 0 for the library default
+// level, 1-4 for the zstd.EncoderLevel a WithZstdLevel value maps to. A new
+// encoder allocates its window and match tables (about 2 MB at the default
+// level), which dominates the cost of compressing a small record batch.
+var zstdEncoderPools [zstd.SpeedBestCompression + 1]sync.Pool
+
+func zstdEncoderPoolIndex(zstdLevel int) int {
+ if zstdLevel <= 0 {
+ return 0
+ }
+ return int(zstd.EncoderLevelFromZstd(zstdLevel))
+}
+
+// releaseCompressors returns the pooled compressors in cs and clears their
+// slots, so calling it twice on the same slice is safe.
+func releaseCompressors(cs []compressor) {
+ for i, c := range cs {
+ z, ok := c.(*zstdCompressor)
+ if !ok {
+ continue
+ }
+ // Reset waits for in-flight block goroutines, also after a
failed
+ // Write, and drops the reference to the last output buffer.
+ z.Reset(nil)
+ z.pool.Put(z)
+ cs[i] = nil
+ }
}
// from zstd.h, ZSTD_COMPRESSBOUND
@@ -74,6 +104,10 @@ func getCompressor(codec flatbuf.CompressionType, zstdLevel
int) compressor {
w.Apply(lz4.ChecksumOption(false),
lz4.BlockSizeOption(lz4.Block64Kb))
return &lz4Compressor{w}
case flatbuf.CompressionTypeZSTD:
+ pool := &zstdEncoderPools[zstdEncoderPoolIndex(zstdLevel)]
+ if c, ok := pool.Get().(*zstdCompressor); ok {
+ return c
+ }
var enc *zstd.Encoder
var err error
if zstdLevel > 0 {
@@ -86,7 +120,7 @@ func getCompressor(codec flatbuf.CompressionType, zstdLevel
int) compressor {
if err != nil {
panic(err)
}
- return zstdCompressor{enc}
+ return &zstdCompressor{Encoder: enc, pool: pool}
}
return nil
}
diff --git a/arrow/ipc/file_writer.go b/arrow/ipc/file_writer.go
index 12b40168..d446f0f0 100644
--- a/arrow/ipc/file_writer.go
+++ b/arrow/ipc/file_writer.go
@@ -302,6 +302,7 @@ func (f *FileWriter) getRecordEncoder() *recordEncoder {
func (f *FileWriter) Close() error {
defer func() { f.encoder = nil }()
+ defer releaseCompressors(f.compressors)
if f.closed {
return f.closeErr
diff --git a/arrow/ipc/writer.go b/arrow/ipc/writer.go
index 97eb76a8..a32a8530 100644
--- a/arrow/ipc/writer.go
+++ b/arrow/ipc/writer.go
@@ -152,6 +152,7 @@ func (w *Writer) getRecordEncoder() *recordEncoder {
func (w *Writer) Close() error {
defer func() { w.encoder = nil }()
+ defer releaseCompressors(w.compressors)
if w.err != nil {
return w.closeAfterFailure()
@@ -1161,6 +1162,7 @@ func GetRecordBatchPayload(batch arrow.RecordBatch, opts
...Option) (Payload, er
cfg := newConfig(opts...)
compressors := make([]compressor, cfg.compressNP)
copy(compressors, cfg.compressors)
+ defer releaseCompressors(compressors)
var (
data = Payload{msg: MessageRecordBatch}
enc = newRecordEncoder(
diff --git a/arrow/ipc/writer_benchmark_test.go
b/arrow/ipc/writer_benchmark_test.go
index 3567ec4f..adb11b2d 100644
--- a/arrow/ipc/writer_benchmark_test.go
+++ b/arrow/ipc/writer_benchmark_test.go
@@ -79,3 +79,39 @@ func BenchmarkWriterRecordEncoderReuse(b *testing.B) {
}
}
}
+
+// BenchmarkWriterZstd measures a fresh compressed writer per record, as when
+// each table or Flight batch gets its own writer, and one writer for a stream.
+func BenchmarkWriterZstd(b *testing.B) {
+ for _, numRows := range []int{256, 4096} {
+ record := benchmarkRecordBatch(1, numRows)
+ defer record.Release()
+
+ b.Run(fmt.Sprintf("writer_per_record/%drows", numRows), func(b
*testing.B) {
+ b.ReportAllocs()
+ b.SetBytes(int64(numRows * arrow.Int32SizeBytes))
+ for b.Loop() {
+ writer := NewWriter(io.Discard,
WithSchema(record.Schema()), WithZstd())
+ if err := writer.Write(record); err != nil {
+ b.Fatal(err)
+ }
+ if err := writer.Close(); err != nil {
+ b.Fatal(err)
+ }
+ }
+ })
+
+ b.Run(fmt.Sprintf("one_writer/%drows", numRows), func(b
*testing.B) {
+ writer := NewWriter(io.Discard,
WithSchema(record.Schema()), WithZstd())
+ defer writer.Close()
+
+ b.ReportAllocs()
+ b.SetBytes(int64(numRows * arrow.Int32SizeBytes))
+ for b.Loop() {
+ if err := writer.Write(record); err != nil {
+ b.Fatal(err)
+ }
+ }
+ })
+ }
+}
diff --git a/arrow/ipc/writer_test.go b/arrow/ipc/writer_test.go
index 2a74a0aa..0b38bd68 100644
--- a/arrow/ipc/writer_test.go
+++ b/arrow/ipc/writer_test.go
@@ -24,6 +24,7 @@ import (
"io"
"math"
"strings"
+ "sync"
"testing"
"time"
@@ -1184,3 +1185,81 @@ func TestWithZstdLevel(t *testing.T) {
assert.True(t, array.RecordEqual(rec, r.RecordBatch()),
"records differ")
})
}
+
+func TestZstdEncoderPool(t *testing.T) {
+ small := benchmarkRecordBatch(1, 16)
+ defer small.Release()
+ large := benchmarkRecordBatch(8, 8192)
+ defer large.Release()
+
+ // assert, not require: the concurrent subtest calls this from other
goroutines.
+ write := func(t *testing.T, rec arrow.RecordBatch, opts ...Option)
[]byte {
+ t.Helper()
+ var buf bytes.Buffer
+ w := NewWriter(&buf, append([]Option{WithSchema(rec.Schema()),
WithZstd()}, opts...)...)
+ assert.NoError(t, w.Write(rec))
+ assert.NoError(t, w.Close())
+ return buf.Bytes()
+ }
+
+ t.Run("Close returns the encoders", func(t *testing.T) {
+ var buf bytes.Buffer
+ w := NewWriter(&buf, WithSchema(large.Schema()), WithZstd(),
WithCompressConcurrency(4))
+ require.NoError(t, w.Write(large))
+ require.NoError(t, w.Close())
+ for i, c := range w.compressors {
+ assert.Nil(t, c, "Writer compressor %d was not returned
to the pool", i)
+ }
+
+ fw, err := NewFileWriter(&buf, WithSchema(large.Schema()),
WithZstd(), WithCompressConcurrency(4))
+ require.NoError(t, err)
+ require.NoError(t, fw.Write(large))
+ require.NoError(t, fw.Close())
+ for i, c := range fw.compressors {
+ assert.Nil(t, c, "FileWriter compressor %d was not
returned to the pool", i)
+ }
+ })
+
+ t.Run("Close puts the encoder back in its pool", func(t *testing.T) {
+ // A fake pool, because sync.Pool may drop any item it is given.
+ w := NewWriter(io.Discard, WithSchema(small.Schema()),
WithZstd())
+ require.NoError(t, w.Write(small))
+ z := w.compressors[0].(*zstdCompressor)
+ pool := &recordingPool{}
+ z.pool = pool
+ require.NoError(t, w.Close())
+ assert.Equal(t, []any{z}, pool.put)
+ })
+
+ t.Run("reused encoder output does not depend on earlier streams",
func(t *testing.T) {
+ first := write(t, small)
+ _ = write(t, large)
+ again := write(t, small)
+ assert.Equal(t, first, again)
+ })
+
+ t.Run("concurrent writers round trip", func(t *testing.T) {
+ var wg sync.WaitGroup
+ for i := range 16 {
+ wg.Go(func() {
+ rec := small
+ if i%2 == 0 {
+ rec = large
+ }
+ r, err := NewReader(bytes.NewReader(write(t,
rec, WithCompressConcurrency(1+i%3))))
+ if !assert.NoError(t, err) {
+ return
+ }
+ defer r.Release()
+ if assert.True(t, r.Next()) {
+ assert.True(t, array.RecordEqual(rec,
r.RecordBatch()))
+ }
+ })
+ }
+ wg.Wait()
+ })
+}
+
+type recordingPool struct{ put []any }
+
+func (p *recordingPool) Put(x any) { p.put = append(p.put, x) }
