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 7c13a602 perf: zstd implement Pooling + DecodeAll with no new reader.
(#1331)
7c13a602 is described below
commit 7c13a602d47616d45ccbbb3f845f3ca44a5926b6
Author: Samuele Resca <[email protected]>
AuthorDate: Mon Oct 12 15:02:16 2026 +0100
perf: zstd implement Pooling + DecodeAll with no new reader. (#1331)
### Rationale for this change
Performance improvement on the zstd path.
### What changes are included in this PR?
- zstd: decoders come from a `sync.Pool` instead of being created for
every record batch, and each compressed buffer is decoded with one
`DecodeAll` call straight into the destination buffer (decoder built
with `WithDecoderConcurrency(1)` and `WithDecodeAllCapLimit(true))`
instead of `Reset(bytes.NewReader)` + `io.ReadFull` on a streaming
decoder.
- lz4 : no functional change. It keeps its original per-batch
`lz4.Reader`; the library already pools its large block buffers
internally. (Only a Decompress method wrapper was added.)
### Are these changes tested?
Added tests + benchmark for zstd (the changed path) - every case
improves
```
SMALL buffers (4 KiB per column)
time : -70.5% to -78.3% in all 8 cases
1 col, 1 batch (stream) 9.89 us -> 2.92 us
1 col, 64 batches (stream) 592.9 us -> 128.7 us
16 cols, 64 batches (file) 3.67 ms -> 894 us
B/op : -60% (1 col, 1 batch), -64% (1 col, 64 batches), -22% (16 cols)
allocs/op : -45% (1 col, 1 batch) to -70% (16 cols, 64 batches)
LARGE buffers
per column time B/op
allocs/op
16 KiB 77.6 us -> 24.0 us (-69%) 306.9 KiB -> 157.1 KiB (-49%)
184 -> 78 (-58%)
256 KiB 796.7 us -> 191.7 us (-76%) 20.6 MiB -> 2.05 MiB (-90%)
185 -> 81 (-56%)
4 MiB 4.63 ms -> 3.05 ms (-34%) 51.1 MiB -> 32.1 MiB (-37%)
189 -> 79 (-58%)
```
### Are there any user-facing changes?
Strictness change for zstd. A buffer whose decompressed size differs
from the uncompressed-length prefix now returns an error. This matches
other arrow libraries (e.g.:
[arrow-cpp](https://github.com/apache/arrow/blob/f3b46fb7d4b8f5fb3fc459860038f0ed8b02bba5/cpp/src/arrow/ipc/reader.cc#L571-L575))
---
arrow/ipc/compression.go | 83 +++++--
arrow/ipc/compression_test.go | 288 +++++++++++++++++++++++++
arrow/ipc/file_reader.go | 4 +-
arrow/ipc/reader_compression_benchmark_test.go | 154 +++++++++++++
4 files changed, 515 insertions(+), 14 deletions(-)
diff --git a/arrow/ipc/compression.go b/arrow/ipc/compression.go
index abdf6313..81bd6fe6 100644
--- a/arrow/ipc/compression.go
+++ b/arrow/ipc/compression.go
@@ -17,7 +17,10 @@
package ipc
import (
+ "bytes"
+ "fmt"
"io"
+ "sync"
"github.com/apache/arrow-go/v18/arrow/internal/debug"
"github.com/apache/arrow-go/v18/arrow/internal/flatbuf"
@@ -81,27 +84,87 @@ func getCompressor(codec flatbuf.CompressionType)
compressor {
}
type decompressor interface {
- io.Reader
- Reset(io.Reader)
+ // Decompress decodes src into dst, which must be sized to the exact
+ // uncompressed length.
+ Decompress(dst, src []byte) error
+ // Close releases the decompressor; it must not be used afterwards.
+ // Implementations may return it to a pool, so calling Close twice hands
+ // the same decompressor to two owners.
Close()
}
+var zstdDecompressorPool = sync.Pool{
+ New: func() any {
+ // WithDecoderConcurrency(1): Each pooled decoder is used by
one goroutine at a time, so a single
+ // block decoder is enough. The default would create up to four
that
+ // could never run in parallel;
+ //
+ // WithDecodeAllCapLimit(true): The cap limit bounds DecodeAll
to cap(dst), so a frame claiming a
+ // larger content size can't allocate beyond the buffer the
caller
+ // sized from the uncompressed length prefix.
+ dec, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1),
zstd.WithDecodeAllCapLimit(true))
+ if err != nil {
+ panic(err)
+ }
+ return &zstdDecompressor{Decoder: dec}
+ },
+}
+
type zstdDecompressor struct {
*zstd.Decoder
}
-func (z *zstdDecompressor) Reset(r io.Reader) {
- if err := z.Decoder.Reset(r); err != nil {
- panic(err)
+func (z *zstdDecompressor) Decompress(dst, src []byte) error {
+ if len(dst) == 0 {
+ return nil
}
+
+ // The decoder was created with WithDecodeAllCapLimit, so it decodes
into
+ // dst's own capacity and fails rather than allocating a larger slice.
+ out, err := z.DecodeAll(src, dst[:0])
+ if err != nil {
+ return err
+ }
+ // Catch cases where the prefix says fewer bytes than the content, but
the content fits in the dst's spare capacity
+ if len(out) != len(dst) {
+ return fmt.Errorf("arrow/ipc: zstd decompressed to %d bytes,
expected %d", len(out), len(dst))
+ }
+ return nil
}
+// Close returns the decoder to the pool. z.Decoder.Close() is deliberately
+// not called: zstd.NewReader(nil) starts no goroutines, so a pooled decoder
+// only holds GC-reclaimable state and there is nothing to tear down. Closing
+// it would make the decoder unusable for the next borrower.
func (z *zstdDecompressor) Close() {
- z.Decoder.Close()
+ zstdDecompressorPool.Put(z)
}
type lz4Decompressor struct {
*lz4.Reader
+ extra [1]byte // probe buffer for trailing data
+}
+
+func (z *lz4Decompressor) Decompress(dst, src []byte) error {
+ if len(dst) == 0 {
+ return nil
+ }
+
+ z.Reset(bytes.NewReader(src))
+ if _, err := io.ReadFull(z.Reader, dst); err != nil {
+ return err
+ }
+ // ReadFull stops once dst is full, so probe for one more byte: if
there is
+ // one, the frame decompresses to more than the length prefix claims,
and we
+ // reject it rather than silently truncate (mirrors the zstd path).
+ n, err := z.Read(z.extra[:])
+ if n > 0 {
+ return fmt.Errorf("arrow/ipc: lz4 decompressed to more than the
expected %d bytes", len(dst))
+ }
+ if err != io.EOF {
+ return err
+ }
+ return nil
}
func (z *lz4Decompressor) Close() {
@@ -111,13 +174,9 @@ func (z *lz4Decompressor) Close() {
func getDecompressor(codec flatbuf.CompressionType) decompressor {
switch codec {
case flatbuf.CompressionTypeLZ4_FRAME:
- return &lz4Decompressor{lz4.NewReader(nil)}
+ return &lz4Decompressor{Reader: lz4.NewReader(nil)}
case flatbuf.CompressionTypeZSTD:
- dec, err := zstd.NewReader(nil)
- if err != nil {
- panic(err)
- }
- return &zstdDecompressor{dec}
+ return zstdDecompressorPool.Get().(*zstdDecompressor)
}
return nil
}
diff --git a/arrow/ipc/compression_test.go b/arrow/ipc/compression_test.go
new file mode 100644
index 00000000..afc997f6
--- /dev/null
+++ b/arrow/ipc/compression_test.go
@@ -0,0 +1,288 @@
+// 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 (
+ "bytes"
+ "encoding/binary"
+ "errors"
+ "fmt"
+ "math/rand"
+ "runtime"
+ "sync"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/internal/flatbuf"
+ "github.com/klauspost/compress/zstd"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+var testCodecs = []struct {
+ name string
+ codec flatbuf.CompressionType
+ opt Option
+}{
+ {"zstd", flatbuf.CompressionTypeZSTD, WithZstd()},
+ {"lz4", flatbuf.CompressionTypeLZ4_FRAME, WithLZ4()},
+}
+
+func compressBuffer(t *testing.T, codec flatbuf.CompressionType, src []byte)
[]byte {
+ t.Helper()
+
+ var out bytes.Buffer
+ c := getCompressor(codec)
+ c.Reset(&out)
+ _, err := c.Write(src)
+ require.NoError(t, err)
+ require.NoError(t, c.Close())
+ return out.Bytes()
+}
+
+// compressibleBytes returns n deterministic bytes drawn from a small alphabet,
+// so they compress but aren't trivially uniform.
+func compressibleBytes(n int, seed int64) []byte {
+ rng := rand.New(rand.NewSource(seed))
+ out := make([]byte, n)
+ for i := range out {
+ out[i] = byte(rng.Intn(16))
+ }
+ return out
+}
+
+func TestDecompressorRoundTrip(t *testing.T) {
+ for _, tc := range testCodecs {
+ t.Run(tc.name, func(t *testing.T) {
+ dec := getDecompressor(tc.codec)
+ defer dec.Close()
+
+ for i, n := range []int{1, 100, 4096, 300_000, 17,
300_000} {
+ want := compressibleBytes(n, int64(i))
+ got := make([]byte, n)
+ require.NoError(t, dec.Decompress(got,
compressBuffer(t, tc.codec, want)), "size %d", n)
+ assert.Equal(t, want, got, "size %d", n)
+ }
+ })
+ }
+}
+
+func TestDecompressorEmptyDestination(t *testing.T) {
+ for _, tc := range testCodecs {
+ t.Run(tc.name, func(t *testing.T) {
+ dec := getDecompressor(tc.codec)
+ defer dec.Close()
+ assert.NoError(t, dec.Decompress(nil, nil))
+ assert.NoError(t, dec.Decompress([]byte{}, []byte{1, 2,
3}))
+ })
+ }
+}
+
+func TestDecompressorSizeMismatch(t *testing.T) {
+ want := compressibleBytes(1000, 1)
+
+ for _, tc := range testCodecs {
+ t.Run(tc.name, func(t *testing.T) {
+ dec := getDecompressor(tc.codec)
+ defer dec.Close()
+ compressed := compressBuffer(t, tc.codec, want)
+
+ // the uncompressed length prefix claims more than the
frame holds
+ assert.Error(t, dec.Decompress(make([]byte,
len(want)+1), compressed))
+
+ // and the decoder is still usable afterwards
+ got := make([]byte, len(want))
+ require.NoError(t, dec.Decompress(got, compressed))
+ assert.Equal(t, want, got)
+ })
+ }
+
+ // zstd frames declare their content size, so a prefix that is too small
+ // is caught instead of silently truncating.
+ t.Run("zstd/short-destination", func(t *testing.T) {
+ dec := getDecompressor(flatbuf.CompressionTypeZSTD)
+ defer dec.Close()
+ assert.Error(t, dec.Decompress(make([]byte, len(want)-1),
compressBuffer(t, flatbuf.CompressionTypeZSTD, want)))
+ })
+}
+
+func TestZstdDecompressorBoundedByDestination(t *testing.T) {
+ // 64 MiB of zeros compresses to a few bytes
+ compressed := compressBuffer(t, flatbuf.CompressionTypeZSTD,
make([]byte, 64<<20))
+ require.Less(t, len(compressed), 64<<10)
+
+ dec := getDecompressor(flatbuf.CompressionTypeZSTD)
+ defer dec.Close()
+
+ dst := make([]byte, 16)
+ err := dec.Decompress(dst, compressed)
+ require.Error(t, err)
+ assert.True(t, errors.Is(err, zstd.ErrDecoderSizeExceeded), "unexpected
error: %v", err)
+}
+
+// A frame header can declare the content size up front. A hostile frame can
+// declare far more than it holds, hoping the decoder allocates that much
+// before it ever reads the (tiny) body. The cap limit has to reject it up
+// front, without allocating: check both the error and the allocation, since a
+// decoder that allocated first and failed afterwards would return an error
too.
+func TestZstdDecompressorRejectsHugeDeclaredSizeWithoutAllocating(t
*testing.T) {
+ const declared = 64 << 20
+
+ // Hand-built 17-byte zstd frame that is well formed but lies about its
+ // size: the header declares 64 MiB of content, the frame holds one
byte.
+ frame := []byte{0x28, 0xb5, 0x2f, 0xfd, 0xe0}
+ frame = binary.LittleEndian.AppendUint64(frame, declared)
+ frame = append(frame, 0x09, 0x00, 0x00, 0x41)
+
+ var hdr zstd.Header
+ require.NoError(t, hdr.Decode(frame))
+ require.True(t, hdr.HasFCS)
+ require.Equal(t, uint64(declared), hdr.FrameContentSize)
+
+ dec := getDecompressor(flatbuf.CompressionTypeZSTD)
+ defer dec.Close()
+
+ // warm up, so decoder state that is created lazily isn't counted below
+ dst := make([]byte, 100)
+ require.NoError(t, dec.Decompress(dst, compressBuffer(t,
flatbuf.CompressionTypeZSTD, compressibleBytes(100, 3))))
+
+ // collect garbage first so GC bookkeeping stays out of the measured
window
+ runtime.GC()
+
+ var before, after runtime.MemStats
+ runtime.ReadMemStats(&before)
+ err := dec.Decompress(dst, frame)
+ runtime.ReadMemStats(&after)
+
+ require.Error(t, err)
+ assert.True(t, errors.Is(err, zstd.ErrDecoderSizeExceeded), "unexpected
error: %v", err)
+ assert.Less(t, after.TotalAlloc-before.TotalAlloc, uint64(1<<20),
+ "decoding a frame declaring %d bytes allocated too much",
declared)
+}
+
+// Reader destinations are allocated in 64-byte aligned blocks, so they usually
+// have spare capacity past their length (e.g. len 100 has cap 128). Decoding
+// must fill exactly len(dst) and reject frames of any other size, including
+// ones that would fit in the spare capacity.
+func TestZstdDecompressorSpareCapacity(t *testing.T) {
+ const n, capacity = 100, 128
+
+ dec := getDecompressor(flatbuf.CompressionTypeZSTD)
+ defer dec.Close()
+
+ compress := func(size int) []byte {
+ return compressBuffer(t, flatbuf.CompressionTypeZSTD,
bytes.Repeat([]byte{0xAB}, size))
+ }
+
+ t.Run("valid frame fills dst and leaves the spare capacity untouched",
func(t *testing.T) {
+ dst := make([]byte, n, capacity)
+ require.NoError(t, dec.Decompress(dst, compress(n)))
+ assert.Equal(t, bytes.Repeat([]byte{0xAB}, n), dst)
+ assert.Equal(t, make([]byte, capacity-n), dst[n:capacity])
+ })
+
+ t.Run("frame larger than dst but within the spare capacity is
rejected", func(t *testing.T) {
+ dst := make([]byte, n, capacity)
+ assert.Error(t, dec.Decompress(dst, compress(capacity-1)))
+ })
+
+ t.Run("frame larger than the spare capacity is rejected", func(t
*testing.T) {
+ dst := make([]byte, n, capacity)
+ err := dec.Decompress(dst, compress(capacity+1))
+ require.Error(t, err)
+ assert.True(t, errors.Is(err, zstd.ErrDecoderSizeExceeded),
"unexpected error: %v", err)
+ })
+
+ t.Run("decoder is still usable after the rejected frames", func(t
*testing.T) {
+ dst := make([]byte, n, capacity)
+ require.NoError(t, dec.Decompress(dst, compress(n)))
+ assert.Equal(t, bytes.Repeat([]byte{0xAB}, n), dst)
+ })
+}
+
+func TestDecompressorRecoversAfterTruncatedInput(t *testing.T) {
+ want := compressibleBytes(4096, 2)
+
+ for _, tc := range testCodecs {
+ t.Run(tc.name, func(t *testing.T) {
+ dec := getDecompressor(tc.codec)
+ defer dec.Close()
+ compressed := compressBuffer(t, tc.codec, want)
+
+ // lz4 frames are written without checksums, so flipped
bytes can go
+ // undetected; a truncated frame is always an error.
+ truncated := compressed[:len(compressed)/2]
+ assert.Error(t, dec.Decompress(make([]byte, len(want)),
truncated))
+
+ got := make([]byte, len(want))
+ require.NoError(t, dec.Decompress(got, compressed))
+ assert.Equal(t, want, got)
+ })
+ }
+}
+
+func TestReadCompressedConcurrently(t *testing.T) {
+ for _, tc := range testCodecs {
+ t.Run(tc.name, func(t *testing.T) {
+ rec := benchmarkRecordBatch(4, 2048)
+ defer rec.Release()
+
+ var buf bytes.Buffer
+ w := NewWriter(&buf, WithSchema(rec.Schema()), tc.opt)
+ for range 8 {
+ require.NoError(t, w.Write(rec))
+ }
+ require.NoError(t, w.Close())
+ data := buf.Bytes()
+
+ var wg sync.WaitGroup
+ errs := make(chan error, 16)
+ for range 16 {
+ wg.Go(func() {
+ for range 10 {
+ rdr, err :=
NewReader(bytes.NewReader(data))
+ if err != nil {
+ errs <- err
+ return
+ }
+ n := 0
+ for rdr.Next() {
+ if
!array.RecordEqual(rec, rdr.RecordBatch()) {
+ errs <-
fmt.Errorf("record %d differs", n)
+ }
+ n++
+ }
+ err = rdr.Err()
+ rdr.Release()
+ if err == nil && n != 8 {
+ err = fmt.Errorf("read
%d batches, want 8", n)
+ }
+ if err != nil {
+ errs <- err
+ return
+ }
+ }
+ })
+ }
+ wg.Wait()
+ close(errs)
+ for err := range errs {
+ assert.NoError(t, err)
+ }
+ })
+ }
+}
diff --git a/arrow/ipc/file_reader.go b/arrow/ipc/file_reader.go
index aba57365..afcc20b3 100644
--- a/arrow/ipc/file_reader.go
+++ b/arrow/ipc/file_reader.go
@@ -601,8 +601,8 @@ func (src *ipcSource) buffer(i int) *memory.Buffer {
if uncompressedSize != -1 {
raw = memory.NewResizableBuffer(src.mem)
raw.Resize(int(uncompressedSize))
- src.codec.Reset(bytes.NewReader(body[8:]))
- if _, err := io.ReadFull(src.codec, raw.Bytes()); err
!= nil {
+ if err := src.codec.Decompress(raw.Bytes(), body[8:]);
err != nil {
+ raw.Release()
panic(err)
}
} else {
diff --git a/arrow/ipc/reader_compression_benchmark_test.go
b/arrow/ipc/reader_compression_benchmark_test.go
new file mode 100644
index 00000000..cfa174b9
--- /dev/null
+++ b/arrow/ipc/reader_compression_benchmark_test.go
@@ -0,0 +1,154 @@
+// 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 (
+ "bytes"
+ "fmt"
+ "io"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+)
+
+func writeBenchmarkStream(b *testing.B, rec arrow.RecordBatch, numBatches int,
opts ...Option) []byte {
+ b.Helper()
+
+ var buf bytes.Buffer
+ w := NewWriter(&buf, append([]Option{WithSchema(rec.Schema())},
opts...)...)
+ for range numBatches {
+ if err := w.Write(rec); err != nil {
+ b.Fatal(err)
+ }
+ }
+ if err := w.Close(); err != nil {
+ b.Fatal(err)
+ }
+ return buf.Bytes()
+}
+
+func writeBenchmarkFile(b *testing.B, rec arrow.RecordBatch, numBatches int,
opts ...Option) []byte {
+ b.Helper()
+
+ var buf bytes.Buffer
+ w, err := NewFileWriter(&buf,
append([]Option{WithSchema(rec.Schema())}, opts...)...)
+ if err != nil {
+ b.Fatal(err)
+ }
+ for range numBatches {
+ if err := w.Write(rec); err != nil {
+ b.Fatal(err)
+ }
+ }
+ if err := w.Close(); err != nil {
+ b.Fatal(err)
+ }
+ return buf.Bytes()
+}
+
+// BenchmarkReaderCompressed measures reading IPC data whose body buffers are
+// compressed, for both the stream Reader and the FileReader.
+func BenchmarkReaderCompressed(b *testing.B) {
+ const numRows = 1024
+
+ codecs := []struct {
+ name string
+ opt Option
+ }{
+ {name: "none"},
+ {name: "zstd", opt: WithZstd()},
+ {name: "lz4", opt: WithLZ4()},
+ }
+
+ for _, codec := range codecs {
+ for _, numColumns := range []int{1, 16} {
+ for _, numBatches := range []int{1, 64} {
+ b.Run(fmt.Sprintf("codec=%s/%dcols/%dbatches",
codec.name, numColumns, numBatches), func(b *testing.B) {
+ rec := benchmarkRecordBatch(numColumns,
numRows)
+ defer rec.Release()
+
+ var opts []Option
+ if codec.opt != nil {
+ opts = append(opts, codec.opt)
+ }
+
+ rawBytes := int64(numColumns * numRows
* arrow.Int32SizeBytes * numBatches)
+
+ b.Run("stream", func(b *testing.B) {
+ data := writeBenchmarkStream(b,
rec, numBatches, opts...)
+
+ b.ReportAllocs()
+ b.SetBytes(rawBytes)
+ for b.Loop() {
+ rdr, err :=
NewReader(bytes.NewReader(data))
+ if err != nil {
+ b.Fatal(err)
+ }
+ n := 0
+ for rdr.Next() {
+ n++
+ }
+ err = rdr.Err()
+ rdr.Release()
+ if err != nil {
+ b.Fatal(err)
+ }
+ if n != numBatches {
+ b.Fatalf("read
%d batches, want %d", n, numBatches)
+ }
+ }
+ b.StopTimer()
+
b.ReportMetric(float64(len(data))/float64(rawBytes), "size/raw")
+ })
+
+ b.Run("file", func(b *testing.B) {
+ data := writeBenchmarkFile(b,
rec, numBatches, opts...)
+
+ b.ReportAllocs()
+ b.SetBytes(rawBytes)
+ for b.Loop() {
+ rdr, err :=
NewFileReader(bytes.NewReader(data))
+ if err != nil {
+ b.Fatal(err)
+ }
+ n := 0
+ for {
+ _, err =
rdr.Read()
+ if err != nil {
+ break
+ }
+ n++
+ }
+ closeErr := rdr.Close()
+ if err != io.EOF {
+ b.Fatal(err)
+ }
+ if closeErr != nil {
+
b.Fatal(closeErr)
+ }
+ if n != numBatches {
+ b.Fatalf("read
%d batches, want %d", n, numBatches)
+ }
+ }
+ b.StopTimer()
+
b.ReportMetric(float64(len(data))/float64(rawBytes), "size/raw")
+ })
+ })
+ }
+ }
+ }
+}
