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 53a747fe feat(parquet): opt-in streaming reads for large data pages
(#880)
53a747fe is described below
commit 53a747fe62f61d6be986dfaa9cf8d1bc90f750dc
Author: Ruihao Chen <[email protected]>
AuthorDate: Tue Jul 14 03:09:38 2026 +0800
feat(parquet): opt-in streaming reads for large data pages (#880)
### Rationale for this change
Addresses #865. Today the reader materializes each full uncompressed data
page before decoding, so peak memory scales with the page's uncompressed size —
a problem for files with very large data pages (e.g. a 1 GB compressed / 1.4 GB
uncompressed page). This adds an **opt-in** path that keeps peak memory near
the requested batch size for the common large-page shape.
Approach was discussed and agreed on the issue.
### What changes are included in this PR
Opt-in via `ReaderProperties.PageStreamingEnabled` (default **false**, no
behavior change). When set, **eligible** data pages skip full-page
decompression and decode their values incrementally; every ineligible page
falls back to the existing materialized path.
**Scope (first cut):**
- Data Page **V1 + V2**, **PLAIN** encoding
- Physical types **`byte_array`** and **`fixed_len_byte_array`**
- Codecs **`UNCOMPRESSED` / `GZIP` / `BROTLI` / `ZSTD`** (explicit
allowlist — Snappy is a raw block whereas its `NewReader` is the framed format,
and LZ4_RAW has no streaming reader)
**Design:**
- `parquet/internal/encoding/streaming`: a small `ValueBuffer` — an
incremental byte source that reads from an `io.Reader`, slides as values are
consumed, grows to fit a single oversized value, and owns the page's stream
(drained + closed on release).
- The materialized `PlainByteArrayDecoder` /
`PlainFixedLenByteArrayDecoder` keep their `[]byte` logic byte-for-byte,
gaining only a `src` field + a one-line guard in `Decode`/`Discard`. Streaming
decode lives beside them; `SetData` resets `src`, so a cached decoder reused
for a materialized page reverts. Other decoders are untouched.
- Page reader: for eligible pages, V2 levels are read raw (uncompressed,
header byte-lengths) and V1 levels are peeled off the decompressed stream using
the same length rules as `LevelDecoder.SetData`; values become a `ValueBuffer`
over the compressed region. On page release the underlying stream is drained +
closed so the reader lands on the next page header even if values were skipped.
- Column reader: the same cached PLAIN decoder is fed via `SetSource`
(streaming) or `SetData` (materialized). `Page`/`DataPage` already return
internal types (e.g. `format.Encoding`), so they can't be implemented outside
`parquet/`; the added **unexported** `valueBuffer`/`valueClip` accessors are
internal plumbing invisible to external consumers.
### Are these changes tested?
Yes.
- All existing encoding tests pass unchanged (materialized parity); the
streaming decoders are additionally exercised with a one-byte-at-a-time reader
and oversized values.
- End-to-end round-trip (`parquet/file/streaming_page_test.go`): streaming
vs materialized reads are asserted identical across V1/V2 and each codec, for
**required** and **nullable** columns (the nullable case covers the V1
def-level peel + streaming `DecodeSpaced`), including values larger than a page
and the stream buffer.
### Are there any user-facing changes?
One new opt-in field, `ReaderProperties.PageStreamingEnabled` (default
false). No behavior change to decoded values; the only interface change is two
unexported accessor methods on `DataPage` (see above).
Follow-ups (out of scope here): fixed-width numeric PLAIN types, raw-Snappy
and LZ4_RAW streaming, and a memory benchmark.
---
parquet/compress/zstd.go | 19 +-
parquet/file/column_reader.go | 44 +-
parquet/file/export_test.go | 21 +
parquet/file/page_reader.go | 307 +++++--
parquet/file/page_reader_internal_test.go | 54 ++
parquet/file/row_group_reader.go | 28 +
parquet/file/streaming_page_test.go | 942 +++++++++++++++++++++
parquet/internal/encoding/byte_array_decoder.go | 72 ++
.../encoding/fixed_len_byte_array_decoder.go | 48 ++
.../internal/encoding/streaming/streaming_test.go | 70 ++
.../internal/encoding/streaming/value_buffer.go | 174 ++++
.../streaming/value_buffer_internal_test.go | 136 +++
.../internal/encoding/streaming_decoder_test.go | 120 +++
parquet/reader_properties.go | 20 +-
14 files changed, 1983 insertions(+), 72 deletions(-)
diff --git a/parquet/compress/zstd.go b/parquet/compress/zstd.go
index a046a48a..c6795c52 100644
--- a/parquet/compress/zstd.go
+++ b/parquet/compress/zstd.go
@@ -112,14 +112,27 @@ func (zstdCodec) Decode(dst, src []byte) []byte {
return dst
}
+var globalDecoderPool = sync.Pool{
+ New: func() interface{} {
+ dec, _ := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1))
+ return dec
+ },
+}
+
func (z *zstdcloser) Close() error {
- z.Decoder.Close()
+ if z.Decoder == nil {
+ return nil
+ }
+ _ = z.Reset(nil)
+ globalDecoderPool.Put(z.Decoder)
+ z.Decoder = nil
return nil
}
func (zstdCodec) NewReader(r io.Reader) io.ReadCloser {
- ret, _ := zstd.NewReader(r)
- return &zstdcloser{ret}
+ dec := globalDecoderPool.Get().(*zstd.Decoder)
+ _ = dec.Reset(r)
+ return &zstdcloser{dec}
}
func (zstdCodec) NewWriter(w io.Writer) io.WriteCloser {
diff --git a/parquet/file/column_reader.go b/parquet/file/column_reader.go
index 48c392b9..b9c96f73 100644
--- a/parquet/file/column_reader.go
+++ b/parquet/file/column_reader.go
@@ -25,6 +25,7 @@ import (
"github.com/apache/arrow-go/v18/internal/utils"
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/internal/encoding"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding/streaming"
"github.com/apache/arrow-go/v18/parquet/internal/encryption"
format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
"github.com/apache/arrow-go/v18/parquet/schema"
@@ -261,9 +262,7 @@ func (c *columnChunkReader) setPageReader(rdr PageReader) {
// Close closes the page raeder and the page if set.
func (c *columnChunkReader) Close() error {
- if c.curPage != nil {
- c.curPage.Release()
- }
+ // c.rdr owns curPage (c.curPage == c.rdr.Page()) and releases it on
Close.
if c.rdr != nil {
return c.rdr.Close()
}
@@ -366,7 +365,7 @@ func (c *columnChunkReader) processPage() (bool, error) {
return true, err
}
- return true, c.initDataDecoder(c.curPage, lvlByteLen)
+ return true, c.initDataDecoder(c.curPage.(DataPage), lvlByteLen)
}
// read a new page from the page reader
@@ -398,7 +397,7 @@ func (c *columnChunkReader) readNewPage() bool {
func (c *columnChunkReader) initLevelDecodersV2(page *DataPageV2) (int64,
error) {
c.numBuffered = int64(page.nvals)
c.numDecoded = 0
- buf := page.Data()
+ buf := page.levelBuffer()
totalLvlLen := int64(page.repLvlByteLen) + int64(page.defLvlByteLen)
if totalLvlLen > int64(len(buf)) {
@@ -428,7 +427,7 @@ func (c *columnChunkReader) initLevelDecodersV1(page
*DataPageV1, repLvlEncoding
c.numBuffered = int64(page.nvals)
c.numDecoded = 0
- buf := page.Data()
+ buf := page.levelBuffer()
maxSize := len(buf)
levelsByteLen := int64(0)
@@ -456,13 +455,7 @@ func (c *columnChunkReader) initLevelDecodersV1(page
*DataPageV1, repLvlEncoding
return levelsByteLen, nil
}
-func (c *columnChunkReader) initDataDecoder(page Page, lvlByteLen int64) error
{
- buf := page.Data()
- if int64(len(buf)) < lvlByteLen {
- return errors.New("parquet: page smaller than size of encoded
levels")
- }
-
- buf = buf[lvlByteLen:]
+func (c *columnChunkReader) initDataDecoder(page DataPage, lvlByteLen int64)
error {
encoding := page.Encoding()
if isDictIndexEncoding(encoding) {
@@ -498,8 +491,24 @@ func (c *columnChunkReader) initDataDecoder(page Page,
lvlByteLen int64) error {
}
c.curEncoding = encoding
- c.curDecoder.SetData(int(c.numBuffered), buf)
- return nil
+ return c.setDecoderData(page, lvlByteLen)
+}
+
+func (c *columnChunkReader) setDecoderData(page DataPage, lvlByteLen int64)
error {
+ if src := page.valueBuffer(); src != nil {
+ sd, ok := c.curDecoder.(streaming.Decoder)
+ if !ok {
+ return fmt.Errorf("parquet: streaming page but %s
decoder does not support streaming", c.descr.PhysicalType())
+ }
+ sd.SetSource(int(c.numBuffered), src)
+ return nil
+ }
+
+ buf := page.Data()
+ if int64(len(buf)) < lvlByteLen {
+ return errors.New("parquet: page smaller than size of encoded
levels")
+ }
+ return c.curDecoder.SetData(int(c.numBuffered), buf[lvlByteLen:])
}
// readDefinitionLevels decodes the definition levels from the page and returns
@@ -720,6 +729,11 @@ func (c *columnChunkReader) readBatchInPage(batchSize
int64, totalRead int64, de
return 0, 0, c.err
}
+ // Streaming holds a whole batch's aliased values at once; clip it near
the buffer size.
+ if clip := c.curPage.(DataPage).valueClip(); clip > 0 && batchSize >
clip {
+ batchSize = clip
+ }
+
ndefs, toRead, err := c.determineNumToRead(batchSize, defLvls, repLvls)
if err != nil {
return 0, 0, err
diff --git a/parquet/file/export_test.go b/parquet/file/export_test.go
new file mode 100644
index 00000000..ada4b966
--- /dev/null
+++ b/parquet/file/export_test.go
@@ -0,0 +1,21 @@
+// 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 file
+
+// StreamingThreshold exposes the package-level streaming page-size threshold
so tests
+// can lower it and exercise streaming on small pages.
+var StreamingThreshold = &streamingThreshold
diff --git a/parquet/file/page_reader.go b/parquet/file/page_reader.go
index 66b8473a..50aef89c 100644
--- a/parquet/file/page_reader.go
+++ b/parquet/file/page_reader.go
@@ -17,16 +17,20 @@
package file
import (
+ "encoding/binary"
"errors"
"fmt"
"io"
+ "math/bits"
"sort"
"sync"
+ "github.com/apache/arrow-go/v18/arrow/bitutil"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/internal/utils"
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/compress"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding/streaming"
"github.com/apache/arrow-go/v18/parquet/internal/encryption"
format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
"github.com/apache/arrow-go/v18/parquet/internal/thrift"
@@ -68,7 +72,11 @@ const (
type Page interface {
// Returns which kind of page this is
Type() PageType
- // Get the raw bytes of this page
+ // Get the raw bytes of this page.
+ //
+ // Note: with PageStreamingEnabled, a streaming data page's Data holds
only the
+ // rep/def level region, not the whole page; decode values through a
column or
+ // record reader instead.
Data() []byte
// return the encoding used for this page, Plain/RLE, etc.
Encoding() format.Encoding
@@ -91,6 +99,43 @@ func (p *page) Data() []byte { return
p.buf.Bytes() }
func (p *page) NumValues() int32 { return p.nvals }
func (p *page) Encoding() format.Encoding { return p.encoding }
+// dataPageBase holds the members shared by DataPageV1/V2 but not dictionary
+// pages: the rep/def level buffer accessor and, in streaming mode, the
+// incremental value source.
+type dataPageBase struct {
+ page
+
+ // valueBuf is only set for streaming mode to read values incrementally.
+ // In such mode, buf only holds the rep+def level region.
+ valueBuf streaming.ValueBuffer
+ // clip bounds a read's batch to keep one Decode's aliased values near
the buffer
+ // size; 0 for non-streaming pages.
+ clip int64
+}
+
+func (d *dataPageBase) valueBuffer() streaming.ValueBuffer { return d.valueBuf
}
+func (d *dataPageBase) levelBuffer() []byte { return
d.buf.Bytes() }
+
+func (d *dataPageBase) valueClip() int64 { return d.clip }
+
+func (d *dataPageBase) setStreamingValues(mem memory.Allocator,
lenUncompressed int, levelBuf *memory.Buffer, valReader, rawSrc io.Reader,
closer io.Closer) {
+ d.buf = levelBuf
+ valueBytes := lenUncompressed - levelBuf.Len()
+ d.valueBuf = streaming.NewStreamBuffer(mem, valReader, valueBytes,
drainAndClose(rawSrc, closer))
+ d.clip = streaming.ClipBatch(valueBytes, int(d.nvals))
+}
+
+// releaseValueStream closes a streaming page's value source (draining +
closing its
+// underlying compressed stream so the file reader lands on the next page
header even
+// when values were skipped). It is a no-op for non-streaming pages.
+func (d *dataPageBase) releaseValueStream() {
+ if d.valueBuf != nil {
+ _ = d.valueBuf.Close()
+ d.valueBuf = nil
+ }
+ d.clip = 0
+}
+
// DataPageConfig is a struct for passing configuration params to data page
creation
// which can be expanded in the future without causing any breaking changes.
type DataPageConfig struct {
@@ -111,6 +156,11 @@ type DataPage interface {
// FirstRowIndex returns the row ordinal within the row group
// to the first row in the data page, or -1 if not set.
FirstRowIndex() int64
+
+ // valueBuffer returns the streaming value source, or nil for a
non-streaming page.
+ valueBuffer() streaming.ValueBuffer
+ // valueClip bounds a read's batch on a streaming page, or 0 for a
non-streaming page.
+ valueClip() int64
}
// Create some pools to use for reusing the data page objects themselves so
that
@@ -133,7 +183,7 @@ var dictPagePool = sync.Pool{
// DataPageV1 represents a DataPage version 1 from the parquet.thrift file
type DataPageV1 struct {
- page
+ dataPageBase
defLvlEncoding format.Encoding
repLvlEncoding format.Encoding
@@ -152,7 +202,7 @@ func NewDataPageV1(buffer *memory.Buffer, num int32,
encoding, defEncoding, repE
dp := dataPageV1Pool.Get().(*DataPageV1)
if dp == nil {
return &DataPageV1{
- page: page{buf: buffer, typ:
format.PageType_DATA_PAGE, nvals: num, encoding: format.Encoding(encoding)},
+ dataPageBase: dataPageBase{page: page{buf: buffer,
typ: format.PageType_DATA_PAGE, nvals: num, encoding:
format.Encoding(encoding)}},
defLvlEncoding: format.Encoding(defEncoding),
repLvlEncoding: format.Encoding(repEncoding),
uncompressedSize: uncompressedSize,
@@ -191,6 +241,10 @@ func NewDataPageV1WithConfig(buffer *memory.Buffer,
defEncoding, repEncoding par
// After calling this function, the object should not be utilized anymore,
otherwise
// conflicts can arise.
func (d *DataPageV1) Release() {
+ if d.buf == nil {
+ return
+ }
+ d.releaseValueStream()
d.buf.Release()
d.buf = nil
dataPageV1Pool.Put(d)
@@ -216,7 +270,7 @@ func (d *DataPageV1) RepetitionLevelEncoding()
parquet.Encoding {
// DataPageV2 is the representation of the V2 data page from the
parquet.thrift spec
type DataPageV2 struct {
- page
+ dataPageBase
nulls int32
nrows int32
@@ -234,7 +288,7 @@ func NewDataPageV2(buffer *memory.Buffer, numValues,
numNulls, numRows int32, en
dp := dataPageV2Pool.Get().(*DataPageV2)
if dp == nil {
return &DataPageV2{
- page: page{buf: buffer, typ:
format.PageType_DATA_PAGE_V2, nvals: numValues, encoding:
format.Encoding(encoding)},
+ dataPageBase: dataPageBase{page: page{buf: buffer,
typ: format.PageType_DATA_PAGE_V2, nvals: numValues, encoding:
format.Encoding(encoding)}},
nulls: numNulls,
nrows: numRows,
defLvlByteLen: defLvlsByteLen,
@@ -277,6 +331,10 @@ func NewDataPageV2WithConfig(buffer *memory.Buffer,
numNulls, numRows int32, def
// After calling this function, the object should not be utilized anymore,
otherwise
// conflicts can arise.
func (d *DataPageV2) Release() {
+ if d.buf == nil {
+ return
+ }
+ d.releaseValueStream()
d.buf.Release()
d.buf = nil
dataPageV2Pool.Put(d)
@@ -358,6 +416,11 @@ type serializedPageReader struct {
mem memory.Allocator
codec compress.Codec
+ // Streaming (PageStreamingEnabled) inputs from schema/properties.
+ columnCanStream bool
+ maxRepLevel int16
+ maxDefLevel int16
+
curPageHdr *format.PageHeader
pageOrd int16
maxPageHeaderSize int
@@ -376,6 +439,11 @@ type serializedPageReader struct {
}
func (p *serializedPageReader) Close() error {
+ if p.curPage != nil {
+ // free a direct PageReader user's streaming value stream
+ p.curPage.Release()
+ p.curPage = nil
+ }
if p.decompressBuffer != nil {
p.decompressBuffer.Release()
p.dictPageBuffer.Release()
@@ -398,6 +466,11 @@ func (p *serializedPageReader) init(compressType
compress.Compression, ctx *Cryp
return err
}
p.codec = codec
+ if _, ok := codec.(compress.StreamingCodec); !ok {
+ // A codec registered via compress.RegisterCodec need not
implement
+ // streaming; disable streaming for this column in that case.
+ p.columnCanStream = false
+ }
if ctx != nil {
p.cryptoCtx = *ctx
@@ -448,9 +521,13 @@ func NewPageReader(r parquet.BufferedReader, nrows int64,
compressType compress.
}
func (p *serializedPageReader) Reset(r parquet.BufferedReader, nrows int64,
compressType compress.Compression, ctx *CryptoContext) {
+ if p.curPage != nil {
+ p.curPage.Release()
+ }
p.rowsSeen, p.pageOrd, p.nrows = 0, 0, nrows
p.curPageHdr, p.curPage, p.err = nil, nil, nil
p.r = r
+ p.columnCanStream = false
p.codec, p.err = compress.GetCodec(compressType)
if p.err != nil {
@@ -723,6 +800,13 @@ func (p *serializedPageReader) SeekToPageWithRow(rowIdx
int64) error {
return fmt.Errorf("parquet: cannot seek column reader to row
index %d", rowIdx)
}
+ // Release before repositioning p.r: a streaming page's
drain-on-release would
+ // otherwise read from the new offset.
+ if p.curPage != nil {
+ p.curPage.Release()
+ p.curPage = nil
+ }
+
var (
oidx metadata.OffsetIndex
err error
@@ -769,6 +853,92 @@ func (p *serializedPageReader) SeekToPageWithRow(rowIdx
int64) error {
return p.err
}
+func (p *serializedPageReader) valueStream(limit io.Reader, compressed bool)
(io.Reader, io.Closer) {
+ if compressed {
+ dec := p.codec.(compress.StreamingCodec).NewReader(limit)
+ return dec, dec
+ }
+ return limit, nil
+}
+
+// drainAndClose returns the cleanup a streaming ValueBuffer runs on Close:
consume
+// any unread compressed bytes for the current page and close the decompressor.
+func drainAndClose(rawSrc io.Reader, closer io.Closer) func() error {
+ return func() error {
+ _, _ = io.Copy(io.Discard, rawSrc)
+ if closer != nil {
+ return closer.Close()
+ }
+ return nil
+ }
+}
+
+// streamingThreshold is the smallest data-page size worth streaming (var so
tests can lower it).
+var streamingThreshold = int64(1 << 20)
+
+func (p *serializedPageReader) pageCanStream(enc format.Encoding,
lenUncompressed int) bool {
+ return p.columnCanStream && enc == format.Encoding_PLAIN &&
p.cryptoCtx.DataDecryptor == nil &&
+ int64(lenUncompressed) > streamingThreshold
+}
+
+// readLevelData reads the V1 rep+def level region, capped at maxBytes (the
+// uncompressed page size) so a bogus length prefix can't force a huge
allocation.
+func (p *serializedPageReader) readLevelData(r io.Reader, repEnc, defEnc
format.Encoding, numValues, maxBytes int) (*memory.Buffer, error) {
+ buf := memory.NewResizableBuffer(p.mem)
+ pos := 0
+ readOne := func(enc format.Encoding, maxLvl int16) error {
+ switch enc {
+ case format.Encoding_RLE:
+ var hdr [4]byte
+ if _, err := io.ReadFull(r, hdr[:]); err != nil {
+ return err
+ }
+ n := int(binary.LittleEndian.Uint32(hdr[:]))
+ if n < 0 || pos+4+n > maxBytes {
+ return fmt.Errorf("parquet: invalid V1 level
region length %d", n)
+ }
+ buf.Resize(pos + 4 + n)
+ copy(buf.Bytes()[pos:], hdr[:])
+ if _, err := io.ReadFull(r,
buf.Bytes()[pos+4:pos+4+n]); err != nil {
+ return err
+ }
+ pos += 4 + n
+ case format.Encoding_BIT_PACKED:
+ bitWidth := bits.Len64(uint64(maxLvl))
+ nbits, ok := utils.Mul(numValues, bitWidth)
+ if !ok {
+ return fmt.Errorf("parquet: V1 level count %d
too large (corrupt page?)", numValues)
+ }
+ n := int(bitutil.BytesForBits(int64(nbits)))
+ if n < 0 || pos+n > maxBytes {
+ return fmt.Errorf("parquet: invalid V1 level
region length %d", n)
+ }
+ buf.Resize(pos + n)
+ if _, err := io.ReadFull(r, buf.Bytes()[pos:pos+n]);
err != nil {
+ return err
+ }
+ pos += n
+ default:
+ return fmt.Errorf("parquet: unsupported V1 level
encoding %s for streaming", enc)
+ }
+ return nil
+ }
+
+ if p.maxRepLevel > 0 {
+ if err := readOne(repEnc, p.maxRepLevel); err != nil {
+ buf.Release()
+ return nil, err
+ }
+ }
+ if p.maxDefLevel > 0 {
+ if err := readOne(defEnc, p.maxDefLevel); err != nil {
+ buf.Release()
+ return nil, err
+ }
+ }
+ return buf, nil
+}
+
func (p *serializedPageReader) Next() bool {
// Loop here because there may be unhandled page types that we skip
until
// finding a page that we do know what to do with
@@ -840,35 +1010,50 @@ func (p *serializedPageReader) Next() bool {
return false
}
- p.dataPageBuffer.ResizeNoShrink(lenUncompressed)
- buf := memory.NewBufferBytes(p.dataPageBuffer.Bytes())
-
firstRowIdx := p.rowsSeen
p.rowsSeen += int64(dataHeader.GetNumValues())
- data, err := p.decompress(p.r, lenCompressed,
buf.Bytes())
- if err != nil {
- p.err = err
- return false
- }
- if len(data) != lenUncompressed {
- p.err = fmt.Errorf("parquet: metadata said %d
bytes uncompressed data page, got %d bytes", lenUncompressed, len(data))
- return false
- }
- // make datapagev1
- p.curPage = &DataPageV1{
- page: page{
- buf: buf,
+ dp := &DataPageV1{
+ dataPageBase: dataPageBase{page: page{
typ: p.curPageHdr.Type,
nvals: dataHeader.GetNumValues(),
encoding: dataHeader.GetEncoding(),
- },
+ }},
defLvlEncoding:
dataHeader.GetDefinitionLevelEncoding(),
repLvlEncoding:
dataHeader.GetRepetitionLevelEncoding(),
uncompressedSize: int32(lenUncompressed),
statistics: extractStats(dataHeader),
firstRowIndex: firstRowIdx,
}
+
+ if p.pageCanStream(dataHeader.GetEncoding(),
lenUncompressed) {
+ limit := io.LimitReader(p.r,
int64(lenCompressed))
+ // V1 compresses the whole body, so the levels
are read through the
+ // codec too (a no-op passthrough for an
uncompressed column).
+ valReader, closer := p.valueStream(limit, true)
+ levelBuf, err := p.readLevelData(valReader,
dataHeader.GetRepetitionLevelEncoding(),
+
dataHeader.GetDefinitionLevelEncoding(), int(dataHeader.GetNumValues()),
lenUncompressed)
+ if err != nil {
+ _ = drainAndClose(limit, closer)()
+ p.err = err
+ return false
+ }
+ dp.setStreamingValues(p.mem, lenUncompressed,
levelBuf, valReader, limit, closer)
+ } else {
+ p.dataPageBuffer.ResizeNoShrink(lenUncompressed)
+ buf :=
memory.NewBufferBytes(p.dataPageBuffer.Bytes())
+ data, err := p.decompress(p.r, lenCompressed,
buf.Bytes())
+ if err != nil {
+ p.err = err
+ return false
+ }
+ if len(data) != lenUncompressed {
+ p.err = fmt.Errorf("parquet: metadata
said %d bytes uncompressed data page, got %d bytes", lenUncompressed, len(data))
+ return false
+ }
+ dp.buf = buf
+ }
+ p.curPage = dp
case format.PageType_DATA_PAGE_V2:
p.pageOrd++
dataHeader := p.curPageHdr.GetDataPageHeaderV2()
@@ -882,11 +1067,7 @@ func (p *serializedPageReader) Next() bool {
return false
}
- p.dataPageBuffer.ResizeNoShrink(lenUncompressed)
- buf := memory.NewBufferBytes(p.dataPageBuffer.Bytes())
-
compressed := dataHeader.GetIsCompressed()
- // extract stats
firstRowIdx := p.rowsSeen
p.rowsSeen += int64(dataHeader.GetNumRows())
levelsBytelen, ok :=
utils.Add(int(dataHeader.GetDefinitionLevelsByteLength()),
int(dataHeader.GetRepetitionLevelsByteLength()))
@@ -894,37 +1075,17 @@ func (p *serializedPageReader) Next() bool {
p.err = errors.New("parquet: levels size too
large (corrupt file?)")
return false
}
-
- if p.cryptoCtx.DataDecryptor != nil {
- got, err := p.readV2Encrypted(p.r,
lenCompressed, lenUncompressed, levelsBytelen, compressed, buf.Bytes())
- if err != nil {
- p.err = err
- return false
- }
- if got != lenUncompressed {
- p.err = fmt.Errorf("parquet: metadata
said %d bytes uncompressed data page, got %d bytes", lenUncompressed, got)
- return false
- }
- } else {
- got, err := p.readV2Unencrypted(p.r,
lenCompressed, lenUncompressed, levelsBytelen, compressed, buf.Bytes())
- if err != nil {
- p.err = err
- return false
- }
- if got != lenUncompressed {
- p.err = fmt.Errorf("parquet: metadata
said %d bytes uncompressed data page, got %d bytes", lenUncompressed, got)
- return false
- }
+ if levelsBytelen > lenCompressed {
+ p.err = fmt.Errorf("parquet: V2 level region %d
exceeds page body %d (corrupt page?)", levelsBytelen, lenCompressed)
+ return false
}
- // make datapage v2
- p.curPage = &DataPageV2{
- page: page{
- buf: buf,
+ dp := &DataPageV2{
+ dataPageBase: dataPageBase{page: page{
typ: p.curPageHdr.Type,
nvals: dataHeader.GetNumValues(),
encoding: dataHeader.GetEncoding(),
- },
+ }},
nulls: dataHeader.GetNumNulls(),
nrows: dataHeader.GetNumRows(),
defLvlByteLen:
dataHeader.GetDefinitionLevelsByteLength(),
@@ -934,6 +1095,48 @@ func (p *serializedPageReader) Next() bool {
statistics: extractStats(dataHeader),
firstRowIndex: firstRowIdx,
}
+
+ if p.pageCanStream(dataHeader.GetEncoding(),
lenUncompressed) {
+ // V2: rep/def levels are uncompressed and
precede the separately
+ // compressed value region. Read the level
bytes directly off the front,
+ // then stream the values from what remains.
+ limit := io.LimitReader(p.r,
int64(lenCompressed))
+ levelBuf := memory.NewResizableBuffer(p.mem)
+ levelBuf.Resize(levelsBytelen)
+ if _, err := io.ReadFull(limit,
levelBuf.Bytes()); err != nil {
+ levelBuf.Release()
+ p.err = err
+ return false
+ }
+ valReader, closer := p.valueStream(limit,
compressed)
+ dp.setStreamingValues(p.mem, lenUncompressed,
levelBuf, valReader, limit, closer)
+ } else {
+ p.dataPageBuffer.ResizeNoShrink(lenUncompressed)
+ buf :=
memory.NewBufferBytes(p.dataPageBuffer.Bytes())
+ if p.cryptoCtx.DataDecryptor != nil {
+ got, err := p.readV2Encrypted(p.r,
lenCompressed, lenUncompressed, levelsBytelen, compressed, buf.Bytes())
+ if err != nil {
+ p.err = err
+ return false
+ }
+ if got != lenUncompressed {
+ p.err = fmt.Errorf("parquet:
metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed,
got)
+ return false
+ }
+ } else {
+ got, err := p.readV2Unencrypted(p.r,
lenCompressed, lenUncompressed, levelsBytelen, compressed, buf.Bytes())
+ if err != nil {
+ p.err = err
+ return false
+ }
+ if got != lenUncompressed {
+ p.err = fmt.Errorf("parquet:
metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed,
got)
+ return false
+ }
+ }
+ dp.buf = buf
+ }
+ p.curPage = dp
default:
// we don't know this page type, we're allowed to skip
non-data pages
continue
diff --git a/parquet/file/page_reader_internal_test.go
b/parquet/file/page_reader_internal_test.go
new file mode 100644
index 00000000..1ccd823e
--- /dev/null
+++ b/parquet/file/page_reader_internal_test.go
@@ -0,0 +1,54 @@
+// 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 file
+
+import (
+ "bytes"
+ "encoding/binary"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
+ "github.com/stretchr/testify/require"
+)
+
+func TestReadLevelDataRLELengthBound(t *testing.T) {
+ p := &serializedPageReader{maxDefLevel: 1, mem: memory.DefaultAllocator}
+
+ // a complete 8-byte RLE level region, so the reject below is the
bound, not EOF
+ data := make([]byte, 4+8)
+ binary.LittleEndian.PutUint32(data[:4], 8)
+
+ _, err := p.readLevelData(bytes.NewReader(data), format.Encoding_RLE,
format.Encoding_RLE, 4, 8)
+ require.Error(t, err)
+
+ out, err := p.readLevelData(bytes.NewReader(data), format.Encoding_RLE,
format.Encoding_RLE, 4, 1024)
+ require.NoError(t, err)
+ require.Equal(t, 12, out.Len())
+ out.Release()
+}
+
+func TestReadLevelDataBitPacked(t *testing.T) {
+ p := &serializedPageReader{maxDefLevel: 1, mem: memory.DefaultAllocator}
+
+ // maxDef 1 -> bitWidth 1; 8 values -> BytesForBits(8) = 1 byte, no
length prefix
+ data := []byte{0b10110100}
+ out, err := p.readLevelData(bytes.NewReader(data), format.Encoding_RLE,
format.Encoding_BIT_PACKED, 8, 1024)
+ require.NoError(t, err)
+ require.Equal(t, data, out.Bytes())
+ out.Release()
+}
diff --git a/parquet/file/row_group_reader.go b/parquet/file/row_group_reader.go
index 295ccc1c..2052a261 100644
--- a/parquet/file/row_group_reader.go
+++ b/parquet/file/row_group_reader.go
@@ -23,6 +23,7 @@ import (
"github.com/apache/arrow-go/v18/internal/utils"
"github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/compress"
"github.com/apache/arrow-go/v18/parquet/internal/encoding"
"github.com/apache/arrow-go/v18/parquet/internal/encryption"
format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
@@ -121,6 +122,7 @@ func (r *RowGroupReader) GetColumnPageReader(i int)
(PageReader, error) {
cryptoMetadata := col.CryptoMetadata()
if cryptoMetadata == nil {
+ descr := r.fileMetadata.Schema.Column(i)
pr := &serializedPageReader{
r: stream,
chunk: col,
@@ -129,6 +131,13 @@ func (r *RowGroupReader) GetColumnPageReader(i int)
(PageReader, error) {
maxPageHeaderSize: defaultMaxPageHeaderSize,
nrows: col.NumValues(),
mem: r.props.Allocator(),
+ // Streaming (PageStreamingEnabled) inputs; only the
unencrypted path is
+ // ever streaming-eligible, so columnCanStream is left
false elsewhere.
+ columnCanStream: r.props.PageStreamingEnabled &&
+ streamablePhysicalType(descr.PhysicalType()) &&
+ streamableCodec(col.Compression()),
+ maxRepLevel: descr.MaxRepetitionLevel(),
+ maxDefLevel: descr.MaxDefinitionLevel(),
}
return pr, pr.init(col.Compression(), nil)
}
@@ -186,3 +195,22 @@ func (r *RowGroupReader) GetColumnPageReader(i int)
(PageReader, error) {
}
return pr, pr.init(col.Compression(), &ctx)
}
+
+func streamablePhysicalType(t parquet.Type) bool {
+ switch t {
+ case parquet.Types.ByteArray, parquet.Types.FixedLenByteArray:
+ return true
+ default:
+ return false
+ }
+}
+
+func streamableCodec(c compress.Compression) bool {
+ switch c {
+ case compress.Codecs.Uncompressed, compress.Codecs.Gzip,
compress.Codecs.Brotli, compress.Codecs.Zstd:
+ return true
+ default:
+ // Snappy/LZ4_RAW register as StreamingCodecs but use block
framing, not a spec-compliant stream.
+ return false
+ }
+}
diff --git a/parquet/file/streaming_page_test.go
b/parquet/file/streaming_page_test.go
new file mode 100644
index 00000000..68d606e6
--- /dev/null
+++ b/parquet/file/streaming_page_test.go
@@ -0,0 +1,942 @@
+// 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 file_test
+
+import (
+ "bytes"
+ "fmt"
+ "os"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/compress"
+ "github.com/apache/arrow-go/v18/parquet/file"
+ "github.com/apache/arrow-go/v18/parquet/schema"
+ "github.com/stretchr/testify/require"
+)
+
+// TestMain lowers the streaming page-size threshold so the streaming tests
exercise
+// the streaming path on their small (fast) pages.
+func TestMain(m *testing.M) {
+ *file.StreamingThreshold = 0
+ os.Exit(m.Run())
+}
+
+func makeStreamTestValues(sizes []int) []parquet.ByteArray {
+ values := make([]parquet.ByteArray, len(sizes))
+ for i, sz := range sizes {
+ buf := make([]byte, sz)
+ buf[0] = byte(i)
+ for j := 1; j < sz; j++ {
+ buf[j] = byte(i*31 + j)
+ }
+ values[i] = buf
+ }
+ return values
+}
+
+func writeByteArrayStream(t *testing.T, values []parquet.ByteArray, defLevels,
repLevels []int16, rep parquet.Repetition, ver parquet.DataPageVersion, codec
compress.Compression, pageSize int64) []byte {
+ t.Helper()
+ sc := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
+ schema.Must(schema.NewPrimitiveNode("data", rep,
parquet.Types.ByteArray, -1, -1)),
+ }, -1)))
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(false), // ensure PLAIN data pages
+ parquet.WithDataPageVersion(ver),
+ parquet.WithCompression(codec),
+ parquet.WithDataPageSize(pageSize),
+ )
+ out := &bytes.Buffer{}
+ w := file.NewParquetWriter(out, sc.Root(), file.WithWriterProps(props))
+ rgw := w.AppendRowGroup()
+ cw, err := rgw.NextColumn()
+ require.NoError(t, err)
+ _, err = cw.(*file.ByteArrayColumnChunkWriter).WriteBatch(values,
defLevels, repLevels)
+ require.NoError(t, err)
+ require.NoError(t, cw.Close())
+ require.NoError(t, rgw.Close())
+ require.NoError(t, w.Close())
+ return out.Bytes()
+}
+
+// writeStreamTestColumn writes a Required byte-array column with small (1
KiB) pages.
+func writeStreamTestColumn(t *testing.T, values []parquet.ByteArray, ver
parquet.DataPageVersion, codec compress.Compression) []byte {
+ return writeByteArrayStream(t, values, nil, nil,
parquet.Repetitions.Required, ver, codec, 1024)
+}
+
+// collectByteArray reads exactly n values from bar via repeated ReadBatch
calls.
+func collectByteArray(t *testing.T, bar *file.ByteArrayColumnChunkReader, n
int) []parquet.ByteArray {
+ t.Helper()
+ out := make([]parquet.ByteArray, n)
+ total := 0
+ for total < n {
+ _, nv, err := bar.ReadBatch(int64(n-total), out[total:], nil,
nil)
+ require.NoError(t, err)
+ if nv == 0 {
+ break
+ }
+ total += nv
+ }
+ require.Equal(t, n, total)
+ return out
+}
+
+func readStreamTestColumn(t *testing.T, data []byte, streaming bool)
[]parquet.ByteArray {
+ t.Helper()
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = streaming
+
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.ByteArrayColumnChunkReader)
+ return collectByteArray(t, bar, int(rdr.NumRows()))
+}
+
+func TestPageStreamingByteArrayRoundTrip(t *testing.T) {
+ // Varied sizes: several small, some > 4096 (stream buffer), some >
1024 (page).
+ values := makeStreamTestValues([]int{65, 100, 5000, 200, 9000, 50,
4096, 1})
+
+ cases := []struct {
+ name string
+ ver parquet.DataPageVersion
+ codec compress.Compression
+ }{
+ {"v1-uncompressed", parquet.DataPageV1,
compress.Codecs.Uncompressed},
+ {"v1-zstd", parquet.DataPageV1, compress.Codecs.Zstd},
+ {"v1-gzip", parquet.DataPageV1, compress.Codecs.Gzip},
+ {"v2-uncompressed", parquet.DataPageV2,
compress.Codecs.Uncompressed},
+ {"v2-zstd", parquet.DataPageV2, compress.Codecs.Zstd},
+ {"v2-brotli", parquet.DataPageV2, compress.Codecs.Brotli},
+ }
+
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ data := writeStreamTestColumn(t, values, tc.ver,
tc.codec)
+
+ materialized := readStreamTestColumn(t, data, false)
+ streamed := readStreamTestColumn(t, data, true)
+
+ require.Len(t, materialized, len(values))
+ for i := range values {
+ require.Truef(t, bytes.Equal(values[i],
materialized[i]), "materialized value %d mismatch", i)
+ require.Truef(t, bytes.Equal(values[i],
streamed[i]), "streamed value %d mismatch", i)
+ }
+ })
+ }
+}
+
+// TestPageStreamingEngaged guards against eligibility silently falling back
to the
+// whole-page read: a required column's streaming page holds only the (empty)
level
+// region, while the materialized page's Data() is the full uncompressed page.
+func TestPageStreamingEngaged(t *testing.T) {
+ values := makeStreamTestValues([]int{65, 5000, 200, 9000, 50})
+ data := writeStreamTestColumn(t, values, parquet.DataPageV1,
compress.Codecs.Zstd)
+
+ firstPageDataLen := func(streaming bool) int {
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = streaming
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ pr, err := rdr.RowGroup(0).GetColumnPageReader(0)
+ require.NoError(t, err)
+ defer pr.Close()
+ require.True(t, pr.Next())
+ return len(pr.Page().Data())
+ }
+
+ require.Zero(t, firstPageDataLen(true), "streaming page Data() should
hold only the empty level region")
+ require.NotZero(t, firstPageDataLen(false), "materialized page Data()
should be the whole page")
+}
+
+func TestPageStreamingAllocatorBalance(t *testing.T) {
+ values := makeStreamTestValues([]int{65, 5000, 200, 9000, 50})
+ data := writeStreamTestColumn(t, values, parquet.DataPageV1,
compress.Codecs.Zstd)
+
+ checked := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer checked.AssertSize(t, 0)
+
+ props := parquet.NewReaderProperties(checked)
+ props.PageStreamingEnabled = true
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ n := int(rdr.NumRows())
+ out := make([]parquet.ByteArray, n)
+ for total := 0; total < n; {
+ _, nv, err :=
cr.(*file.ByteArrayColumnChunkReader).ReadBatch(int64(n-total), out[total:],
nil, nil)
+ require.NoError(t, err)
+ if nv == 0 {
+ break
+ }
+ total += nv
+ }
+ require.NoError(t, cr.Close())
+}
+
+func TestPageStreamingLeveledAllocatorBalance(t *testing.T) {
+ nullDef := make([]int16, 40)
+ var nullVals []parquet.ByteArray
+ for i := 0; i < 40; i++ {
+ if i%3 == 0 {
+ continue
+ }
+ nullDef[i] = 1
+ buf := make([]byte, 40+i*120)
+ for j := range buf {
+ buf[j] = byte(i*7 + j)
+ }
+ nullVals = append(nullVals, buf)
+ }
+
+ cases := []struct {
+ name string
+ rep parquet.Repetition
+ values []parquet.ByteArray
+ defLevels []int16
+ repLevels []int16
+ }{
+ {"nullable", parquet.Repetitions.Optional, nullVals, nullDef,
nil},
+ {"repeated", parquet.Repetitions.Repeated,
+ makeStreamTestValues([]int{40, 5000, 60, 9000, 70, 80}),
+ []int16{1, 1, 1, 1, 1, 1}, []int16{0, 1, 1, 0, 0, 1}},
+ }
+ for _, tc := range cases {
+ for _, ver := range
[]parquet.DataPageVersion{parquet.DataPageV1, parquet.DataPageV2} {
+ t.Run(fmt.Sprintf("%s/v%d", tc.name, ver+1), func(t
*testing.T) {
+ data := writeByteArrayStream(t, tc.values,
tc.defLevels, tc.repLevels, tc.rep, ver, compress.Codecs.Zstd, 1024)
+
+ checked :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer checked.AssertSize(t, 0)
+
+ props := parquet.NewReaderProperties(checked)
+ props.PageStreamingEnabled = true
+ rdr, err :=
file.NewParquetReader(bytes.NewReader(data), file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.ByteArrayColumnChunkReader)
+ vbuf := make([]parquet.ByteArray, 64)
+ dbuf := make([]int16, 64)
+ rbuf := make([]int16, 64)
+ for {
+ total, _, err := bar.ReadBatch(64,
vbuf, dbuf, rbuf)
+ require.NoError(t, err)
+ if total == 0 {
+ break
+ }
+ }
+ require.NoError(t, cr.Close())
+ })
+ }
+ }
+}
+
+func TestPageStreamingIneligibleCodecFallback(t *testing.T) {
+ values := makeStreamTestValues([]int{65, 5000, 200, 9000, 50})
+
+ for _, codec := range []compress.Compression{compress.Codecs.Snappy,
compress.Codecs.Lz4Raw} {
+ for _, ver := range
[]parquet.DataPageVersion{parquet.DataPageV1, parquet.DataPageV2} {
+ data := writeStreamTestColumn(t, values, ver, codec)
+
+ // PageStreamingEnabled is set, but Snappy/LZ4_RAW are
not in the allowlist,
+ // so pages must fall back to whole-page reads and
match materialized.
+ materialized := readStreamTestColumn(t, data, false)
+ streamed := readStreamTestColumn(t, data, true)
+
+ require.Len(t, streamed, len(values))
+ for i := range values {
+ require.Truef(t, bytes.Equal(values[i],
streamed[i]), "codec %v ver %v value %d", codec, ver, i)
+ require.Truef(t, bytes.Equal(materialized[i],
streamed[i]), "codec %v ver %v value %d: streamed != materialized", codec, ver,
i)
+ }
+ }
+ }
+}
+
+func readNullableStreamColumn(t *testing.T, data []byte, streaming bool,
numRows int) (defLevels []int16, values []parquet.ByteArray) {
+ t.Helper()
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = streaming
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.ByteArrayColumnChunkReader)
+
+ defLevels = make([]int16, 0, numRows)
+ values = make([]parquet.ByteArray, 0, numRows)
+ for len(defLevels) < numRows {
+ vbuf := make([]parquet.ByteArray, numRows)
+ dbuf := make([]int16, numRows)
+ total, nv, err := bar.ReadBatch(int64(numRows-len(defLevels)),
vbuf, dbuf, nil)
+ require.NoError(t, err)
+ if total == 0 {
+ break
+ }
+ defLevels = append(defLevels, dbuf[:total]...)
+ values = append(values, vbuf[:nv]...)
+ }
+ return defLevels, values
+}
+
+func TestPageStreamingByteArrayNullable(t *testing.T) {
+ const numRows = 40
+ defLevels := make([]int16, numRows)
+ var values []parquet.ByteArray
+ for i := 0; i < numRows; i++ {
+ if i%3 == 0 {
+ defLevels[i] = 0 // null
+ continue
+ }
+ defLevels[i] = 1
+ sz := 40 + i*120 // varied sizes, some > page / stream buffer
+ buf := make([]byte, sz)
+ buf[0] = byte(i)
+ for j := 1; j < sz; j++ {
+ buf[j] = byte(i*7 + j)
+ }
+ values = append(values, buf)
+ }
+
+ cases := []struct {
+ name string
+ ver parquet.DataPageVersion
+ codec compress.Compression
+ }{
+ {"v1-zstd", parquet.DataPageV1, compress.Codecs.Zstd},
+ {"v2-gzip", parquet.DataPageV2, compress.Codecs.Gzip},
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ data := writeByteArrayStream(t, values, defLevels, nil,
parquet.Repetitions.Optional, tc.ver, tc.codec, 1024)
+
+ mDef, mVals := readNullableStreamColumn(t, data, false,
numRows)
+ sDef, sVals := readNullableStreamColumn(t, data, true,
numRows)
+
+ require.Equal(t, mDef, sDef, "def levels differ")
+ require.Equal(t, defLevels, sDef, "streamed def levels
wrong")
+ require.Equal(t, len(values), len(sVals))
+ for i := range values {
+ require.Truef(t, bytes.Equal(mVals[i],
sVals[i]), "value %d differs streaming vs materialized", i)
+ require.Truef(t, bytes.Equal(values[i],
sVals[i]), "streamed value %d wrong", i)
+ }
+ })
+ }
+}
+
+func readStreamTestColumnSkip(t *testing.T, data []byte, streaming bool, skip
int64) []parquet.ByteArray {
+ t.Helper()
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = streaming
+
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.ByteArrayColumnChunkReader)
+
+ skipped, err := bar.Skip(skip)
+ require.NoError(t, err)
+ require.Equal(t, skip, skipped)
+
+ return collectByteArray(t, bar, int(rdr.NumRows())-int(skip))
+}
+
+func TestPageStreamingByteArraySkip(t *testing.T) {
+ // several small values share a page so the skip discards within a
streaming page,
+ // with some large values spanning pages too.
+ values := makeStreamTestValues([]int{40, 50, 60, 70, 80, 90, 5000, 100,
9000, 40, 55, 66})
+
+ cases := []struct {
+ name string
+ ver parquet.DataPageVersion
+ codec compress.Compression
+ }{
+ {"v1-zstd", parquet.DataPageV1, compress.Codecs.Zstd},
+ {"v2-gzip", parquet.DataPageV2, compress.Codecs.Gzip},
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ data := writeStreamTestColumn(t, values, tc.ver,
tc.codec)
+
+ for _, skip := range []int64{3, 8} { // 3 within page
1; 8 crosses page boundaries
+ materialized := readStreamTestColumnSkip(t,
data, false, skip)
+ streamed := readStreamTestColumnSkip(t, data,
true, skip)
+
+ require.Len(t, streamed, len(values)-int(skip))
+ for i := range streamed {
+ require.Truef(t,
bytes.Equal(values[int(skip)+i], streamed[i]), "skip %d value %d", skip, i)
+ require.Truef(t,
bytes.Equal(materialized[i], streamed[i]), "skip %d value %d streamed !=
materialized", skip, i)
+ }
+ }
+ })
+ }
+}
+
+func readRepeatedStreamColumn(t *testing.T, data []byte, streaming bool)
(defLevels, repLevels []int16, values []parquet.ByteArray) {
+ t.Helper()
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = streaming
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.ByteArrayColumnChunkReader)
+
+ const batchSize = 64
+ for {
+ vbuf := make([]parquet.ByteArray, batchSize)
+ dbuf := make([]int16, batchSize)
+ rbuf := make([]int16, batchSize)
+ total, nv, err := bar.ReadBatch(batchSize, vbuf, dbuf, rbuf)
+ require.NoError(t, err)
+ if total == 0 {
+ break
+ }
+ defLevels = append(defLevels, dbuf[:total]...)
+ repLevels = append(repLevels, rbuf[:total]...)
+ values = append(values, vbuf[:nv]...)
+ }
+ return defLevels, repLevels, values
+}
+
+func TestPageStreamingRepeated(t *testing.T) {
+ // lists: [v0 v1 v2] [v3] [v4 v5]; rep levels mark list boundaries.
+ values := makeStreamTestValues([]int{40, 5000, 60, 9000, 70, 80})
+ repLevels := []int16{0, 1, 1, 0, 0, 1}
+ defLevels := []int16{1, 1, 1, 1, 1, 1}
+
+ cases := []struct {
+ name string
+ ver parquet.DataPageVersion
+ codec compress.Compression
+ }{
+ {"v1-zstd", parquet.DataPageV1, compress.Codecs.Zstd},
+ {"v2-gzip", parquet.DataPageV2, compress.Codecs.Gzip},
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ data := writeByteArrayStream(t, values, defLevels,
repLevels, parquet.Repetitions.Repeated, tc.ver, tc.codec, 1024)
+
+ mDef, mRep, mVals := readRepeatedStreamColumn(t, data,
false)
+ sDef, sRep, sVals := readRepeatedStreamColumn(t, data,
true)
+
+ require.Equal(t, defLevels, sDef)
+ require.Equal(t, repLevels, sRep)
+ require.Equal(t, mDef, sDef)
+ require.Equal(t, mRep, sRep)
+ require.Len(t, sVals, len(values))
+ for i := range values {
+ require.Truef(t, bytes.Equal(values[i],
sVals[i]), "value %d", i)
+ require.Truef(t, bytes.Equal(mVals[i],
sVals[i]), "value %d streamed != materialized", i)
+ }
+ })
+ }
+}
+
+func readStreamTestColumnSeekAfterRead(t *testing.T, data []byte, streaming
bool, preRead int, seekRow int64) []parquet.ByteArray {
+ t.Helper()
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = streaming
+
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.ByteArrayColumnChunkReader)
+
+ // Read a few values first so a streaming page is current and partially
consumed
+ // when the seek repositions the underlying reader.
+ if preRead > 0 {
+ pre := make([]parquet.ByteArray, preRead)
+ _, _, err := bar.ReadBatch(int64(preRead), pre, nil, nil)
+ require.NoError(t, err)
+ }
+ require.NoError(t, bar.SeekToRow(seekRow))
+
+ return collectByteArray(t, bar, int(rdr.NumRows())-int(seekRow))
+}
+
+func TestPageStreamingSeekAfterRead(t *testing.T) {
+ values := makeStreamTestValues([]int{40, 50, 60, 70, 5000, 90, 9000,
40, 55, 66, 77, 88})
+ const (
+ preRead = 2
+ seekRow = int64(8)
+ )
+
+ cases := []struct {
+ name string
+ ver parquet.DataPageVersion
+ codec compress.Compression
+ }{
+ {"v1-zstd", parquet.DataPageV1, compress.Codecs.Zstd},
+ {"v2-gzip", parquet.DataPageV2, compress.Codecs.Gzip},
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ data := writeStreamTestColumn(t, values, tc.ver,
tc.codec)
+
+ mat := readStreamTestColumnSeekAfterRead(t, data,
false, preRead, seekRow)
+ str := readStreamTestColumnSeekAfterRead(t, data, true,
preRead, seekRow)
+
+ require.Len(t, str, len(values)-int(seekRow))
+ for i := range str {
+ require.Truef(t,
bytes.Equal(values[int(seekRow)+i], str[i]), "value %d after seek", i)
+ require.Truef(t, bytes.Equal(mat[i], str[i]),
"value %d streamed != materialized", i)
+ }
+ })
+ }
+}
+
+// Large pages (> stream buffer) partially read leave undrained bytes at
SeekToRow —
+// the case that actually triggers the seek drain fix (unlike the small-page
test).
+func TestPageStreamingSeekLargePagePartial(t *testing.T) {
+ sizes := make([]int, 40)
+ for i := range sizes {
+ sizes[i] = 64 << 10 // 64 KiB -> ~16 values per >1 MiB page
+ }
+ values := makeStreamTestValues(sizes)
+ data := writeByteArrayStream(t, values, nil, nil,
parquet.Repetitions.Required, parquet.DataPageV1, compress.Codecs.Uncompressed,
1<<20)
+
+ const (
+ preRead = 20 // into page 2
+ seekRow = int64(35)
+ )
+ mat := readStreamTestColumnSeekAfterRead(t, data, false, preRead,
seekRow)
+ str := readStreamTestColumnSeekAfterRead(t, data, true, preRead,
seekRow)
+
+ require.Len(t, str, len(values)-int(seekRow))
+ for i := range str {
+ require.Truef(t, bytes.Equal(values[int(seekRow)+i], str[i]),
"value %d after seek", i)
+ require.Truef(t, bytes.Equal(mat[i], str[i]), "value %d
streamed != materialized", i)
+ }
+}
+
+func TestPageStreamingPageReaderCloseFreesStream(t *testing.T) {
+ values := makeStreamTestValues([]int{65, 5000, 200})
+ data := writeStreamTestColumn(t, values, parquet.DataPageV1,
compress.Codecs.Zstd)
+
+ checked := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer checked.AssertSize(t, 0)
+
+ props := parquet.NewReaderProperties(checked)
+ props.PageStreamingEnabled = true
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ pr, err := rdr.RowGroup(0).GetColumnPageReader(0)
+ require.NoError(t, err)
+ require.True(t, pr.Next()) // advance to a streaming data page
+ require.NoError(t, pr.Close())
+}
+
+func makeFLBAValues(count, width int) []parquet.FixedLenByteArray {
+ vals := make([]parquet.FixedLenByteArray, count)
+ for i := range vals {
+ b := make([]byte, width)
+ b[0] = byte(i)
+ for j := 1; j < width; j++ {
+ b[j] = byte(i*13 + j)
+ }
+ vals[i] = b
+ }
+ return vals
+}
+
+func writeFLBAStreamColumn(t *testing.T, values []parquet.FixedLenByteArray,
width int, ver parquet.DataPageVersion, codec compress.Compression) []byte {
+ t.Helper()
+ sc := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
+ schema.Must(schema.NewPrimitiveNode("data",
parquet.Repetitions.Required, parquet.Types.FixedLenByteArray, -1,
int32(width))),
+ }, -1)))
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(false),
+ parquet.WithDataPageVersion(ver),
+ parquet.WithCompression(codec),
+ parquet.WithDataPageSize(1024),
+ )
+ out := &bytes.Buffer{}
+ w := file.NewParquetWriter(out, sc.Root(), file.WithWriterProps(props))
+ rgw := w.AppendRowGroup()
+ cw, err := rgw.NextColumn()
+ require.NoError(t, err)
+ _, err =
cw.(*file.FixedLenByteArrayColumnChunkWriter).WriteBatch(values, nil, nil)
+ require.NoError(t, err)
+ require.NoError(t, cw.Close())
+ require.NoError(t, rgw.Close())
+ require.NoError(t, w.Close())
+ return out.Bytes()
+}
+
+func readFLBAStreamColumnSkip(t *testing.T, data []byte, streaming bool, skip
int64) []parquet.FixedLenByteArray {
+ t.Helper()
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = streaming
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.FixedLenByteArrayColumnChunkReader)
+ if skip > 0 {
+ skipped, err := bar.Skip(skip)
+ require.NoError(t, err)
+ require.Equal(t, skip, skipped)
+ }
+ n := int(rdr.NumRows()) - int(skip)
+ out := make([]parquet.FixedLenByteArray, n)
+ total := 0
+ for total < n {
+ _, nv, err := bar.ReadBatch(int64(n-total), out[total:], nil,
nil)
+ require.NoError(t, err)
+ if nv == 0 {
+ break
+ }
+ total += nv
+ }
+ require.Equal(t, n, total)
+ return out
+}
+
+func TestPageStreamingFLBASkip(t *testing.T) {
+ const width = 24
+ values := makeFLBAValues(80, width) // multiple 1 KiB pages
+
+ cases := []struct {
+ name string
+ ver parquet.DataPageVersion
+ codec compress.Compression
+ }{
+ {"v1-zstd", parquet.DataPageV1, compress.Codecs.Zstd},
+ {"v2-gzip", parquet.DataPageV2, compress.Codecs.Gzip},
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ data := writeFLBAStreamColumn(t, values, width, tc.ver,
tc.codec)
+
+ full := readFLBAStreamColumnSkip(t, data, true, 0)
+ require.Len(t, full, len(values))
+ for i := range values {
+ require.Truef(t, bytes.Equal(values[i],
full[i]), "roundtrip value %d", i)
+ }
+
+ for _, skip := range []int64{5, 50} { // 50 crosses a
page boundary
+ mat := readFLBAStreamColumnSkip(t, data, false,
skip)
+ str := readFLBAStreamColumnSkip(t, data, true,
skip)
+ require.Len(t, str, len(values)-int(skip))
+ for i := range str {
+ require.Truef(t,
bytes.Equal(values[int(skip)+i], str[i]), "skip %d value %d", skip, i)
+ require.Truef(t, bytes.Equal(mat[i],
str[i]), "skip %d value %d streamed != materialized", skip, i)
+ }
+ }
+ })
+ }
+}
+
+// TestPageStreamingPeakMemoryBounded is the core memory-goal check: a single
page
+// far larger than the stream buffer, read in small batches, must keep the
+// allocator footprint near the buffer, not grow to the whole page.
+func TestPageStreamingPeakMemoryBounded(t *testing.T) {
+ sizes := make([]int, 2000)
+ for i := range sizes {
+ sizes[i] = 2048 // ~4 MiB of values in one page
+ }
+ values := makeStreamTestValues(sizes)
+ data := writeByteArrayStream(t, values, nil, nil,
parquet.Repetitions.Required, parquet.DataPageV1, compress.Codecs.Zstd, 8<<20)
+
+ checked := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer checked.AssertSize(t, 0)
+ props := parquet.NewReaderProperties(checked)
+ props.PageStreamingEnabled = true
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.ByteArrayColumnChunkReader)
+
+ const batch = 64
+ out := make([]parquet.ByteArray, batch)
+ peak, total := 0, 0
+ for {
+ _, nv, err := bar.ReadBatch(batch, out, nil, nil)
+ require.NoError(t, err)
+ if nv == 0 {
+ break
+ }
+ total += nv
+ if a := checked.CurrentAlloc(); a > peak {
+ peak = a
+ }
+ }
+ require.NoError(t, bar.Close())
+ require.Equal(t, len(values), total)
+
+ t.Logf("peak allocator use: %d bytes (page ~%d bytes)", peak,
len(values)*2048)
+ require.Less(t, peak, 2<<20, "streaming peak should stay near the
buffer, not the page")
+}
+
+func readInPage(t *testing.T, data []byte, batch int64) (defLevels, repLevels
[]int16, values []parquet.ByteArray) {
+ t.Helper()
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = true
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ cr, err := rdr.RowGroup(0).Column(0)
+ require.NoError(t, err)
+ bar := cr.(*file.ByteArrayColumnChunkReader)
+
+ dbuf := make([]int16, batch)
+ rbuf := make([]int16, batch)
+ vbuf := make([]parquet.ByteArray, batch)
+ for {
+ total, nv, err := bar.ReadBatchInPage(batch, vbuf, dbuf, rbuf)
+ require.NoError(t, err)
+ if total == 0 {
+ break
+ }
+ defLevels = append(defLevels, dbuf[:total]...)
+ repLevels = append(repLevels, rbuf[:total]...)
+ for i := 0; i < nv; i++ {
+ values = append(values, append([]byte(nil),
vbuf[i]...)) // clone the alias
+ }
+ }
+ return defLevels, repLevels, values
+}
+
+// Large values force per-page batch clipping (multiple Decodes) on a nullable
column.
+func TestPageStreamingReadBatchInPageNullable(t *testing.T) {
+ const numRows = 60
+ defLevels := make([]int16, numRows)
+ var values []parquet.ByteArray
+ for i := 0; i < numRows; i++ {
+ if i%3 == 0 {
+ continue // null
+ }
+ defLevels[i] = 1
+ sz := 80 << 10 // 80 KiB: several per 1 MiB window, so
DecodeStreaming returns short
+ b := make([]byte, sz)
+ b[0] = byte(i)
+ for j := 1; j < sz; j++ {
+ b[j] = byte(i*7 + j)
+ }
+ values = append(values, b)
+ }
+
+ for _, tc := range []struct {
+ name string
+ ver parquet.DataPageVersion
+ codec compress.Compression
+ }{
+ {"v1-zstd", parquet.DataPageV1, compress.Codecs.Zstd},
+ {"v2-gzip", parquet.DataPageV2, compress.Codecs.Gzip},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ data := writeByteArrayStream(t, values, defLevels, nil,
parquet.Repetitions.Optional, tc.ver, tc.codec, 8<<20)
+ for _, batch := range []int64{7, 64} {
+ gotDef, _, gotVals := readInPage(t, data, batch)
+ require.Equalf(t, defLevels, gotDef, "batch %d
def levels", batch)
+ require.Lenf(t, gotVals, len(values), "batch %d
value count", batch)
+ for i := range values {
+ require.Truef(t, bytes.Equal(values[i],
gotVals[i]), "batch %d value %d", batch, i)
+ }
+ }
+ })
+ }
+}
+
+func TestPageStreamingReadBatchInPageRequired(t *testing.T) {
+ sizes := make([]int, 40)
+ for i := range sizes {
+ sizes[i] = 80 << 10 // large, so DecodeStreaming returns short
against the window
+ }
+ values := makeStreamTestValues(sizes)
+ data := writeByteArrayStream(t, values, nil, nil,
parquet.Repetitions.Required, parquet.DataPageV1, compress.Codecs.Zstd, 8<<20)
+
+ for _, batch := range []int64{5, 64} {
+ _, _, got := readInPage(t, data, batch)
+ require.Lenf(t, got, len(values), "batch %d value count", batch)
+ for i := range values {
+ require.Truef(t, bytes.Equal(values[i], got[i]), "batch
%d value %d", batch, i)
+ }
+ }
+}
+
+func TestPageStreamingReadBatchInPageRepeated(t *testing.T) {
+ sizes := make([]int, 30)
+ for i := range sizes {
+ sizes[i] = 80 << 10
+ }
+ values := makeStreamTestValues(sizes)
+ defLevels := make([]int16, len(values))
+ repLevels := make([]int16, len(values))
+ for i := range values {
+ defLevels[i] = 1
+ if i%3 != 0 {
+ repLevels[i] = 1 // 0 starts a new list, 1 continues it
+ }
+ }
+ data := writeByteArrayStream(t, values, defLevels, repLevels,
parquet.Repetitions.Repeated, parquet.DataPageV1, compress.Codecs.Zstd, 8<<20)
+
+ for _, batch := range []int64{7, 64} {
+ gotDef, gotRep, gotVals := readInPage(t, data, batch)
+ require.Equalf(t, defLevels, gotDef, "batch %d def levels",
batch)
+ require.Equalf(t, repLevels, gotRep, "batch %d rep levels",
batch)
+ require.Lenf(t, gotVals, len(values), "batch %d value count",
batch)
+ for i := range values {
+ require.Truef(t, bytes.Equal(values[i], gotVals[i]),
"batch %d value %d", batch, i)
+ }
+ }
+}
+
+// TestPageStreamingThreshold verifies streamingThreshold gates streaming: a
page below
+// the threshold falls back to a whole-page read.
+func TestPageStreamingThreshold(t *testing.T) {
+ values := makeStreamTestValues([]int{65, 5000, 200}) // small page,
well under 1 MiB
+ data := writeStreamTestColumn(t, values, parquet.DataPageV1,
compress.Codecs.Zstd)
+
+ firstPageDataLen := func(threshold int64) int {
+ defer func(old int64) { *file.StreamingThreshold = old
}(*file.StreamingThreshold)
+ *file.StreamingThreshold = threshold
+
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = true
+ rdr, err := file.NewParquetReader(bytes.NewReader(data),
file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ pr, err := rdr.RowGroup(0).GetColumnPageReader(0)
+ require.NoError(t, err)
+ defer pr.Close()
+ require.True(t, pr.Next())
+ return len(pr.Page().Data())
+ }
+
+ require.NotZero(t, firstPageDataLen(1<<20), "small page must not stream
at the 1 MiB threshold")
+ require.Zero(t, firstPageDataLen(0), "with threshold 0 the small page
streams (empty level region)")
+}
+
+// TestPageStreamingRecordReader drives the record reader (the pqarrow read
path) over a
+// streaming column in small batches, so a Decode's aliased values must
survive until the
+// builder copies them and be reclaimed by the next Decode's Recycle. Large
values force
+// chunk rotation, so a broken alias contract would corrupt the output.
+func TestPageStreamingRecordReader(t *testing.T) {
+ cases := []struct {
+ name string
+ rep parquet.Repetition
+ defLevel int16
+ }{
+ {"required", parquet.Repetitions.Required, 0},
+ {"optional", parquet.Repetitions.Optional, 1},
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ const nrows = 300
+ var (
+ sizes []int
+ defLevels []int16
+ wantNull []bool
+ )
+ for i := 0; i < nrows; i++ {
+ if tc.defLevel == 1 && i%5 == 0 {
+ defLevels = append(defLevels, 0)
+ wantNull = append(wantNull, true)
+ continue
+ }
+ sizes = append(sizes, 3000) // large enough to
force chunk rotation
+ defLevels = append(defLevels, tc.defLevel)
+ wantNull = append(wantNull, false)
+ }
+ values := makeStreamTestValues(sizes)
+ var dl []int16
+ if tc.defLevel == 1 {
+ dl = defLevels
+ }
+ data := writeByteArrayStream(t, values, dl, nil,
tc.rep, parquet.DataPageV1, compress.Codecs.Zstd, 8<<20)
+
+ props :=
parquet.NewReaderProperties(memory.DefaultAllocator)
+ props.PageStreamingEnabled = true
+ rdr, err :=
file.NewParquetReader(bytes.NewReader(data), file.WithReadProps(props))
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ pr, err := rdr.RowGroup(0).GetColumnPageReader(0)
+ require.NoError(t, err)
+ descr := rdr.MetaData().Schema.Column(0)
+
+ rr := file.NewRecordReader(descr,
file.LevelInfo{NullSlotUsage: 1, DefLevel: tc.defLevel},
arrow.BinaryTypes.Binary, memory.DefaultAllocator, rdr.BufferPool())
+ defer rr.Release()
+ rr.SetPageReader(pr)
+ rr.Reset()
+ require.NoError(t, rr.Reserve(nrows))
+ for rr.HasMore() {
+ read, err := rr.ReadRecords(37) // small batch
=> many Decode/Recycle boundaries
+ require.NoError(t, err)
+ if read == 0 {
+ break
+ }
+ }
+
+ var got []parquet.ByteArray
+ var gotNull []bool
+ for _, chunk := range
rr.(file.BinaryRecordReader).GetBuilderChunks() {
+ bin := chunk.(*array.Binary)
+ for i := 0; i < bin.Len(); i++ {
+ gotNull = append(gotNull, bin.IsNull(i))
+ got = append(got, bin.Value(i))
+ }
+ chunk.Release()
+ }
+
+ require.Equal(t, nrows, len(got))
+ require.Equal(t, wantNull, gotNull)
+ vi := 0
+ for i := 0; i < nrows; i++ {
+ if wantNull[i] {
+ continue
+ }
+ require.Truef(t, bytes.Equal(values[vi],
got[i]), "value at row %d corrupted", i)
+ vi++
+ }
+ })
+ }
+}
diff --git a/parquet/internal/encoding/byte_array_decoder.go
b/parquet/internal/encoding/byte_array_decoder.go
index e568ef92..10112b07 100644
--- a/parquet/internal/encoding/byte_array_decoder.go
+++ b/parquet/internal/encoding/byte_array_decoder.go
@@ -19,12 +19,14 @@ package encoding
import (
"encoding/binary"
"errors"
+ "fmt"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/internal/utils"
"github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding/streaming"
pqutils "github.com/apache/arrow-go/v18/parquet/internal/utils"
)
@@ -37,6 +39,9 @@ import (
// bytes being the raw data of the byte array.
type PlainByteArrayDecoder struct {
decoder
+
+ // src drives the streaming path (PageStreamingEnabled) when non-nil.
+ src streaming.ValueBuffer
}
// Type returns parquet.Types.ByteArray for this decoder
@@ -44,7 +49,70 @@ func (PlainByteArrayDecoder) Type() parquet.Type {
return parquet.Types.ByteArray
}
+func (pbad *PlainByteArrayDecoder) SetSource(nvals int, src
streaming.ValueBuffer) {
+ pbad.src, pbad.nvals = src, nvals
+}
+
+func (pbad *PlainByteArrayDecoder) SetData(nvals int, data []byte) error {
+ pbad.src = nil
+ return pbad.decoder.SetData(nvals, data)
+}
+
+// decodeStreaming fills out by aliasing each value in the buffer (no copy).
Recycle at
+// entry reclaims the previous batch, so the returned slices are valid only
until the
+// next Decode call; callers consume or copy them before then.
+func (pbad *PlainByteArrayDecoder) decodeStreaming(out []parquet.ByteArray)
(int, error) {
+ pbad.src.Recycle()
+ max := utils.Min(len(out), pbad.nvals)
+ for i := 0; i < max; i++ {
+ hdr, err := pbad.src.Fill(4)
+ if err != nil {
+ return i, fmt.Errorf("parquet: reading bytearray: %w",
err)
+ }
+ byteLen := int(int32(binary.LittleEndian.Uint32(hdr[:4])))
+ if byteLen < 0 {
+ return i, errors.New("parquet: invalid BYTE_ARRAY
value")
+ }
+ pbad.src.Advance(4)
+ // Fill rejects a byteLen past the value region before
allocating.
+ val, err := pbad.src.Fill(byteLen)
+ if err != nil {
+ return i, fmt.Errorf("parquet: reading bytearray: %w",
err)
+ }
+ // 3-index cap stops appends bleeding into the next value; a
zero-length slice
+ // of the non-nil buffer stays non-nil (empty value != null).
+ out[i] = val[:byteLen:byteLen]
+ pbad.src.Advance(byteLen)
+ }
+ pbad.nvals -= max
+ return max, nil
+}
+
+func (pbad *PlainByteArrayDecoder) discardStreaming(n int) (int, error) {
+ n = min(n, pbad.nvals)
+ for i := 0; i < n; i++ {
+ pbad.src.Recycle()
+ buf, err := pbad.src.Fill(4)
+ if err != nil {
+ return i, fmt.Errorf("parquet: skipping bytearray: %w",
err)
+ }
+ valueLen := int(int32(binary.LittleEndian.Uint32(buf[:4])))
+ if valueLen < 0 {
+ return i, errors.New("parquet: invalid BYTE_ARRAY
value")
+ }
+ pbad.src.Advance(4)
+ if err := pbad.src.Skip(valueLen); err != nil {
+ return i, fmt.Errorf("parquet: skipping bytearray: %w",
err)
+ }
+ }
+ pbad.nvals -= n
+ return n, nil
+}
+
func (pbad *PlainByteArrayDecoder) Discard(n int) (int, error) {
+ if pbad.src != nil {
+ return pbad.discardStreaming(n)
+ }
n = min(n, pbad.nvals)
// we have to skip the length of each value by first checking
// the length of the value and then skipping that many bytes
@@ -74,6 +142,10 @@ func (pbad *PlainByteArrayDecoder) Discard(n int) (int,
error) {
//
// Returns the number of values that were decoded.
func (pbad *PlainByteArrayDecoder) Decode(out []parquet.ByteArray) (int,
error) {
+ if pbad.src != nil {
+ return pbad.decodeStreaming(out)
+ }
+
max := utils.Min(len(out), pbad.nvals)
for i := 0; i < max; i++ {
diff --git a/parquet/internal/encoding/fixed_len_byte_array_decoder.go
b/parquet/internal/encoding/fixed_len_byte_array_decoder.go
index 8b197e7b..89c47632 100644
--- a/parquet/internal/encoding/fixed_len_byte_array_decoder.go
+++ b/parquet/internal/encoding/fixed_len_byte_array_decoder.go
@@ -23,11 +23,15 @@ import (
"github.com/apache/arrow-go/v18/internal/utils"
"github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding/streaming"
)
// PlainFixedLenByteArrayDecoder is a plain encoding decoder for Fixed Length
Byte Arrays
type PlainFixedLenByteArrayDecoder struct {
decoder
+
+ // src drives the streaming path (PageStreamingEnabled) when non-nil.
+ src streaming.ValueBuffer
}
// Type returns the physical type this decoder operates on, FixedLength Byte
Arrays
@@ -35,7 +39,47 @@ func (PlainFixedLenByteArrayDecoder) Type() parquet.Type {
return parquet.Types.FixedLenByteArray
}
+func (pflba *PlainFixedLenByteArrayDecoder) SetSource(nvals int, src
streaming.ValueBuffer) {
+ pflba.src, pflba.nvals = src, nvals
+}
+
+func (pflba *PlainFixedLenByteArrayDecoder) SetData(nvals int, data []byte)
error {
+ pflba.src = nil
+ return pflba.decoder.SetData(nvals, data)
+}
+
+// decodeStreaming aliases each value in the buffer; see the byte-array
decoder for the
+// Recycle/alias lifetime.
+func (pflba *PlainFixedLenByteArrayDecoder) decodeStreaming(out
[]parquet.FixedLenByteArray) (int, error) {
+ pflba.src.Recycle()
+ max := utils.Min(len(out), pflba.nvals)
+ w := pflba.typeLen
+ for i := 0; i < max; i++ {
+ val, err := pflba.src.Fill(w)
+ if err != nil {
+ return i, errors.New("parquet: eof exception")
+ }
+ out[i] = val[:w:w]
+ pflba.src.Advance(w)
+ }
+ pflba.nvals -= max
+ return max, nil
+}
+
+func (pflba *PlainFixedLenByteArrayDecoder) discardStreaming(n int) (int,
error) {
+ n = min(n, pflba.nvals)
+ pflba.src.Recycle()
+ if err := pflba.src.Skip(n * pflba.typeLen); err != nil {
+ return 0, errors.New("parquet: eof exception")
+ }
+ pflba.nvals -= n
+ return n, nil
+}
+
func (pflba *PlainFixedLenByteArrayDecoder) Discard(n int) (int, error) {
+ if pflba.src != nil {
+ return pflba.discardStreaming(n)
+ }
n = min(n, pflba.nvals)
numBytesNeeded := n * pflba.typeLen
if numBytesNeeded > len(pflba.data) || numBytesNeeded > math.MaxInt32 {
@@ -51,6 +95,10 @@ func (pflba *PlainFixedLenByteArrayDecoder) Discard(n int)
(int, error) {
// values to decode or the length of out has been filled. Then returns the
total number of values
// that were decoded.
func (pflba *PlainFixedLenByteArrayDecoder) Decode(out
[]parquet.FixedLenByteArray) (int, error) {
+ if pflba.src != nil {
+ return pflba.decodeStreaming(out)
+ }
+
max := utils.Min(len(out), pflba.nvals)
numBytesNeeded := max * pflba.typeLen
if numBytesNeeded > len(pflba.data) || numBytesNeeded > math.MaxInt32 {
diff --git a/parquet/internal/encoding/streaming/streaming_test.go
b/parquet/internal/encoding/streaming/streaming_test.go
new file mode 100644
index 00000000..c2b7ae2b
--- /dev/null
+++ b/parquet/internal/encoding/streaming/streaming_test.go
@@ -0,0 +1,70 @@
+// 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 streaming_test
+
+import (
+ "bytes"
+ "io"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding/streaming"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestStreamBuffer(t *testing.T) {
+ data := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
+ vb := streaming.NewStreamBuffer(memory.DefaultAllocator,
bytes.NewReader(data), len(data), nil)
+
+ buf, err := vb.Fill(4)
+ require.NoError(t, err)
+ require.GreaterOrEqual(t, len(buf), 4)
+ assert.Equal(t, []byte{0, 1, 2, 3}, buf[:4])
+
+ vb.Advance(4)
+ buf, err = vb.Fill(4)
+ require.NoError(t, err)
+ require.GreaterOrEqual(t, len(buf), 4)
+ assert.Equal(t, []byte{4, 5, 6, 7}, buf[:4])
+
+ vb.Advance(4)
+ _, err = vb.Fill(4)
+ assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+}
+
+func TestClipBatch(t *testing.T) {
+ const buf = streaming.DefaultBufferSize
+ // Value region larger than the buffer: clip so rows*avgWidth stays
near the buffer.
+ assert.Equal(t, int64(buf/1024), streaming.ClipBatch(4*buf, 4096)) //
avg 1024 bytes
+ assert.Equal(t, int64(1), streaming.ClipBatch(4*buf, 2)) //
huge values: one per decode
+ // Whole region fits one buffer, or no values: no clip (returns nvals).
+ assert.Equal(t, int64(1_000_000), streaming.ClipBatch(buf/2,
1_000_000)) // small values
+ assert.Equal(t, int64(1_000_000), streaming.ClipBatch(0, 1_000_000))
// all-null page
+ assert.Equal(t, int64(0), streaming.ClipBatch(0, 0))
+}
+
+func TestStreamBufferGrowsForOversizedValue(t *testing.T) {
+ data := bytes.Repeat([]byte{0xAB}, (1<<20)+1024) // larger than the
default buffer
+ vb := streaming.NewStreamBuffer(memory.DefaultAllocator,
bytes.NewReader(data), len(data), nil)
+
+ buf, err := vb.Fill(len(data))
+ require.NoError(t, err)
+ require.GreaterOrEqual(t, len(buf), len(data))
+ assert.Equal(t, data, buf[:len(data)])
+ require.NoError(t, vb.Close())
+}
diff --git a/parquet/internal/encoding/streaming/value_buffer.go
b/parquet/internal/encoding/streaming/value_buffer.go
new file mode 100644
index 00000000..4bfc7875
--- /dev/null
+++ b/parquet/internal/encoding/streaming/value_buffer.go
@@ -0,0 +1,174 @@
+// 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 streaming
+
+import (
+ "io"
+
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+// ValueBuffer is the byte source a streaming decoder reads; bytes from Fill
stay valid
+// until the next Recycle, so a decoder can alias them instead of copying.
+type ValueBuffer interface {
+ // Fill returns >= need contiguous bytes from the cursor, valid until
the next Recycle;
+ // io.ErrUnexpectedEOF if fewer remain (incl. a need past the value
region).
+ Fill(need int) ([]byte, error)
+ // Advance marks the first n bytes from the cursor as consumed.
+ Advance(n int)
+ // Skip discards the next n bytes without buffering them, for the
discard/seek path.
+ Skip(n int) error
+ // Recycle reclaims everything handed out since the last Recycle; the
caller must be
+ // done with the previous batch's aliases. A decoder calls it when
entering Decode.
+ Recycle()
+ io.Closer
+}
+
+// Decoder is a decoder that can read from a ValueBuffer via SetSource.
+type Decoder interface {
+ SetSource(nvals int, src ValueBuffer)
+}
+
+// DefaultBufferSize is the steady-state chunk size and the clip target.
+const DefaultBufferSize = 1 << 20
+
+// streamBuffer serves the value stream from reusable chunks; aliases stay
valid until Recycle.
+type streamBuffer struct {
+ r io.Reader
+ onClose func() error
+ mem memory.Allocator
+ chunkSize int
+ cur []byte // chunk currently being filled (off == n on Fill
entry)
+ off, n int // consumed cursor / filled end within cur
+ live [][]byte // earlier chunks still backing this Decode's aliases
+ remaining int // value-region bytes not yet consumed; bounds Fill
vs corrupt lengths
+}
+
+// NewStreamBuffer returns a ValueBuffer over r, bounded to valueBytes (the
value region).
+func NewStreamBuffer(mem memory.Allocator, r io.Reader, valueBytes int,
onClose func() error) ValueBuffer {
+ size := DefaultBufferSize
+ if valueBytes > 0 && valueBytes < size {
+ size = valueBytes
+ }
+ return &streamBuffer{mem: mem, r: r, onClose: onClose, chunkSize: size,
cur: mem.Allocate(size), remaining: valueBytes}
+}
+
+// ClipBatch estimates rows per Decode that keep a batch's aliased values near
the buffer
+// size, for a page whose value region exceeds it. Row count is a conservative
proxy for
+// value count; skewed sizes may overshoot. Returns nvals (no clip) when the
whole region
+// already fits one buffer.
+func ClipBatch(valueBytes, nvals int) int64 {
+ if valueBytes <= DefaultBufferSize || nvals <= 0 {
+ return int64(nvals)
+ }
+ avg := valueBytes / nvals
+ if avg < 1 {
+ avg = 1
+ }
+ if clip := int64(DefaultBufferSize / avg); clip > 1 {
+ return clip
+ }
+ return 1
+}
+
+func (s *streamBuffer) Advance(n int) {
+ s.off += n
+ s.remaining -= n
+}
+
+func (s *streamBuffer) Skip(n int) error {
+ if n > s.remaining {
+ return io.ErrUnexpectedEOF
+ }
+ s.remaining -= n
+ buffered := s.n - s.off
+ if buffered >= n {
+ s.off += n
+ return nil
+ }
+ n -= buffered
+ // Nothing here is aliased, so drop the cursor and discard the rest
from the reader.
+ s.off, s.n = 0, 0
+ _, err := io.CopyN(io.Discard, s.r, int64(n))
+ return err
+}
+
+func (s *streamBuffer) Fill(need int) ([]byte, error) {
+ if s.n-s.off >= need {
+ return s.cur[s.off:s.n], nil
+ }
+ if need > s.remaining {
+ return nil, io.ErrUnexpectedEOF
+ }
+ if need > len(s.cur)-s.off {
+ s.rotate(need)
+ }
+ for s.n-s.off < need {
+ m, err := s.r.Read(s.cur[s.n : s.off+need])
+ s.n += m
+ if err != nil {
+ if err == io.EOF {
+ break
+ }
+ return nil, err
+ }
+ }
+ if s.n-s.off < need {
+ return nil, io.ErrUnexpectedEOF
+ }
+ return s.cur[s.off:s.n], nil
+}
+
+// rotate retires the full chunk (still backing aliases) and starts fresh,
carrying any unconsumed tail.
+func (s *streamBuffer) rotate(need int) {
+ tail := s.n - s.off
+ s.live = append(s.live, s.cur)
+ size := s.chunkSize
+ if need > size {
+ size = need
+ }
+ nc := s.mem.Allocate(size)
+ copy(nc, s.cur[s.off:s.n])
+ s.cur, s.off, s.n = nc, 0, tail
+}
+
+func (s *streamBuffer) Recycle() {
+ for _, b := range s.live {
+ s.mem.Free(b)
+ }
+ s.live = s.live[:0]
+ // Rewind the primary chunk; swap out an oversized one so steady state
stays small.
+ if len(s.cur) != s.chunkSize {
+ s.mem.Free(s.cur)
+ s.cur = s.mem.Allocate(s.chunkSize)
+ }
+ s.off, s.n = 0, 0
+}
+
+func (s *streamBuffer) Close() error {
+ for _, b := range s.live {
+ s.mem.Free(b)
+ }
+ if s.cur != nil {
+ s.mem.Free(s.cur)
+ }
+ s.live, s.cur = nil, nil
+ if s.onClose != nil {
+ return s.onClose()
+ }
+ return nil
+}
diff --git a/parquet/internal/encoding/streaming/value_buffer_internal_test.go
b/parquet/internal/encoding/streaming/value_buffer_internal_test.go
new file mode 100644
index 00000000..b10d54a0
--- /dev/null
+++ b/parquet/internal/encoding/streaming/value_buffer_internal_test.go
@@ -0,0 +1,136 @@
+// 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 streaming
+
+import (
+ "bytes"
+ "io"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+// newTestBuffer builds a streamBuffer with a small chunk so a few reads force
the
+// rotate/oversized paths.
+func newTestBuffer(mem memory.Allocator, data []byte, chunkSize int)
*streamBuffer {
+ return &streamBuffer{mem: mem, r: bytes.NewReader(data), chunkSize:
chunkSize, cur: mem.Allocate(chunkSize), remaining: len(data)}
+}
+
+// fillValue reads the next need bytes as a value: Fill then Advance,
returning the
+// aliased slice.
+func fillValue(t *testing.T, s *streamBuffer, need int) []byte {
+ t.Helper()
+ b, err := s.Fill(need)
+ require.NoError(t, err)
+ v := b[:need:need]
+ s.Advance(need)
+ return v
+}
+
+// TestStreamBufferAliasesStableUntilRecycle checks that values aliased across
a chunk
+// rotation stay valid until Recycle: the filled chunk is set aside, not
overwritten.
+func TestStreamBufferAliasesStableUntilRecycle(t *testing.T) {
+ data := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}
+ s := newTestBuffer(memory.DefaultAllocator, data, 8)
+
+ v1 := fillValue(t, s, 4) // [0..3]
+ v2 := fillValue(t, s, 4) // [4..7], fills the chunk
+ v3 := fillValue(t, s, 4) // [8..11], forces a fresh chunk
+
+ assert.Equal(t, []byte{0, 1, 2, 3}, v1, "aliased value corrupted by
rotation")
+ assert.Equal(t, []byte{4, 5, 6, 7}, v2, "aliased value corrupted by
rotation")
+ assert.Equal(t, []byte{8, 9, 10, 11}, v3)
+
+ require.NoError(t, s.Close())
+}
+
+// TestStreamBufferOversizedValue checks that a value larger than the chunk
allocates a
+// fresh chunk to hold it, and Recycle swaps back to the steady-state size.
+func TestStreamBufferOversizedValue(t *testing.T) {
+ big := bytes.Repeat([]byte{0xAB}, 20)
+ data := append([]byte{1, 2, 3, 4}, big...)
+ s := newTestBuffer(memory.DefaultAllocator, data, 8)
+
+ assert.Equal(t, []byte{1, 2, 3, 4}, fillValue(t, s, 4))
+ assert.Equal(t, big, fillValue(t, s, 20))
+ assert.Greater(t, len(s.cur), 8, "buffer should grow for the oversized
value")
+
+ s.Recycle()
+ assert.Equal(t, 8, len(s.cur), "buffer should shrink back to the
steady-state size")
+
+ require.NoError(t, s.Close())
+}
+
+// TestStreamBufferSkip checks Skip discards both buffered and not-yet-read
bytes.
+func TestStreamBufferSkip(t *testing.T) {
+ data := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
+ s := newTestBuffer(memory.DefaultAllocator, data, 8)
+
+ assert.Equal(t, []byte{0, 1}, fillValue(t, s, 2))
+ require.NoError(t, s.Skip(6)) // 2 buffered + 4 straight from the reader
+ assert.Equal(t, []byte{8, 9}, fillValue(t, s, 2))
+
+ require.NoError(t, s.Close())
+}
+
+// TestStreamBufferRejectsOversizedNeed checks a need past the value region (a
corrupt
+// length prefix) is rejected before allocating.
+func TestStreamBufferRejectsOversizedNeed(t *testing.T) {
+ checked := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ s := newTestBuffer(checked, make([]byte, 10), 8)
+
+ _, err := s.Fill(1 << 30)
+ require.ErrorIs(t, err, io.ErrUnexpectedEOF)
+
+ require.NoError(t, s.Close())
+ checked.AssertSize(t, 0) // the 1 GiB was never allocated
+}
+
+// TestStreamBufferRecycleReclaims checks that repeated Recycle keeps memory
bounded and
+// frees the extra chunks a rotation created, so nothing leaks across Decodes.
+func TestStreamBufferRecycleReclaims(t *testing.T) {
+ checked := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ s := newTestBuffer(checked, bytes.Repeat([]byte{1}, 1000), 8)
+
+ for i := 0; i < 50; i++ {
+ fillValue(t, s, 5) // off=5
+ fillValue(t, s, 5) // forces a rotation: old chunk goes to live
+ require.NotEmpty(t, s.live)
+ s.Recycle() // must free the retired chunk
+ require.Empty(t, s.live)
+ }
+
+ require.NoError(t, s.Close())
+ checked.AssertSize(t, 0)
+}
+
+func TestStreamBufferRotateCarriesTail(t *testing.T) {
+ data := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
+ s := newTestBuffer(memory.DefaultAllocator, data, 8)
+ defer s.Close()
+
+ b, err := s.Fill(6)
+ require.NoError(t, err)
+ assert.Equal(t, data[:6], b[:6])
+ s.Advance(2) // leave an unconsumed tail so the next Fill rotates with
off != n
+
+ b, err = s.Fill(7)
+ require.NoError(t, err)
+ assert.Equal(t, data[2:9], b[:7])
+}
diff --git a/parquet/internal/encoding/streaming_decoder_test.go
b/parquet/internal/encoding/streaming_decoder_test.go
new file mode 100644
index 00000000..84e4a7e0
--- /dev/null
+++ b/parquet/internal/encoding/streaming_decoder_test.go
@@ -0,0 +1,120 @@
+// 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 encoding_test
+
+import (
+ "bytes"
+ "encoding/binary"
+ "testing"
+ "testing/iotest"
+
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding/streaming"
+ "github.com/apache/arrow-go/v18/parquet/schema"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func plainByteArrayBytes(vals []parquet.ByteArray) []byte {
+ var raw bytes.Buffer
+ var l [4]byte
+ for _, v := range vals {
+ binary.LittleEndian.PutUint32(l[:], uint32(len(v)))
+ raw.Write(l[:])
+ raw.Write(v)
+ }
+ return raw.Bytes()
+}
+
+func TestPlainByteArrayDecoderSetSource(t *testing.T) {
+ vals := []parquet.ByteArray{[]byte("hello"), []byte(""),
[]byte("world!!"), bytes.Repeat([]byte("x"), 5000)}
+
+ raw := plainByteArrayBytes(vals)
+ dec := encoding.NewDecoder(parquet.Types.ByteArray,
parquet.Encodings.Plain, nil, memory.DefaultAllocator)
+ sd, ok := dec.(streaming.Decoder)
+ require.True(t, ok, "PlainByteArrayDecoder must implement
streaming.Decoder")
+ sd.SetSource(len(vals),
streaming.NewStreamBuffer(memory.DefaultAllocator,
iotest.OneByteReader(bytes.NewReader(raw)), len(raw), nil))
+
+ out := make([]parquet.ByteArray, len(vals))
+ n, err := dec.(encoding.ByteArrayDecoder).Decode(out)
+ require.NoError(t, err)
+ require.Equal(t, len(vals), n)
+ for i := range vals {
+ assert.Equal(t, []byte(vals[i]), []byte(out[i]))
+ }
+}
+
+func TestPlainFixedLenByteArrayDecoderSetSource(t *testing.T) {
+ const w = 3
+ descr := schema.NewColumn(schema.NewFixedLenByteArrayNode("f",
parquet.Repetitions.Required, w, -1), 0, 0)
+ vals := []parquet.FixedLenByteArray{[]byte("abc"), []byte("def"),
bytes.Repeat([]byte("z"), w)}
+ raw := make([]byte, 0, len(vals)*w)
+ for _, v := range vals {
+ raw = append(raw, v...)
+ }
+
+ dec := encoding.NewDecoder(parquet.Types.FixedLenByteArray,
parquet.Encodings.Plain, descr, memory.DefaultAllocator)
+ sd, ok := dec.(streaming.Decoder)
+ require.True(t, ok, "PlainFixedLenByteArrayDecoder must implement
streaming.Decoder")
+ sd.SetSource(len(vals),
streaming.NewStreamBuffer(memory.DefaultAllocator,
iotest.OneByteReader(bytes.NewReader(raw)), len(raw), nil))
+
+ out := make([]parquet.FixedLenByteArray, len(vals))
+ n, err := dec.(encoding.FixedLenByteArrayDecoder).Decode(out)
+ require.NoError(t, err)
+ require.Equal(t, len(vals), n)
+ for i := range vals {
+ assert.Equal(t, []byte(vals[i]), []byte(out[i]))
+ }
+}
+
+func TestFixedWidthDecoderNotStreaming(t *testing.T) {
+ dec := encoding.NewDecoder(parquet.Types.Int64,
parquet.Encodings.Plain, nil, memory.DefaultAllocator)
+ _, ok := dec.(streaming.Decoder)
+ assert.False(t, ok)
+}
+
+func TestStreamingByteArrayDecodeErrors(t *testing.T) {
+ decode := func(raw []byte) error {
+ dec := encoding.NewDecoder(parquet.Types.ByteArray,
parquet.Encodings.Plain, nil, memory.DefaultAllocator)
+ dec.(streaming.Decoder).SetSource(1,
streaming.NewStreamBuffer(memory.DefaultAllocator, bytes.NewReader(raw),
len(raw), nil))
+ _, err :=
dec.(encoding.ByteArrayDecoder).Decode(make([]parquet.ByteArray, 1))
+ return err
+ }
+
+ t.Run("truncated value", func(t *testing.T) {
+ var prefix [4]byte
+ binary.LittleEndian.PutUint32(prefix[:], 100) // claims 100
bytes
+ require.Error(t, decode(append(prefix[:],
bytes.Repeat([]byte{0xAB}, 30)...)))
+ })
+ t.Run("negative length", func(t *testing.T) {
+ var prefix [4]byte
+ binary.LittleEndian.PutUint32(prefix[:], 0xFFFFFFFF) // -1 as
int32
+ require.Error(t, decode(prefix[:]))
+ })
+}
+
+func TestStreamingFixedLenByteArrayDecodeTruncated(t *testing.T) {
+ const w = 4
+ descr := schema.NewColumn(schema.NewFixedLenByteArrayNode("f",
parquet.Repetitions.Required, w, -1), 0, 0)
+ raw := []byte{0x01, 0x02} // 2 of 4 bytes
+ dec := encoding.NewDecoder(parquet.Types.FixedLenByteArray,
parquet.Encodings.Plain, descr, memory.DefaultAllocator)
+ dec.(streaming.Decoder).SetSource(1,
streaming.NewStreamBuffer(memory.DefaultAllocator, bytes.NewReader(raw),
len(raw), nil))
+ _, err :=
dec.(encoding.FixedLenByteArrayDecoder).Decode(make([]parquet.FixedLenByteArray,
1))
+ require.Error(t, err)
+}
diff --git a/parquet/reader_properties.go b/parquet/reader_properties.go
index 5b119dcf..36c7529e 100644
--- a/parquet/reader_properties.go
+++ b/parquet/reader_properties.go
@@ -44,6 +44,22 @@ type ReaderProperties struct {
// this to true can optimize memory usage for the reader. Additionally,
this can decrease
// the amount of data retrieved when only needs to access small
portions of the parquet file.
BufferedStreamEnabled bool
+ // PageStreamingEnabled, when true, decodes eligible data pages
incrementally instead
+ // of decoding the whole uncompressed page at once. A decode aliases
the batch's
+ // values directly in a reusable stream buffer (min(1 MiB, page size))
rather than
+ // materializing the page, so the read batch is clipped toward the
average value
+ // width to keep one decode's values near the buffer size. Peak memory
is then
+ // best-effort roughly that buffer plus the largest single value, the
caller's batch
+ // and the level buffers; skewed value sizes may overshoot the target.
Eligible pages
+ // are PLAIN-encoded V1/V2 data pages larger than 1 MiB (smaller pages
fit in a single
+ // buffer, so streaming only adds overhead), of a streaming-capable,
unencrypted codec
+ // (UNCOMPRESSED/GZIP/BROTLI/ZSTD) for a supported physical type; every
other page is
+ // read whole as before. Default false (no behavior change to decoded
values).
+ //
+ // Note: a streaming page's raw Page.Data() (via GetColumnPageReader)
returns just its
+ // level region; the decoded values read through
RowGroup.Column()/ReadBatch are the
+ // same as without streaming.
+ PageStreamingEnabled bool
}
type BufferedReader interface {
@@ -62,7 +78,7 @@ func NewReaderProperties(alloc memory.Allocator)
*ReaderProperties {
if alloc == nil {
alloc = memory.DefaultAllocator
}
- return &ReaderProperties{alloc, DefaultBufSize, nil, false}
+ return &ReaderProperties{alloc: alloc, BufferSize: DefaultBufSize}
}
// Allocator returns the allocator that the properties were initialized with
@@ -74,7 +90,7 @@ func (r *ReaderProperties) Allocator() memory.Allocator {
return r.alloc }
// 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
+ return utils.NewBufferedReader(io.NewSectionReader(source,
start, nbytes), int(min(r.BufferSize, nbytes))), nil
}
data := make([]byte, nbytes)
