Correctly unencrypt V2 data pages

Author: daniel-adam-tfs

parquet/file/file_reader_test.go

@@ -515,6 +515,36 @@ func (p *serializedPageReader) decompress(rd io.Reader, lenCompressed int, buf [
 	return p.codec.Decode(buf, data), nil
 }
 
+func (p *serializedPageReader) decompressV2(rd io.Reader, lenCompressed int, levelsBytelen int, compressed bool, buf []byte) ([]byte, error) {
+	// if encrypted, we need to decrypt before decompressing
+	p.decompressBuffer.ResizeNoShrink(lenCompressed)
+	b := bytes.NewBuffer(p.decompressBuffer.Bytes()[:0])
+	if _, err := io.CopyN(b, rd, int64(lenCompressed)); err != nil {
+		return nil, err
+	}
+	data := p.decompressBuffer.Bytes()
+	if p.cryptoCtx.DataDecryptor != nil {
+		data = p.cryptoCtx.DataDecryptor.Decrypt(p.decompressBuffer.Bytes())
+		// encrypted + uncompressed -> just copy the decrypted data to output buffer
+		if !compressed {
+			copy(buf, data)
+			return buf, nil
+		}
+	} else {
+		// unencryped + uncompressed -> nothing to be done, data is already read into the buffer
+		if !compressed {
+			return buf, nil
+		}
+	}
+
+	// definition + repetition levels are always uncompressed
+	if levelsBytelen > 0 {
+		copy(buf, data[:levelsBytelen])
+		data = data[levelsBytelen:]
+	}
+	return p.codec.Decode(buf[levelsBytelen:], data), nil
+}
+
 type dataheader interface {
 	IsSetStatistics() bool
 	GetStatistics() *format.Statistics
@@ -628,7 +658,6 @@ func (p *serializedPageReader) readPageHeader(rd parquet.BufferedReader, hdr *fo
 			}
 			continue
 		}
-
 		rd.Discard(len(view) - int(remaining) + extra)
 		break
 	}
@@ -812,15 +841,9 @@ func (p *serializedPageReader) Next() bool {
 				return false
 			}
 
-			if compressed {
-				if levelsBytelen > 0 {
-					io.ReadFull(p.r, buf.Bytes()[:levelsBytelen])
-				}
-				if _, p.err = p.decompress(p.r, lenCompressed-levelsBytelen, buf.Bytes()[levelsBytelen:]); p.err != nil {
-					return false
-				}
-			} else {
-				io.ReadFull(p.r, buf.Bytes())
+			if _, err := p.decompressV2(p.r, lenCompressed, levelsBytelen, compressed, buf.Bytes()); err != nil {
+				p.err = err
+				return false
 			}
 
 			if buf.Len() != lenUncompressed {