Add BufferedReaderV2 interface for advanced buffer management

- Introduce BufferedReaderV2 interface extending BufferedReader with Buffered()
  and Free() methods.
- Enables explicit resource management and buffer introspection for readers using
  custom allocators.
- Update GetStreamV2 to return BufferedReaderV2, allowing callers to free
  allocated buffers when done.
- Keep BufferedReader and BufferedReader for backwards compatibility

Author: daniel-adam-tfs

internal/utils/buf_reader.go

@@ -40,7 +40,7 @@ type byteReader struct {
 
 // NewByteReader creates a new ByteReader instance from the given byte slice.
 // It wraps the bytes.NewReader function to implement BufferedReader interface.
-// It is considered not to own the underlying byte slice, so the Free method is a no-op.
+// It is considered not to own the underlying byte slice.
 func NewByteReader(buf []byte) *byteReader {
 	r := bytes.NewReader(buf)
 	return &byteReader{
@@ -113,7 +113,11 @@ func (r *byteReader) BufferSize() int { return len(r.buf) }
 
 func (r *byteReader) Buffered() int { return len(r.buf) - r.pos }
 
-func (r *byteReader) Free() {}
+func (r *byteReader) Free() {
+	r.r = nil
+	r.buf = nil
+	r.pos = 0
+}
 
 // bytesBufferReader is a byte slice with a bytes reader wrapped around it.
 // It uses an allocator to allocate and free the underlying byte slice.

parquet/file/page_reader.go

@@ -349,7 +349,7 @@ func (d *DictionaryPage) Release() {
 func (d *DictionaryPage) IsSorted() bool { return d.sorted }
 
 type serializedPageReader struct {
-	r             parquet.BufferedReader
+	r             parquet.BufferedReaderV2
 	chunk         *metadata.ColumnChunkMetaData
 	colIdx        int
 	pgIndexReader *metadata.RowGroupPageIndexReader
@@ -421,6 +421,25 @@ func (p *serializedPageReader) init(compressType compress.Compression, ctx *Cryp
 	return nil
 }
 
+type bufferedReaderV2Adapter struct {
+	parquet.BufferedReader
+}
+
+func (b *bufferedReaderV2Adapter) Buffered() int {
+	return 0
+}
+
+func (b *bufferedReaderV2Adapter) Free() {
+	// no-op
+}
+
+func getBufferedReaderV2(r parquet.BufferedReader) parquet.BufferedReaderV2 {
+	if brV2, ok := r.(parquet.BufferedReaderV2); ok {
+		return brV2
+	}
+	return &bufferedReaderV2Adapter{BufferedReader: r}
+}
+
 // NewPageReader returns a page reader for the data which can be read from the provided reader and compression.
 //
 // Deprecated: This function isn't properly safe for public API use and should not be utilized
@@ -436,7 +455,7 @@ func NewPageReader(r parquet.BufferedReader, nrows int64, compressType compress.
 	}
 
 	rdr := &serializedPageReader{
-		r:                 r,
+		r:                 getBufferedReaderV2(r),
 		maxPageHeaderSize: defaultMaxPageHeaderSize,
 		nrows:             nrows,
 		mem:               mem,
@@ -458,7 +477,10 @@ func NewPageReader(r parquet.BufferedReader, nrows int64, compressType compress.
 func (p *serializedPageReader) Reset(r parquet.BufferedReader, nrows int64, compressType compress.Compression, ctx *CryptoContext) {
 	p.rowsSeen, p.pageOrd, p.nrows = 0, 0, nrows
 	p.curPageHdr, p.curPage, p.err = nil, nil, nil
-	p.r = r
+	if p.r != nil && p.r != r {
+		p.r.Free()
+	}
+	p.r = getBufferedReaderV2(r)
 
 	p.codec, p.err = compress.GetCodec(compressType)
 	if p.err != nil {
@@ -508,26 +530,6 @@ func (p *serializedPageReader) Page() Page {
 	return p.curPage
 }
 
-func (p *serializedPageReader) decompress(rd io.Reader, lenCompressed int, buf []byte) ([]byte, error) {
-	p.decompressBuffer.ResizeNoShrink(lenCompressed)
-
-	// Read directly into the memory.Buffer's backing slice
-	n, err := io.ReadFull(rd, p.decompressBuffer.Bytes()[:lenCompressed])
-	if err != nil {
-		return nil, err
-	}
-	if n != lenCompressed {
-		return nil, fmt.Errorf("parquet: expected to read %d compressed bytes, got %d", lenCompressed, n)
-	}
-
-	data := p.decompressBuffer.Bytes()[:lenCompressed]
-	if p.cryptoCtx.DataDecryptor != nil {
-		data = p.cryptoCtx.DataDecryptor.Decrypt(data)
-	}
-
-	return p.codec.Decode(buf, data), nil
-}
-
 type dataheader interface {
 	IsSetStatistics() bool
 	GetStatistics() *format.Statistics
@@ -613,7 +615,7 @@ func (p *serializedPageReader) GetDictionaryPage() (*DictionaryPage, error) {
 	return nil, nil
 }
 
-func (p *serializedPageReader) readPageHeader(rd parquet.BufferedReader, hdr *format.PageHeader) error {
+func (p *serializedPageReader) readPageHeader(rd parquet.BufferedReaderV2, hdr *format.PageHeader) error {
 	allowedPgSz := defaultPageHeaderSize
 	for {
 		view, err := rd.Peek(allowedPgSz)
@@ -699,7 +701,7 @@ func (p *serializedPageReader) SeekToPageWithRow(rowIdx int64) error {
 
 // readOrStealData attempts to steal data from the buffered reader if enough is buffered,
 // otherwise reads from the underlying reader into the provided buffer.
-func (p *serializedPageReader) readOrStealData(r parquet.BufferedReader, lenCompressed int, buffer *memory.Buffer) ([]byte, error) {
+func (p *serializedPageReader) readOrStealData(r parquet.BufferedReaderV2, lenCompressed int, buffer *memory.Buffer) ([]byte, error) {
 	// if enough data is buffered, steal it to avoid an extra copy
 	if r.Buffered() >= lenCompressed {
 		data, err := r.Peek(lenCompressed)
@@ -727,7 +729,7 @@ func (p *serializedPageReader) readOrStealData(r parquet.BufferedReader, lenComp
 }
 
 func (p *serializedPageReader) getPageBytesV1(
-	r parquet.BufferedReader, isCompressed bool, lenCompressed, lenUncompressed int, buffer *memory.Buffer,
+	r parquet.BufferedReaderV2, isCompressed bool, lenCompressed, lenUncompressed int, buffer *memory.Buffer,
 ) ([]byte, error) {
 	// 8 possible cases:
 	// 1. enough data buffered (r.Buffered() >= lenCompressed)
@@ -773,7 +775,7 @@ func (p *serializedPageReader) getPageBytesV1(
 	return data, nil
 }
 
-func (p *serializedPageReader) readV2UnencryptedCompressedWithLevels(r parquet.BufferedReader, lenCompressed, lenUncompressed, levelsBytelen int, buffer *memory.Buffer) ([]byte, error) {
+func (p *serializedPageReader) readV2UnencryptedCompressedWithLevels(r parquet.BufferedReaderV2, lenCompressed, lenUncompressed, levelsBytelen int, buffer *memory.Buffer) ([]byte, error) {
 	// Special case: unencrypted + compressed + has levels
 	// Read levels directly into output buffer, compressed data into decompress buffer
 	buffer.ResizeNoShrink(lenUncompressed)
@@ -800,7 +802,7 @@ func (p *serializedPageReader) readV2UnencryptedCompressedWithLevels(r parquet.B
 }
 
 func (p *serializedPageReader) getPageBytesV2(
-	r parquet.BufferedReader, isCompressed bool, lenCompressed, lenUncompressed, levelsBytelen int, buffer *memory.Buffer,
+	r parquet.BufferedReaderV2, isCompressed bool, lenCompressed, lenUncompressed, levelsBytelen int, buffer *memory.Buffer,
 ) ([]byte, error) {
 	// Special case: unencrypted + compressed + has levels - read levels and compressed data separately
 	if r.Buffered() < lenCompressed && p.cryptoCtx.DataDecryptor == nil && isCompressed && levelsBytelen > 0 {

parquet/file/row_group_reader.go

@@ -114,7 +114,7 @@ func (r *RowGroupReader) GetColumnPageReader(i int) (PageReader, error) {
 		colLen += padding
 	}
 
-	stream, err := r.props.GetStream(r.r, colStart, colLen)
+	stream, err := r.props.GetStreamV2(r.r, colStart, colLen)
 	if err != nil {
 		return nil, err
 	}

parquet/internal/testutils/pagebuilder.go

@@ -136,7 +136,8 @@ type DictionaryPageBuilder struct {
 func NewDictionaryPageBuilder(d *schema.Column) *DictionaryPageBuilder {
 	return &DictionaryPageBuilder{
 		encoding.NewEncoder(d.PhysicalType(), parquet.Encodings.Plain, true, d, mem).(encoding.DictEncoder),
-		0, false}
+		0, false,
+	}
 }
 
 func (d *DictionaryPageBuilder) AppendValues(values interface{}) encoding.Buffer {
@@ -243,9 +244,12 @@ func (m *MockPageReader) Err() error {
 }
 
 func (m *MockPageReader) Reset(parquet.BufferedReader, int64, compress.Compression, *file.CryptoContext) {
+	// no-op
 }
 
-func (m *MockPageReader) SetMaxPageHeaderSize(int) {}
+func (m *MockPageReader) SetMaxPageHeaderSize(int) {
+	// no-op
+}
 
 func (m *MockPageReader) Page() file.Page {
 	return m.TestData().Get("pages").Data().([]file.Page)[m.curpage-1]
@@ -278,8 +282,8 @@ func (m *MockPageReader) Next() bool {
 }
 
 func PaginatePlain(version parquet.DataPageVersion, d *schema.Column, values reflect.Value, defLevels, repLevels []int16,
-	maxDef, maxRep int16, lvlsPerPage int, valuesPerPage []int, enc parquet.Encoding) []file.Page {
-
+	maxDef, maxRep int16, lvlsPerPage int, valuesPerPage []int, enc parquet.Encoding,
+) []file.Page {
 	var (
 		npages      = len(valuesPerPage)
 		defLvlStart = 0

parquet/reader_properties.go

@@ -46,15 +46,20 @@ type ReaderProperties struct {
 	BufferedStreamEnabled bool
 }
 
+type BufferedReaderV2 interface {
+	BufferedReader
+	// Buffered returns the number of bytes already read and stored in the buffer
+	Buffered() int
+	// Free releases any resources held by the BufferedReader
+	Free()
+}
+
 type BufferedReader interface {
 	Peek(int) ([]byte, error)
 	Discard(int) (int, error)
 	Outer() utils.Reader
-	// Buffered returns the number of bytes already read and stored in the buffer
-	Buffered() int
 	BufferSize() int
 	Reset(utils.Reader)
-	Free()
 	io.Reader
 }
 
@@ -76,6 +81,29 @@ func (r *ReaderProperties) Allocator() memory.Allocator { return r.alloc }
 // If BufferedStreamEnabled is true, it creates an io.SectionReader, otherwise it will read the entire section
 // into a buffer in memory and return a bytes.NewReader for that buffer.
 func (r *ReaderProperties) GetStream(source io.ReaderAt, start, nbytes int64) (BufferedReader, error) {
+	if r.BufferedStreamEnabled {
+		return utils.NewBufferedReader(io.NewSectionReader(source, start, nbytes), int(r.BufferSize), nil), nil
+	}
+
+	data := make([]byte, nbytes)
+	n, err := source.ReadAt(data, start)
+	if err != nil {
+		return nil, fmt.Errorf("parquet: tried reading from file, but got error: %w", err)
+	}
+	if n != int(nbytes) {
+		return nil, fmt.Errorf("parquet: tried reading %d bytes starting at position %d from file but only got %d", nbytes, start, n)
+	}
+
+	return utils.NewByteReader(data), nil
+}
+
+// GetStreamV2 returns a section of the underlying reader based on whether or not BufferedStream is enabled.
+//
+// If BufferedStreamEnabled is true, it creates an io.SectionReader, otherwise it will read the entire section
+// into a buffer in memory and return a bytes.NewReader for that buffer.
+// In comparison with GetStream, this version uses r.alloc to allocate the buffer for reading data and returns BufferedReaderV2,
+// to allow freeing the allocated buffer when no longer needed with the Free() method.
+func (r *ReaderProperties) GetStreamV2(source io.ReaderAt, start, nbytes int64) (BufferedReaderV2, error) {
 	if r.BufferedStreamEnabled {
 		return utils.NewBufferedReader(io.NewSectionReader(source, start, nbytes), int(r.BufferSize), r.alloc), nil
 	}

parquet/reader_writer_properties_test.go

@@ -89,16 +89,16 @@ func TestReaderPropsGetStreamWithAllocator(t *testing.T) {
 
 	// no leak on success
 	props := parquet.NewReaderProperties(pool)
-	bufRdr, err := props.GetStream(rdr, 0, int64(len(data)))
+	bufRdr, err := props.GetStreamV2(rdr, 0, int64(len(data)))
 	assert.NoError(t, err)
 	bufRdr.Free()
 
 	// no leak on reader error
-	_, err = props.GetStream(&mockReaderAt{}, 0, 10)
+	_, err = props.GetStreamV2(&mockReaderAt{}, 0, 10)
 	assert.Error(t, err)
 
 	// no leak on insufficient read
-	_, err = props.GetStream(rdr, 0, int64(len(data)+10))
+	_, err = props.GetStreamV2(rdr, 0, int64(len(data)+10))
 	assert.Error(t, err)
 }
 
@@ -113,19 +113,19 @@ func TestReaderPropsGetStreamBufferedWithAllocator(t *testing.T) {
 	props.BufferedStreamEnabled = true
 
 	buf := make([]byte, len(data))
-	bufRdr, err := props.GetStream(rdr, 0, int64(len(data)))
+	bufRdr, err := props.GetStreamV2(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)
+	bufRdr, err = props.GetStreamV2(&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))
+	bufRdr, err = props.GetStreamV2(rdr, 0, int64(len(data)+10))
 	assert.NoError(t, err)
 	n, err := bufRdr.Read(buf)
 	assert.NoError(t, err)