Read uncompressed data directly into the page buffer

Author: daniel-adam-tfs

parquet/file/page_reader.go

@@ -374,6 +374,8 @@ type serializedPageReader struct {
 	dataPageBuffer   *memory.Buffer
 	dictPageBuffer   *memory.Buffer
 	err              error
+
+	isCompressed bool
 }
 
 func (p *serializedPageReader) Close() error {
@@ -402,6 +404,7 @@ func (p *serializedPageReader) init(compressType compress.Compression, ctx *Cryp
 		return err
 	}
 	p.codec = codec
+	p.isCompressed = compressType != compress.Codecs.Uncompressed
 
 	if ctx != nil {
 		p.cryptoCtx = *ctx
@@ -444,6 +447,7 @@ func NewPageReader(r parquet.BufferedReader, nrows int64, compressType compress.
 		dictPageBuffer:   memory.NewResizableBuffer(mem),
 	}
 	rdr.decompressBuffer.ResizeNoShrink(defaultPageHeaderSize)
+	rdr.isCompressed = compressType != compress.Codecs.Uncompressed
 	if ctx != nil {
 		rdr.cryptoCtx = *ctx
 		rdr.initDecryption()
@@ -460,6 +464,8 @@ func (p *serializedPageReader) Reset(r parquet.BufferedReader, nrows int64, comp
 	if p.err != nil {
 		return
 	}
+	p.isCompressed = compressType != compress.Codecs.Uncompressed
+
 	if ctx != nil {
 		p.cryptoCtx = *ctx
 		p.initDecryption()
@@ -502,6 +508,20 @@ func (p *serializedPageReader) Page() Page {
 	return p.curPage
 }
 
+func (p *serializedPageReader) readUncompressed(rd io.Reader, lenUncompressed int, buf []byte) ([]byte, error) {
+	n, err := io.ReadFull(rd, buf[:lenUncompressed])
+	if err != nil {
+		return nil, err
+	}
+	if n != lenUncompressed {
+		return nil, fmt.Errorf("parquet: expected to read %d bytes but only read %d", lenUncompressed, n)
+	}
+	if p.cryptoCtx.DataDecryptor != nil {
+		buf = p.cryptoCtx.DataDecryptor.Decrypt(buf)
+	}
+	return buf[:lenUncompressed], nil
+}
+
 func (p *serializedPageReader) decompress(rd io.Reader, lenCompressed int, buf []byte) ([]byte, error) {
 	p.decompressBuffer.ResizeNoShrink(lenCompressed)
 	data := p.decompressBuffer.Bytes()
@@ -735,7 +755,13 @@ func (p *serializedPageReader) Next() bool {
 			p.dictPageBuffer.ResizeNoShrink(lenUncompressed)
 			buf := memory.NewBufferBytes(p.dictPageBuffer.Bytes())
 
-			data, err := p.decompress(p.r, lenCompressed, buf.Bytes())
+			var data []byte
+			var err error
+			if p.isCompressed {
+				data, err = p.decompress(p.r, lenCompressed, buf.Bytes())
+			} else {
+				data, err = p.readUncompressed(p.r, lenCompressed, buf.Bytes())
+			}
 			if err != nil {
 				p.err = err
 				return false
@@ -770,7 +796,13 @@ func (p *serializedPageReader) Next() bool {
 			firstRowIdx := p.rowsSeen
 			p.rowsSeen += int64(dataHeader.GetNumValues())
 
-			data, err := p.decompress(p.r, lenCompressed, buf.Bytes())
+			var data []byte
+			var err error
+			if p.isCompressed {
+				data, err = p.decompress(p.r, lenCompressed, buf.Bytes())
+			} else {
+				data, err = p.readUncompressed(p.r, lenCompressed, buf.Bytes())
+			}
 			if err != nil {
 				p.err = err
 				return false