samueleresca commented on code in PR #1331:
URL: https://github.com/apache/arrow-go/pull/1331#discussion_r4212007856
##########
arrow/ipc/compression.go:
##########
@@ -81,29 +84,72 @@ 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
}
+func (z *lz4Decompressor) Decompress(dst, src []byte) error {
+ z.Reset(bytes.NewReader(src))
+ _, err := io.ReadFull(z.Reader, dst)
Review Comment:
Fixed:
https://github.com/apache/arrow-go/pull/1331/commits/6dc67ed1365096a02bc3cb177f9d0ac848b20608
##########
arrow/ipc/reader_compression_benchmark_test.go:
##########
@@ -0,0 +1,156 @@
+// 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)
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
Review Comment:
Fixed:
https://github.com/apache/arrow-go/pull/1331/commits/6dc67ed1365096a02bc3cb177f9d0ac848b20608
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
