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 08c90684 perf(parquet/compress): reuse Brotli writers across encodes
(#1254)
08c90684 is described below
commit 08c906849b23edb15e05b5840608e96c69b3bff1
Author: Minh Vu <[email protected]>
AuthorDate: Wed Sep 2 19:55:20 2026 +0200
perf(parquet/compress): reuse Brotli writers across encodes (#1254)
## Summary
- **Reuse Brotli writers** for one-shot `Encode` and `EncodeLevel`
calls.
- Keep a separate `sync.Pool` for each supported compression level.
- Reset pooled writers with a nil destination before putting them back.
- Keep the streaming writer APIs unchanged.
- Add serial and parallel benchmarks plus sequential and concurrent
correctness tests.
## Benchmark
Compared with `main` at `6b039a76`:
| Case | Main | This PR | Main allocations | This PR allocations |
| --- | ---: | ---: | ---: | ---: |
| compressible / 256 KiB | ~0.78 ms/op | ~0.51 ms/op | ~22 allocs/op,
~11.9 MB/op | 1 alloc/op, ~10 KB/op |
| semi-random / 256 KiB | ~2.04 ms/op | ~1.77 ms/op | ~21 allocs/op,
~12.0 MB/op | 1 alloc/op, ~37 KB/op |
The parallel benchmark also stays at one allocation per operation in the
pooled path.
## Checks
- `go test ./parquet/compress -count=1`
- `go test -race ./parquet/compress`
- `go vet ./parquet/compress`
- Parquet package tests with `PARQUET_TEST_DATA` set
---
parquet/compress/brotli.go | 55 +++++++++++++++-
parquet/compress/brotli_benchmark_test.go | 105 ++++++++++++++++++++++++++++++
parquet/compress/brotli_internal_test.go | 63 ++++++++++++++++++
parquet/compress/compress_test.go | 60 +++++++++++++++++
4 files changed, 281 insertions(+), 2 deletions(-)
diff --git a/parquet/compress/brotli.go b/parquet/compress/brotli.go
index fdc0346e..ca058ed1 100644
--- a/parquet/compress/brotli.go
+++ b/parquet/compress/brotli.go
@@ -27,6 +27,24 @@ import (
type brotliCodec struct{}
+const (
+ // Brotli qualities above the default retain substantially larger
encoder
+ // workspaces. Keep those writers transient so a high-quality request
does
+ // not permanently increase the process memory footprint.
+ maxPooledBrotliLevel = brotli.DefaultCompression
+ brotliWriterPoolSize = 1
+)
+
+var brotliWriterPools = newBrotliWriterPools()
+
+func newBrotliWriterPools() [brotli.BestCompression + 1]chan *brotli.Writer {
+ var pools [brotli.BestCompression + 1]chan *brotli.Writer
+ for level := brotli.BestSpeed; level <= maxPooledBrotliLevel; level++ {
+ pools[level] = make(chan *brotli.Writer, brotliWriterPoolSize)
+ }
+ return pools
+}
+
func (brotliCodec) NewReader(r io.Reader) io.ReadCloser {
return io.NopCloser(brotli.NewReader(r))
}
@@ -41,15 +59,48 @@ func (b brotliCodec) EncodeLevel(dst, src []byte, level
int) []byte {
dst = make([]byte, 0, maxlen)
}
buf := bytes.NewBuffer(dst[:0])
- w := brotli.NewWriterLevel(buf, level)
+ pool := brotliWriterPool(level)
+ var w *brotli.Writer
+ if pool != nil {
+ select {
+ case w = <-pool:
+ w.Reset(buf)
+ default:
+ }
+ }
+ if w == nil {
+ w = brotli.NewWriterLevel(buf, level)
+ }
_, err := w.Write(src)
if err != nil {
+ releaseBrotliWriter(pool, w)
panic(err)
}
if err := w.Close(); err != nil {
+ releaseBrotliWriter(pool, w)
panic(err)
}
- return buf.Bytes()
+ compressed := buf.Bytes()
+ releaseBrotliWriter(pool, w)
+ return compressed
+}
+
+func brotliWriterPool(level int) chan *brotli.Writer {
+ if level < brotli.BestSpeed || level > maxPooledBrotliLevel {
+ return nil
+ }
+ return brotliWriterPools[level]
+}
+
+func releaseBrotliWriter(pool chan *brotli.Writer, w *brotli.Writer) {
+ if pool == nil {
+ return
+ }
+ w.Reset(nil)
+ select {
+ case pool <- w:
+ default:
+ }
}
func (b brotliCodec) Encode(dst, src []byte) []byte {
diff --git a/parquet/compress/brotli_benchmark_test.go
b/parquet/compress/brotli_benchmark_test.go
new file mode 100644
index 00000000..c1071550
--- /dev/null
+++ b/parquet/compress/brotli_benchmark_test.go
@@ -0,0 +1,105 @@
+// 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 compress_test
+
+import (
+ "fmt"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/parquet/compress"
+)
+
+func BenchmarkBrotliEncodeLevel(b *testing.B) {
+ dataCases := []struct {
+ name string
+ data []byte
+ }{
+ {name: "compressible/64KiB", data: makeCompressibleData(64 *
1024)},
+ {name: "compressible/256KiB", data: makeCompressibleData(256 *
1024)},
+ {name: "compressible/1MiB", data: makeCompressibleData(1024 *
1024)},
+ {name: "semi-random/64KiB", data: makeSemiRandomBrotliData(64 *
1024)},
+ {name: "semi-random/256KiB", data: makeSemiRandomBrotliData(256
* 1024)},
+ {name: "semi-random/1MiB", data: makeSemiRandomBrotliData(1024
* 1024)},
+ }
+ levels := []struct {
+ name string
+ level int
+ }{
+ {name: "level=1", level: 1},
+ {name: "level=default", level:
compress.DefaultCompressionLevel},
+ {name: "level=9", level: 9},
+ {name: "level=11", level: 11},
+ }
+
+ codec, err := compress.GetCodec(compress.Codecs.Brotli)
+ if err != nil {
+ b.Fatal(err)
+ }
+
+ for _, dataCase := range dataCases {
+ for _, level := range levels {
+ b.Run(fmt.Sprintf("%s/%s", dataCase.name, level.name),
func(b *testing.B) {
+ dst := make([]byte,
int(codec.CompressBound(int64(len(dataCase.data)))))
+ b.SetBytes(int64(len(dataCase.data)))
+ b.ReportAllocs()
+ b.ResetTimer()
+ for range b.N {
+ codec.EncodeLevel(dst, dataCase.data,
level.level)
+ }
+ })
+ }
+ }
+}
+
+func BenchmarkBrotliEncodeLevelParallel(b *testing.B) {
+ data := makeSemiRandomBrotliData(256 * 1024)
+ codec, err := compress.GetCodec(compress.Codecs.Brotli)
+ if err != nil {
+ b.Fatal(err)
+ }
+
+ for _, level := range []struct {
+ name string
+ level int
+ }{
+ {name: "level=1", level: 1},
+ {name: "level=default", level:
compress.DefaultCompressionLevel},
+ {name: "level=9", level: 9},
+ {name: "level=11", level: 11},
+ } {
+ b.Run(level.name, func(b *testing.B) {
+ b.SetBytes(int64(len(data)))
+ b.ReportAllocs()
+ b.ResetTimer()
+ b.RunParallel(func(pb *testing.PB) {
+ dst := make([]byte,
int(codec.CompressBound(int64(len(data)))))
+ for pb.Next() {
+ codec.EncodeLevel(dst, data,
level.level)
+ }
+ })
+ })
+ }
+}
+
+func makeSemiRandomBrotliData(size int) []byte {
+ data := makeRandomData(size)
+ pattern := []byte("parquet-page-data-pattern-0123456789abcdef")
+ for i := 0; i < len(data)/4; i += len(pattern) {
+ copy(data[i:], pattern)
+ }
+ return data
+}
diff --git a/parquet/compress/brotli_internal_test.go
b/parquet/compress/brotli_internal_test.go
new file mode 100644
index 00000000..29946c08
--- /dev/null
+++ b/parquet/compress/brotli_internal_test.go
@@ -0,0 +1,63 @@
+// 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 compress
+
+import (
+ "bytes"
+ "testing"
+
+ "github.com/andybalholm/brotli"
+ "github.com/stretchr/testify/assert"
+)
+
+func TestBrotliWriterPoolMemoryPolicy(t *testing.T) {
+ for level := brotli.BestSpeed; level <= brotli.BestCompression; level++
{
+ pool := brotliWriterPool(level)
+ if level <= maxPooledBrotliLevel {
+ assert.NotNil(t, pool, "level %d should be pooled",
level)
+ assert.Equal(t, brotliWriterPoolSize, cap(pool), "level
%d pool size", level)
+ } else {
+ assert.Nil(t, pool, "level %d should not retain its
workspace", level)
+ }
+ }
+
+ for _, level := range []int{brotli.BestSpeed - 1,
brotli.BestCompression + 1} {
+ assert.Nil(t, brotliWriterPool(level), "invalid level %d should
not be pooled", level)
+ }
+}
+
+func TestBrotliHighQualityWritersAreNotRetained(t *testing.T) {
+ const pattern = "parquet-page-data-pattern-0123456789abcdef"
+ src := bytes.Repeat([]byte(pattern),
(256*1024+len(pattern)-1)/len(pattern))
+ codec := brotliCodec{}
+
+ for level := maxPooledBrotliLevel + 1; level <= brotli.BestCompression;
level++ {
+ codec.EncodeLevel(nil, src, level)
+ assert.Nil(t, brotliWriterPool(level), "level %d should not
retain its workspace", level)
+ }
+}
+
+func TestReleaseBrotliWriterIsBounded(t *testing.T) {
+ pool := make(chan *brotli.Writer, brotliWriterPoolSize)
+ w1 := brotli.NewWriterLevel(nil, brotli.DefaultCompression)
+ w2 := brotli.NewWriterLevel(nil, brotli.DefaultCompression)
+
+ releaseBrotliWriter(pool, w1)
+ releaseBrotliWriter(pool, w2)
+
+ assert.Len(t, pool, brotliWriterPoolSize)
+}
diff --git a/parquet/compress/compress_test.go
b/parquet/compress/compress_test.go
index 9f6e38d3..5fe8f867 100644
--- a/parquet/compress/compress_test.go
+++ b/parquet/compress/compress_test.go
@@ -259,6 +259,66 @@ func TestBrotliCompressBound(t *testing.T) {
}
}
+func TestBrotliEncodeLevelReuse(t *testing.T) {
+ codec, err := compress.GetCodec(compress.Codecs.Brotli)
+ assert.NoError(t, err)
+
+ tests := []struct {
+ name string
+ src []byte
+ }{
+ {name: "empty", src: nil},
+ {name: "one byte", src: []byte{0xab}},
+ {name: "compressible", src: makeCompressibleData(4 * 1024)},
+ {name: "random", src: makeRandomData(4 * 1024)},
+ }
+ levels := []int{0, compress.DefaultCompressionLevel, 5, 11, -2, 12}
+
+ for _, tt := range tests {
+ dst := make([]byte, 0,
int(codec.CompressBound(int64(len(tt.src)))))
+ for _, level := range levels {
+ compressed := codec.EncodeLevel(dst, tt.src, level)
+ decoded, err := compress.Decode(codec, nil, compressed)
+ if !assert.NoError(t, err, "%s, level %d", tt.name,
level) {
+ continue
+ }
+ assert.True(t, bytes.Equal(tt.src, decoded), "%s, level
%d", tt.name, level)
+ dst = compressed[:0]
+ }
+ }
+}
+
+func TestBrotliEncodeLevelConcurrent(t *testing.T) {
+ codec, err := compress.GetCodec(compress.Codecs.Brotli)
+ assert.NoError(t, err)
+
+ src := makeCompressibleData(4 * 1024)
+ levels := []int{0, compress.DefaultCompressionLevel, 9, 11}
+
+ const workers = 16
+ var wg sync.WaitGroup
+ wg.Add(workers)
+ for i := 0; i < workers; i++ {
+ go func(i int) {
+ defer wg.Done()
+ for j := 0; j < len(levels); j++ {
+ level := levels[(i+j)%len(levels)]
+ compressed := codec.EncodeLevel(nil, src, level)
+ decoded, err := compress.Decode(codec, nil,
compressed)
+ if err != nil {
+ t.Errorf("level %d: decode failed: %v",
level, err)
+ return
+ }
+ if !bytes.Equal(src, decoded) {
+ t.Errorf("level %d: decoded data
differs", level)
+ return
+ }
+ }
+ }(i)
+ }
+ wg.Wait()
+}
+
func TestCompressReaderWriter(t *testing.T) {
tests := []struct {
c compress.Compression
