samueleresca commented on code in PR #1331:
URL: https://github.com/apache/arrow-go/pull/1331#discussion_r4066203695
##########
arrow/ipc/compression.go:
##########
@@ -81,29 +84,63 @@ func getCompressor(codec flatbuf.CompressionType)
compressor {
}
type decompressor interface {
- io.Reader
- Reset(io.Reader)
+ Decompress(dst, src []byte) error
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
Review Comment:
See updated PR description
--
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]
