Skip to content

Commit 8bd5514

Browse files
Matthew Topolzeroshade
Matthew Topol
and
zeroshade
committed
ARROW-16187: [Go][Parquet] Properly utilize BufferedStream and buffer size when reading
Currently the BufferSize in the `ReaderProperties` isn't utilized properly when enabling BufferedStreams. This fixes that issue so that enabling `BufferedStream` reading via the properties will correctly utilize the given buffer size when reading. The default buffersize is currently 16K, so reads that are larger than that will ignore the buffering and just pull directly from the underlying reader when BufferedStream is enabled, pulling the entire page or otherwise from the reader if BufferedStream is not enabled. The buffer size can be set larger so that controlled reads can improve performance on high-latency readers without having to use the memory to read the entire column/page/row group into memory. Closes apache#12876 from zeroshade/arrow-16187-parquet-buffered-stream Lead-authored-by: Matthew Topol <mtopol@factset.com> Co-authored-by: Matthew Topol <zotthewizard@gmail.com> Signed-off-by: Matthew Topol <mtopol@factset.com>
1 parent cbf7d76 commit 8bd5514

8 files changed

Lines changed: 276 additions & 27 deletions

File tree

‎go/internal/utils/buf_reader.go‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing, software
12+
// distributed under the License is distributed on an "AS IS" BASIS,
13+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
// See the License for the specific language governing permissions and
15+
// limitations under the License.
16+
17+
package utils
18+
19+
import (
20+
"errors"
21+
"io"
22+
)
23+
24+
// bufferedReader is similar to bufio.Reader except
25+
// it will expand the buffer if necessary when asked to Peek
26+
// more bytes than are in the buffer
27+
type bufferedReader struct {
28+
bufferSz int
29+
buf []byte
30+
r, w int
31+
rd io.Reader
32+
err error
33+
}
34+
35+
func NewBufferedReader(rd io.Reader, sz int) *bufferedReader {
36+
b, ok := rd.(*bufferedReader)
37+
if ok && len(b.buf) >= sz {
38+
return b
39+
}
40+
41+
r := &bufferedReader{
42+
rd: rd,
43+
}
44+
r.resizeBuffer(sz)
45+
return r
46+
}
47+
48+
func (b *bufferedReader) resetBuffer() {
49+
if b.buf == nil {
50+
b.buf = make([]byte, b.bufferSz)
51+
} else if b.bufferSz > cap(b.buf) {
52+
buf := b.buf
53+
b.buf = make([]byte, b.bufferSz)
54+
copy(b.buf, buf)
55+
} else {
56+
b.buf = b.buf[:b.bufferSz]
57+
}
58+
}
59+
60+
func (b *bufferedReader) resizeBuffer(newSize int) {
61+
b.bufferSz = newSize
62+
b.resetBuffer()
63+
}
64+
65+
func (b *bufferedReader) fill() {
66+
// slide existing data to the beginning
67+
if b.r > 0 {
68+
copy(b.buf, b.buf[b.r:b.w])
69+
b.w -= b.r
70+
b.r = 0
71+
}
72+
73+
if b.w >= len(b.buf) {
74+
panic("parquet/bufio: tried to fill full buffer")
75+
}
76+
77+
n, err := io.ReadAtLeast(b.rd, b.buf[b.w:], 1)
78+
if n < 0 {
79+
panic("negative read")
80+
}
81+
82+
b.w += n
83+
b.err = err
84+
}
85+
86+
func (b *bufferedReader) readErr() error {
87+
err := b.err
88+
b.err = nil
89+
return err
90+
}
91+
92+
func (b *bufferedReader) Buffered() int { return b.w - b.r }
93+
94+
func (b *bufferedReader) SetBufferSize(newSize int) error {
95+
if newSize <= 0 {
96+
return errors.New("buffer size should be positive")
97+
}
98+
99+
if b.w >= newSize {
100+
return errors.New("cannot shrink read buffer if buffered data remains")
101+
}
102+
103+
b.resizeBuffer(newSize)
104+
return nil
105+
}
106+
107+
func (b *bufferedReader) Peek(n int) ([]byte, error) {
108+
if n < 0 {
109+
return nil, errors.New("parquet/bufio: negative count")
110+
}
111+
112+
if n > len(b.buf) {
113+
if err := b.SetBufferSize(n); err != nil {
114+
return nil, err
115+
}
116+
}
117+
118+
for b.w-b.r < n && b.w-b.r < len(b.buf) && b.err == nil {
119+
b.fill() // b.w-b.r < len(b.buf) => buffer is not full
120+
}
121+
122+
return b.buf[b.r : b.r+n], b.readErr()
123+
}
124+
125+
func (b *bufferedReader) Discard(n int) (discarded int, err error) {
126+
if n < 0 {
127+
return 0, errors.New("negative count")
128+
}
129+
130+
if n == 0 {
131+
return
132+
}
133+
134+
remain := n
135+
for {
136+
skip := b.Buffered()
137+
if skip == 0 {
138+
b.fill()
139+
skip = b.Buffered()
140+
}
141+
if skip > remain {
142+
skip = remain
143+
}
144+
b.r += skip
145+
remain -= skip
146+
if remain == 0 {
147+
return n, nil
148+
}
149+
if b.err != nil {
150+
return n - remain, b.readErr()
151+
}
152+
}
153+
}
154+
155+
func (b *bufferedReader) Read(p []byte) (n int, err error) {
156+
n = len(p)
157+
if n == 0 {
158+
if b.Buffered() > 0 {
159+
return 0, nil
160+
}
161+
return 0, b.readErr()
162+
}
163+
164+
if b.r == b.w {
165+
if b.err != nil {
166+
return 0, b.readErr()
167+
}
168+
if len(p) >= len(b.buf) {
169+
// large read, empty buffer
170+
// read directly into p to avoid extra copy
171+
n, b.err = b.rd.Read(p)
172+
if n < 0 {
173+
panic("negative read")
174+
}
175+
return n, b.readErr()
176+
}
177+
178+
// one read
179+
// don't use b.fill
180+
b.r, b.w = 0, 0
181+
n, b.err = b.rd.Read(b.buf)
182+
if n < 0 {
183+
panic("negative read")
184+
}
185+
if n == 0 {
186+
return 0, b.readErr()
187+
}
188+
b.w += n
189+
}
190+
191+
// copy as much as we can
192+
n = copy(p, b.buf[b.r:b.w])
193+
b.r += n
194+
return n, nil
195+
}

‎go/parquet/file/column_writer_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424

2525
"github.com/apache/arrow/go/v8/arrow/bitutil"
2626
"github.com/apache/arrow/go/v8/arrow/memory"
27+
arrutils "github.com/apache/arrow/go/v8/internal/utils"
2728
"github.com/apache/arrow/go/v8/parquet"
2829
"github.com/apache/arrow/go/v8/parquet/compress"
2930
"github.com/apache/arrow/go/v8/parquet/file"
@@ -233,7 +234,7 @@ func (p *PrimitiveWriterTestSuite) SetupTest() {
233234

234235
func (p *PrimitiveWriterTestSuite) buildReader(nrows int64, compression compress.Compression) file.ColumnChunkReader {
235236
p.readbuffer = p.sink.Finish()
236-
pagereader, _ := file.NewPageReader(bytes.NewReader(p.readbuffer.Bytes()), nrows, compression, mem, nil)
237+
pagereader, _ := file.NewPageReader(arrutils.NewBufferedReader(bytes.NewReader(p.readbuffer.Bytes()), p.readbuffer.Len()), nrows, compression, mem, nil)
237238
return file.NewColumnReader(p.descr, pagereader, mem)
238239
}
239240

‎go/parquet/file/file_reader_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"testing"
2525

2626
"github.com/apache/arrow/go/v8/arrow/memory"
27+
"github.com/apache/arrow/go/v8/internal/utils"
2728
"github.com/apache/arrow/go/v8/parquet/compress"
2829
"github.com/apache/arrow/go/v8/parquet/file"
2930
"github.com/apache/arrow/go/v8/parquet/internal/encoding"
@@ -101,7 +102,7 @@ func (p *PageSerdeSuite) SetupTest() {
101102
func (p *PageSerdeSuite) InitSerializedPageReader(nrows int64, codec compress.Compression) {
102103
p.EndStream()
103104

104-
p.pageReader, _ = file.NewPageReader(bytes.NewReader(p.buffer.Bytes()), nrows, codec, memory.DefaultAllocator, nil)
105+
p.pageReader, _ = file.NewPageReader(utils.NewBufferedReader(bytes.NewReader(p.buffer.Bytes()), p.buffer.Len()), nrows, codec, memory.DefaultAllocator, nil)
105106
}
106107

107108
func (p *PageSerdeSuite) WriteDataPageHeader(maxSerialized int, uncompressed, compressed int32) {

‎go/parquet/file/page_reader.go‎

Lines changed: 9 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ type PageReader interface {
4646
// nil if there was no error and you just hit the end of the page
4747
Err() error
4848
// Reset allows reusing a page reader
49-
Reset(r io.ReadSeeker, nrows int64, compressType compress.Compression, ctx *CryptoContext)
49+
Reset(r parquet.BufferedReader, nrows int64, compressType compress.Compression, ctx *CryptoContext)
5050
}
5151

5252
// Page is an interface for handling DataPages or Dictionary Pages
@@ -289,7 +289,7 @@ func (d *DictionaryPage) Release() {
289289
func (d *DictionaryPage) IsSorted() bool { return d.sorted }
290290

291291
type serializedPageReader struct {
292-
r io.ReadSeeker
292+
r parquet.BufferedReader
293293
nrows int64
294294
rowsSeen int64
295295
mem memory.Allocator
@@ -310,7 +310,7 @@ type serializedPageReader struct {
310310
}
311311

312312
// NewPageReader returns a page reader for the data which can be read from the provided reader and compression.
313-
func NewPageReader(r io.ReadSeeker, nrows int64, compressType compress.Compression, mem memory.Allocator, ctx *CryptoContext) (PageReader, error) {
313+
func NewPageReader(r parquet.BufferedReader, nrows int64, compressType compress.Compression, mem memory.Allocator, ctx *CryptoContext) (PageReader, error) {
314314
if mem == nil {
315315
mem = memory.NewGoAllocator()
316316
}
@@ -336,10 +336,10 @@ func NewPageReader(r io.ReadSeeker, nrows int64, compressType compress.Compressi
336336
return rdr, nil
337337
}
338338

339-
func (p *serializedPageReader) Reset(r io.ReadSeeker, nrows int64, compressType compress.Compression, ctx *CryptoContext) {
340-
p.rowsSeen, p.pageOrd = 0, 0
339+
func (p *serializedPageReader) Reset(r parquet.BufferedReader, nrows int64, compressType compress.Compression, ctx *CryptoContext) {
340+
p.rowsSeen, p.pageOrd, p.nrows = 0, 0, nrows
341341
p.curPageHdr, p.curPage, p.err = nil, nil, nil
342-
p.r, p.nrows = r, nrows
342+
p.r = r
343343

344344
p.codec, p.err = compress.GetCodec(compressType)
345345
if p.err != nil {
@@ -390,7 +390,6 @@ func (p *serializedPageReader) Page() Page {
390390
}
391391

392392
func (p *serializedPageReader) decompress(lenCompressed int, buf []byte) ([]byte, error) {
393-
p.decompressBuffer.Reset()
394393
p.decompressBuffer.Grow(lenCompressed)
395394
if _, err := io.CopyN(&p.decompressBuffer, p.r, int64(lenCompressed)); err != nil {
396395
return nil, err
@@ -444,28 +443,19 @@ func (p *serializedPageReader) Next() bool {
444443
p.err = nil
445444

446445
for p.rowsSeen < p.nrows {
447-
// headerSize := 0
448446
allowedPgSz := defaultPageHeaderSize
449-
450-
start, _ := p.r.Seek(0, io.SeekCurrent)
451447
p.decompressBuffer.Reset()
452-
// Page headers can be very large because of page statistics
453-
// We try to deserialize a larger buffer progressively
454-
// until a maximum allowed header limit
455448
for {
456-
n, err := io.CopyN(&p.decompressBuffer, p.r, int64(allowedPgSz))
457-
// view, err := p.r.Peek(allowedPgSz)
449+
view, err := p.r.Peek(allowedPgSz)
458450
if err != nil && err != io.EOF {
459451
p.err = err
460452
return false
461453
}
462454

463-
if n == 0 {
455+
if len(view) == 0 {
464456
return false
465457
}
466458

467-
view := p.decompressBuffer.Bytes()
468-
469459
extra := 0
470460
if p.cryptoCtx.MetaDecryptor != nil {
471461
p.updateDecryption(p.cryptoCtx.MetaDecryptor, encryption.DictPageHeaderModule, p.dataPageHeaderAad)
@@ -483,7 +473,7 @@ func (p *serializedPageReader) Next() bool {
483473
continue
484474
}
485475

486-
p.r.Seek(start+int64(len(view)-int(remaining)+extra), io.SeekStart)
476+
p.r.Discard(len(view) - int(remaining) + extra)
487477
break
488478
}
489479

@@ -519,7 +509,6 @@ func (p *serializedPageReader) Next() bool {
519509
return false
520510
}
521511

522-
// p.buf.Resize(lenUncompressed)
523512
// make dictionary page
524513
p.curPage = &DictionaryPage{
525514
page: page{

‎go/parquet/internal/encryption/aes.go‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,52 @@ func newAesDecryptor(alg parquet.Cipher, metadata bool) *aesDecryptor {
191191
// the length of the plaintext after decryption.
192192
func (a *aesDecryptor) CiphertextSizeDelta() int { return a.ciphertextSizeDelta }
193193

194+
// DecryptFrom
195+
func (a *aesDecryptor) DecryptFrom(r io.Reader, key, aad []byte) []byte {
196+
block, err := aes.NewCipher(key)
197+
if err != nil {
198+
panic(err)
199+
}
200+
201+
var writtenCiphertextLen uint32
202+
if err := binary.Read(r, binary.LittleEndian, &writtenCiphertextLen); err != nil {
203+
panic(err)
204+
}
205+
206+
cipherText := make([]byte, writtenCiphertextLen)
207+
if n, err := io.ReadFull(r, cipherText); n != int(writtenCiphertextLen) || err != nil {
208+
panic(err)
209+
}
210+
211+
nonce := cipherText[:NonceLength]
212+
cipherText = cipherText[NonceLength:]
213+
if a.mode == gcmMode {
214+
aead, err := cipher.NewGCM(block)
215+
if err != nil {
216+
panic(err)
217+
}
218+
219+
plain, err := aead.Open(cipherText[:0], nonce, cipherText, aad)
220+
if err != nil {
221+
panic(err)
222+
}
223+
return plain
224+
}
225+
226+
// Parquet CTR IVs are comprised of a 12-byte nonce and a 4-byte initial
227+
// counter field.
228+
// The first 31 bits of the initial counter field are set to 0, the last bit
229+
// is set to 1.
230+
iv := make([]byte, ctrIVLen)
231+
copy(iv, nonce)
232+
iv[ctrIVLen-1] = 1
233+
234+
stream := cipher.NewCTR(block, iv)
235+
// dst := make([]byte, len(cipherText))
236+
stream.XORKeyStream(cipherText, cipherText)
237+
return cipherText
238+
}
239+
194240
// Decrypt returns the plaintext version of the given ciphertext when decrypted
195241
// with the provided key and AAD security bytes.
196242
func (a *aesDecryptor) Decrypt(cipherText, key, aad []byte) []byte {

0 commit comments

Comments
 (0)