zeroshade commented on code in PR #880:
URL: https://github.com/apache/arrow-go/pull/880#discussion_r3554173406
##########
parquet/file/row_group_reader.go:
##########
@@ -186,3 +195,21 @@ 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:
Review Comment:
Non-blocking nit: Snappy actually implements `StreamingCodec`
(compress/snappy.go) yet is intentionally excluded from this allowlist
(parquet's Snappy framing is block-based, not a spec-compliant stream). A
one-line comment here explaining *why* Snappy/LZ4_RAW are excluded despite
being registered streaming codecs would save a future maintainer from
"helpfully" adding them.
##########
parquet/internal/encoding/streaming/value_buffer.go:
##########
@@ -0,0 +1,171 @@
+// 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; no
tail since off == n.
+func (s *streamBuffer) rotate(need int) {
Review Comment:
Non-blocking, but worth hardening before more decoders adopt this: `rotate`
assumes `off == n` (per the comment) and so drops the buffered tail
`cur[off:n]` when it starts a fresh chunk. That invariant holds for every
current caller because they all do strictly matched
`Fill(need)`/`Advance(need)`. But `Fill`'s first branch can return *more* than
`need`, so a future "peek then partial-advance" decoder would silently drop
un-consumed bytes here — I reproduced both corruption and a spurious
`unexpected EOF` with a `Fill(6)`/`Advance(2)`/`Fill(7)` probe against this
buffer.
Suggest either carrying the tail into the new chunk on rotate, or making the
invariant fail loud, e.g. `if s.off != s.n { panic("streamBuffer.rotate:
unconsumed tail") }`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
