fixup! Ensure that decompressBuffer doesn't get reallocated by io.CopyN

Author: daniel-adam-tfs

internal/utils/buf_reader.go

@@ -122,6 +122,9 @@ type bytesBufferReader struct {
 
 // NewBytesBufferReader creates a new bytesBufferReader with the given size and allocator.
 func NewBytesBufferReader(size int, alloc memory.Allocator) *bytesBufferReader {
+	if alloc == nil {
+		alloc = memory.DefaultAllocator
+	}
 	buf := alloc.Allocate(size)
 	return &bytesBufferReader{
 		alloc: alloc,
@@ -159,6 +162,9 @@ type bufferedReader struct {
 // except Peek will expand the internal buffer if needed rather than return
 // an error.
 func NewBufferedReader(rd Reader, sz int, alloc memory.Allocator) *bufferedReader {
+	if alloc == nil {
+		alloc = memory.DefaultAllocator
+	}
 	r := &bufferedReader{
 		alloc: alloc,
 		rd:    rd,

parquet/file/page_reader.go

@@ -17,7 +17,6 @@
 package file
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
 	"io"
@@ -504,21 +503,16 @@ func (p *serializedPageReader) Page() Page {
 }
 
 func (p *serializedPageReader) decompress(rd io.Reader, lenCompressed int, buf []byte) ([]byte, error) {
-	// As of go1.25.3: There is an issue when bytes.Buffer and io.CopyN are used together. io.CopyN
-	// uses io.LimitReader, which does an additional read on the underlying reader to determine EOF.
-	// However, bytes.Buffer always attempts to read at least bytes.MinRead (which is 512 bytes) from the
-	// underlying reader, even if there is less data available than that. So even if there are no more bytes,
-	// the buffer must have at least bytes.MinRead capacity remaining to avoid a relocation.
-	if p.decompressBuffer.Cap() < lenCompressed+bytes.MinRead {
-		p.decompressBuffer.Reserve(lenCompressed + bytes.MinRead)
-	}
 	p.decompressBuffer.ResizeNoShrink(lenCompressed)
-	b := bytes.NewBuffer(p.decompressBuffer.Bytes()[:0])
-	if _, err := io.CopyN(b, rd, int64(lenCompressed)); err != nil {
+	data := p.decompressBuffer.Bytes()
+	n, err := io.ReadFull(rd, data)
+	if err != nil {
 		return nil, err
 	}
+	if n != lenCompressed {
+		return nil, fmt.Errorf("parquet: expected to read %d bytes but only read %d", lenCompressed, n)
+	}
 
-	data := p.decompressBuffer.Bytes()
 	if p.cryptoCtx.DataDecryptor != nil {
 		data = p.cryptoCtx.DataDecryptor.Decrypt(p.decompressBuffer.Bytes())
 	}
@@ -563,6 +557,7 @@ func (p *serializedPageReader) GetDictionaryPage() (*DictionaryPage, error) {
 			io.NewSectionReader(p.r.Outer(), p.dictOffset-p.baseOffset, p.dataOffset-p.baseOffset),
 			readBufSize,
 			p.mem)
+		defer rd.Free()
 		if err := p.readPageHeader(rd, hdr); err != nil {
 			return nil, err
 		}
@@ -774,6 +769,7 @@ func (p *serializedPageReader) Next() bool {
 
 			firstRowIdx := p.rowsSeen
 			p.rowsSeen += int64(dataHeader.GetNumValues())
+
 			data, err := p.decompress(p.r, lenCompressed, buf.Bytes())
 			if err != nil {
 				p.err = err

parquet/file/row_group_reader.go

@@ -134,11 +134,13 @@ func (r *RowGroupReader) GetColumnPageReader(i int) (PageReader, error) {
 	}
 
 	if r.fileDecryptor == nil {
+		stream.Free()
 		return nil, xerrors.New("column in rowgroup is encrypted, but no file decryptor")
 	}
 
 	const encryptedRowGroupsLimit = 32767
 	if i > encryptedRowGroupsLimit {
+		stream.Free()
 		return nil, xerrors.New("encrypted files cannot contain more than 32767 column chunks")
 	}
 

parquet/reader_writer_properties_test.go

@@ -18,6 +18,7 @@ package parquet_test
 
 import (
 	"bytes"
+	"errors"
 	"testing"
 
 	"github.com/apache/arrow-go/v18/arrow/memory"
@@ -67,7 +68,67 @@ func TestReaderPropsGetStreamInsufficient(t *testing.T) {
 	buf := memory.NewBufferBytes([]byte(data))
 	rdr := bytes.NewReader(buf.Bytes())
 
-	props := parquet.NewReaderProperties(nil)
-	_, err := props.GetStream(rdr, 12, 15)
+	props1 := parquet.NewReaderProperties(nil)
+	_, err := props1.GetStream(rdr, 12, 15)
+	assert.Error(t, err)
+}
+
+type mockReaderAt struct{}
+
+func (m *mockReaderAt) ReadAt(p []byte, off int64) (int, error) {
+	return 0, errors.New("mock error")
+}
+
+func TestReaderPropsGetStreamWithAllocator(t *testing.T) {
+	pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	defer pool.AssertSize(t, 0)
+
+	data := "data to read"
+	buf := memory.NewBufferBytes([]byte(data))
+	rdr := bytes.NewReader(buf.Bytes())
+
+	// no leak on success
+	props := parquet.NewReaderProperties(pool)
+	bufRdr, err := props.GetStream(rdr, 0, int64(len(data)))
+	assert.NoError(t, err)
+	bufRdr.Free()
+
+	// no leak on reader error
+	_, err = props.GetStream(&mockReaderAt{}, 0, 10)
+	assert.Error(t, err)
+
+	// no leak on insufficient read
+	_, err = props.GetStream(rdr, 0, int64(len(data)+10))
 	assert.Error(t, err)
 }
+
+func TestReaderPropsGetStreamBufferedWithAllocator(t *testing.T) {
+	pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	defer pool.AssertSize(t, 0)
+
+	data := "data to read"
+	rdr := bytes.NewReader(memory.NewBufferBytes([]byte(data)).Bytes())
+
+	props := parquet.NewReaderProperties(pool)
+	props.BufferedStreamEnabled = true
+
+	buf := make([]byte, len(data))
+	bufRdr, err := props.GetStream(rdr, 0, int64(len(data)))
+	assert.NoError(t, err)
+	_, err = bufRdr.Read(buf)
+	assert.NoError(t, err)
+	bufRdr.Free()
+
+	bufRdr, err = props.GetStream(&mockReaderAt{}, 0, 10)
+	assert.NoError(t, err)
+	_, err = bufRdr.Read(buf)
+	assert.Error(t, err)
+	bufRdr.Free()
+
+	bufRdr, err = props.GetStream(rdr, 0, int64(len(data)+10))
+	assert.NoError(t, err)
+	n, err := bufRdr.Read(buf)
+	assert.NoError(t, err)
+	assert.NotEqual(t, len(data)+10, n)
+	bufRdr.Free()
+}