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 1476769e fix(parquet): reject truncated bit-packed batches (#943)
1476769e is described below
commit 1476769ecda35b4527f2d8f1adc276257a9f19f4
Author: Minh Vu <[email protected]>
AuthorDate: Thu Jul 16 20:20:20 2026 +0200
fix(parquet): reject truncated bit-packed batches (#943)
---
parquet/file/column_reader.go | 49 ++++---
parquet/file/column_writer_v2_row_boundary_test.go | 6 +-
parquet/file/record_reader.go | 17 ++-
parquet/internal/encoding/boolean_decoder.go | 19 ++-
parquet/internal/encoding/decoder.go | 4 +-
parquet/internal/encoding/encoding_test.go | 20 +++
parquet/internal/encoding/levels.go | 20 ++-
parquet/internal/encoding/levels_test.go | 32 ++++-
parquet/internal/utils/bit_packing_avx2_amd64.go | 15 ++-
parquet/internal/utils/bit_packing_default.go | 16 ++-
parquet/internal/utils/bit_packing_neon_arm64.go | 15 ++-
parquet/internal/utils/bit_reader.go | 16 ++-
parquet/internal/utils/bit_reader_test.go | 149 ++++++++++++++++++++-
parquet/internal/utils/rle.go | 29 ++--
parquet/internal/utils/typed_rle_dict.go | 10 +-
15 files changed, 342 insertions(+), 75 deletions(-)
diff --git a/parquet/file/column_reader.go b/parquet/file/column_reader.go
index 53cc284a..cb4c5930 100644
--- a/parquet/file/column_reader.go
+++ b/parquet/file/column_reader.go
@@ -126,12 +126,14 @@ type ColumnChunkReader interface {
// when this returns 0, you're at the end of a page.
numAvailValues() int64
// read the definition levels and return the number of definitions,
- // and the number of values to be read (number of def levels == maxdef
level)
+ // the number of values to be read (number of def levels == maxdef
level),
+ // and any decoding error.
// it also populates the passed in slice which should be sized
appropriately.
- readDefinitionLevels(levels []int16) (int, int64)
+ readDefinitionLevels(levels []int16) (int, int64, error)
// read the repetition levels and return the number of repetition
levels read
+ // and any decoding error.
// also populates the passed in slice, which should be sized
appropriately.
- readRepetitionLevels(levels []int16) int
+ readRepetitionLevels(levels []int16) (int, error)
// a column is made up of potentially multiple pages across potentially
multiple
// row groups. A PageReader allows looping through the pages in a
single row group.
// When moving to another row group for reading, use setPageReader to
re-use the
@@ -531,10 +533,10 @@ func (c *columnChunkReader) setDecoderData(page DataPage,
lvlByteLen int64) erro
// (the number of levels that are equal to the max definition level).
//
// If the max definition level is 0, the assumption is that there no nulls in
the
-// column and therefore no definition levels to read, so it will always return
0, 0
-func (c *columnChunkReader) readDefinitionLevels(levels []int16) (totalDecoded
int, valuesToRead int64) {
+// column and therefore no definition levels to read, so it will always return
0, 0, nil.
+func (c *columnChunkReader) readDefinitionLevels(levels []int16) (totalDecoded
int, valuesToRead int64, err error) {
if c.descr.MaxDefinitionLevel() == 0 {
- return 0, 0
+ return 0, 0, nil
}
return c.definitionDecoder.Decode(levels)
@@ -545,17 +547,17 @@ func (c *columnChunkReader) readDefinitionLevels(levels
[]int16) (totalDecoded i
// slice).
//
// If max repetition level is 0, it is assumed there are no repetition levels,
-// and thus will always return 0.
-func (c *columnChunkReader) readRepetitionLevels(levels []int16) int {
+// and thus will always return 0, nil.
+func (c *columnChunkReader) readRepetitionLevels(levels []int16) (int, error) {
if c.descr.MaxRepetitionLevel() == 0 {
- return 0
+ return 0, nil
}
if len(c.repLvlBuffer) > 0 {
- return copy(levels, c.repLvlBuffer[c.numDecoded:])
+ return copy(levels, c.repLvlBuffer[c.numDecoded:]), nil
}
- nlevels, _ := c.repetitionDecoder.Decode(levels)
- return nlevels
+ nlevels, _, err := c.repetitionDecoder.Decode(levels)
+ return nlevels, err
}
// determineNumToRead reads the definition levels (and optionally populates
the repetition levels)
@@ -580,13 +582,20 @@ func (c *columnChunkReader) determineNumToRead(batchLen
int64, defLvls, repLvls
if defLvls == nil {
defLvls = c.getDefLvlBuffer(size)
}
- ndefs, toRead = c.readDefinitionLevels(defLvls[:size])
+ ndefs, toRead, err = c.readDefinitionLevels(defLvls[:size])
+ if err != nil {
+ return
+ }
} else {
toRead = size
}
if c.descr.MaxRepetitionLevel() > 0 && repLvls != nil {
- nreps := c.readRepetitionLevels(repLvls[:size])
+ nreps, repErr := c.readRepetitionLevels(repLvls[:size])
+ if repErr != nil {
+ err = repErr
+ return
+ }
if defLvls != nil && ndefs != nreps {
err = errors.New("parquet: number of decoded rep/def
levels did not match")
}
@@ -663,7 +672,11 @@ func (c *columnChunkReader) skipRows(nrows int64) error {
// the repetition levels for the entire page at once
and then go
// through them to find the right row.
repLvls := c.getRepLvlBuffer(c.numBuffered)
- nreps, _ := c.repetitionDecoder.Decode(repLvls)
+ nreps, _, err := c.repetitionDecoder.Decode(repLvls)
+ if err != nil {
+ c.err = err
+ return err
+ }
rowsSkipped := int64(0)
levelsToSkip := -1
@@ -686,7 +699,11 @@ func (c *columnChunkReader) skipRows(nrows int64) error {
var valuesToSkip int64
if c.descr.MaxDefinitionLevel() > 0 {
defLvls :=
c.getDefLvlBuffer(int64(levelsToSkip))
- _, valuesToSkip =
c.readDefinitionLevels(defLvls)
+ _, valuesToSkip, err =
c.readDefinitionLevels(defLvls)
+ if err != nil {
+ c.err = err
+ return err
+ }
} else {
valuesToSkip = int64(levelsToSkip)
}
diff --git a/parquet/file/column_writer_v2_row_boundary_test.go
b/parquet/file/column_writer_v2_row_boundary_test.go
index aa8e4b39..1492a937 100644
--- a/parquet/file/column_writer_v2_row_boundary_test.go
+++ b/parquet/file/column_writer_v2_row_boundary_test.go
@@ -133,7 +133,8 @@ func assertPagesStartOnRowBoundary(t *testing.T, raw
[]byte) {
require.NoError(t, dec.SetDataV2(v2.RepetitionLevelByteLen(),
maxRep, int(v2.NumValues()), v2.Data()))
levels := make([]int16, v2.NumValues())
- n, _ := dec.Decode(levels)
+ n, _, err := dec.Decode(levels)
+ require.NoError(t, err)
require.Greater(t, n, 0, "page %d decoded zero repetition
levels", v2Count)
require.Zerof(t, levels[0],
"DataPageV2 #%d starts mid-row (first repetition level
= %d, want 0)", v2Count, levels[0])
@@ -225,7 +226,8 @@ func TestWriteSpacedV2NullableListWideRow(t *testing.T) {
var dec encoding.LevelDecoder
require.NoError(t, dec.SetDataV2(v2.RepetitionLevelByteLen(),
maxRep, int(v2.NumValues()), v2.Data()))
levels := make([]int16, v2.NumValues())
- n, _ := dec.Decode(levels)
+ n, _, err := dec.Decode(levels)
+ require.NoError(t, err)
require.Greater(t, n, 0)
require.Zero(t, levels[0], "wide row page starts mid-row")
}
diff --git a/parquet/file/record_reader.go b/parquet/file/record_reader.go
index 3b98cdf0..87ae20f3 100644
--- a/parquet/file/record_reader.go
+++ b/parquet/file/record_reader.go
@@ -664,15 +664,26 @@ func (rr *recordReader) ReadRecords(numRecords int64)
(int64, error) {
defLevels := rr.DefLevels()[int(rr.levelsWritten):]
levelsRead := 0
+ var err error
// not present for non-repeated fields
if rr.Descriptor().MaxRepetitionLevel() > 0 {
repLevels :=
rr.RepLevels()[int(rr.levelsWritten):]
- levelsRead, _ =
rr.readDefinitionLevels(defLevels[:batchSize])
- if
rr.readRepetitionLevels(repLevels[:batchSize]) != levelsRead {
+ levelsRead, _, err =
rr.readDefinitionLevels(defLevels[:batchSize])
+ if err != nil {
+ return recordsRead, err
+ }
+ repetitionLevelsRead, err :=
rr.readRepetitionLevels(repLevels[:batchSize])
+ if err != nil {
+ return recordsRead, err
+ }
+ if repetitionLevelsRead != levelsRead {
return 0, errors.New("number of decoded
rep/def levels did not match")
}
} else if rr.Descriptor().MaxDefinitionLevel() > 0 {
- levelsRead, _ =
rr.readDefinitionLevels(defLevels[:batchSize])
+ levelsRead, _, err =
rr.readDefinitionLevels(defLevels[:batchSize])
+ if err != nil {
+ return recordsRead, err
+ }
}
if levelsRead == 0 {
diff --git a/parquet/internal/encoding/boolean_decoder.go
b/parquet/internal/encoding/boolean_decoder.go
index 4d1fa98b..f46a8288 100644
--- a/parquet/internal/encoding/boolean_decoder.go
+++ b/parquet/internal/encoding/boolean_decoder.go
@@ -235,16 +235,21 @@ func (dec *RleBooleanDecoder) Decode(out []bool) (int,
error) {
for n > 0 {
batch := shared_utils.Min(len(buf), n)
- decoded := dec.rleDec.GetBatch(buf[:batch])
- if decoded != batch {
- return max - n, io.ErrUnexpectedEOF
- }
+ decoded, err := dec.rleDec.GetBatch(buf[:batch])
- for i := 0; i < batch; i++ {
+ for i := 0; i < decoded; i++ {
out[i] = buf[i] != 0
}
- n -= batch
- out = out[batch:]
+ n -= decoded
+ out = out[decoded:]
+ if err != nil {
+ dec.nvals -= max - n
+ return max - n, err
+ }
+ if decoded != batch {
+ dec.nvals -= max - n
+ return max - n, io.ErrUnexpectedEOF
+ }
}
dec.nvals -= max
diff --git a/parquet/internal/encoding/decoder.go
b/parquet/internal/encoding/decoder.go
index c45b021d..9c124da2 100644
--- a/parquet/internal/encoding/decoder.go
+++ b/parquet/internal/encoding/decoder.go
@@ -167,7 +167,7 @@ func (d *dictDecoder[T]) DecodeIndices(numValues int, bldr
array.Builder) (int,
d.idxScratchSpace = d.idxScratchSpace[:n]
}
- n = d.idxDecoder.GetBatch(d.idxScratchSpace)
+ n, err := d.idxDecoder.GetBatch(d.idxScratchSpace)
toAppend := make([]int, n)
for i, v := range d.idxScratchSpace {
@@ -175,7 +175,7 @@ func (d *dictDecoder[T]) DecodeIndices(numValues int, bldr
array.Builder) (int,
}
bldr.(*array.BinaryDictionaryBuilder).AppendIndices(toAppend, nil)
d.nvals -= n
- return n, nil
+ return n, err
}
func (d *dictDecoder[T]) DecodeIndicesSpaced(numValues, nullCount int,
validBits []byte, offset int64, bldr array.Builder) (int, error) {
diff --git a/parquet/internal/encoding/encoding_test.go
b/parquet/internal/encoding/encoding_test.go
index 26e87aa6..3ec9c4d2 100644
--- a/parquet/internal/encoding/encoding_test.go
+++ b/parquet/internal/encoding/encoding_test.go
@@ -18,7 +18,9 @@ package encoding_test
import (
"bufio"
+ "encoding/binary"
"fmt"
+ "io"
"os"
"path"
"reflect"
@@ -605,6 +607,24 @@ func TestDictEncoding(t *testing.T) {
}
}
+func TestDictionaryDecoderRejectsTruncatedPackedLiteralRun(t *testing.T) {
+ column := schema.NewColumn(schema.NewInt32Node("int32",
parquet.Repetitions.Required, -1), 0, 0)
+ dictionaryData := make([]byte, 8)
+ binary.LittleEndian.PutUint32(dictionaryData, 10)
+ binary.LittleEndian.PutUint32(dictionaryData[4:], 20)
+
+ dictionary := encoding.NewDecoder(parquet.Types.Int32,
parquet.Encodings.Plain, column, memory.DefaultAllocator)
+ require.NoError(t, dictionary.SetData(2, dictionaryData))
+
+ decoder := encoding.NewDictDecoder(parquet.Types.Int32, column,
memory.DefaultAllocator)
+ decoder.SetDict(dictionary)
+ indices := append([]byte{1, 17}, make([]byte, 7)...)
+ require.NoError(t, decoder.SetData(64, indices))
+
+ _, err := decoder.(encoding.Int32Decoder).Decode(make([]int32, 64))
+ assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+}
+
func TestWriteDeltaBitPackedInt32(t *testing.T) {
column := schema.NewColumn(schema.NewInt32Node("int32",
parquet.Repetitions.Required, -1), 0, 0)
diff --git a/parquet/internal/encoding/levels.go
b/parquet/internal/encoding/levels.go
index a63e9ab8..86c9b84f 100644
--- a/parquet/internal/encoding/levels.go
+++ b/parquet/internal/encoding/levels.go
@@ -21,6 +21,7 @@ import (
"encoding/binary"
"errors"
"fmt"
+ "io"
"math/bits"
"github.com/apache/arrow-go/v18/arrow/bitutil"
@@ -252,13 +253,14 @@ func (l *LevelDecoder) SetDataV2(nbytes int32, maxLvl
int16, nbuffered int, data
// Decode decodes the bytes that were set with SetData into the slice of levels
// returning the total number of levels that were decoded and the number of
// values which had a level equal to the max level, indicating how many
physical
-// values exist to be read.
-func (l *LevelDecoder) Decode(levels []int16) (int, int64) {
+// values exist to be read, along with any decoding error.
+func (l *LevelDecoder) Decode(levels []int16) (int, int64, error) {
var (
buf [1024]uint64
totaldecoded int
decoded int
valsToRead int64
+ err error
)
n := shared_utils.Min(int64(l.remaining), int64(len(levels)))
@@ -266,13 +268,13 @@ func (l *LevelDecoder) Decode(levels []int16) (int,
int64) {
batch := shared_utils.Min(1024, n)
switch l.encoding {
case format.Encoding_RLE:
- decoded = l.rle.GetBatch(buf[:batch])
+ decoded, err = l.rle.GetBatch(buf[:batch])
case format.Encoding_BIT_PACKED:
- decoded, _ = l.bit.GetBatch(uint(l.bitWidth),
buf[:batch])
+ decoded, err = l.bit.GetBatch(uint(l.bitWidth),
buf[:batch])
}
l.remaining -= decoded
totaldecoded += decoded
- n -= batch
+ n -= int64(decoded)
for idx, val := range buf[:decoded] {
lvl := int16(val)
@@ -282,7 +284,13 @@ func (l *LevelDecoder) Decode(levels []int16) (int, int64)
{
}
}
levels = levels[decoded:]
+ if err != nil {
+ return totaldecoded, valsToRead, err
+ }
+ if decoded != int(batch) {
+ return totaldecoded, valsToRead, io.ErrUnexpectedEOF
+ }
}
- return totaldecoded, valsToRead
+ return totaldecoded, valsToRead, nil
}
diff --git a/parquet/internal/encoding/levels_test.go
b/parquet/internal/encoding/levels_test.go
index b9378b66..52827d86 100644
--- a/parquet/internal/encoding/levels_test.go
+++ b/parquet/internal/encoding/levels_test.go
@@ -18,6 +18,7 @@ package encoding_test
import (
"encoding/binary"
+ "io"
"strconv"
"testing"
@@ -89,7 +90,8 @@ func verifyDecodingLvls(t *testing.T, enc parquet.Encoding,
maxLvl int16, input
// try multiple decoding on a single setdata call
for ct := 0; ct < decodeCount; ct++ {
offset := ct * numInnerLevels
- lvlCount, _ = decoder.Decode(output[:numInnerLevels])
+ lvlCount, _, err = decoder.Decode(output[:numInnerLevels])
+ assert.NoError(t, err)
assert.Equal(t, numInnerLevels, lvlCount)
assert.Equal(t, input[offset:offset+numInnerLevels],
output[:numInnerLevels])
}
@@ -101,12 +103,14 @@ func verifyDecodingLvls(t *testing.T, enc
parquet.Encoding, maxLvl int16, input
)
if remaining > 0 {
- lvlCount, _ = decoder.Decode(output[:remaining])
+ lvlCount, _, err = decoder.Decode(output[:remaining])
+ assert.NoError(t, err)
assert.Equal(t, remaining, lvlCount)
assert.Equal(t, input[levelsCompleted:], output[:remaining])
}
// test decode zero values
- lvlCount, _ = decoder.Decode(output[:1])
+ lvlCount, _, err = decoder.Decode(output[:1])
+ assert.NoError(t, err)
assert.Zero(t, lvlCount)
}
@@ -124,7 +128,8 @@ func verifyDecodingMultipleSetData(t *testing.T, enc
parquet.Encoding, max int16
assert.Len(t, output, numLevels)
_, err := decoder.SetData(enc, max, numLevels, buf[ct])
assert.NoError(t, err)
- lvlCount, _ = decoder.Decode(output)
+ lvlCount, _, err = decoder.Decode(output)
+ assert.NoError(t, err)
assert.Equal(t, numLevels, lvlCount)
assert.Equal(t, input[offset:offset+numLevels], output)
}
@@ -167,6 +172,22 @@ func TestLevelsDecodeMultipleBitWidth(t *testing.T) {
}
}
+func TestLevelDecoderPropagatesTruncatedRleLiteralRun(t *testing.T) {
+ encoded := append([]byte{17}, make([]byte, 7)...)
+ data := make([]byte, 4, 4+len(encoded))
+ binary.LittleEndian.PutUint32(data, uint32(len(encoded)))
+ data = append(data, encoded...)
+
+ var decoder encoding.LevelDecoder
+ _, err := decoder.SetData(parquet.Encodings.RLE, 1, 64, data)
+ assert.NoError(t, err)
+
+ levels := make([]int16, 64)
+ n, _, err := decoder.Decode(levels)
+ assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+ assert.Equal(t, 32, n)
+}
+
func TestLevelsDecodeMultipleSetData(t *testing.T) {
t.Parallel()
@@ -286,7 +307,8 @@ func TestEncodeDecodeLevels(t *testing.T) {
assert.NoError(t, err)
var levelOut [numToEncode]int16
- total, vals := decoder.Decode(levelOut[:])
+ total, vals, err := decoder.Decode(levelOut[:])
+ assert.NoError(t, err)
assert.EqualValues(t, numToEncode, total)
assert.EqualValues(t, numones, vals)
assert.Equal(t, levels, levelOut[:])
diff --git a/parquet/internal/utils/bit_packing_avx2_amd64.go
b/parquet/internal/utils/bit_packing_avx2_amd64.go
index 5f1923fa..64fa86d9 100644
--- a/parquet/internal/utils/bit_packing_avx2_amd64.go
+++ b/parquet/internal/utils/bit_packing_avx2_amd64.go
@@ -31,23 +31,30 @@ var bufferPool = sync.Pool{New: func() interface{} { return
&bytes.Buffer{} }}
//go:noescape
func _unpack32_avx2(in, out unsafe.Pointer, batchSize, nbits int) (num int)
-func unpack32Avx2(in io.Reader, out []uint32, nbits int) int {
+func unpack32Avx2(in io.Reader, out []uint32, nbits int) (int, error) {
batch := len(out) / 32 * 32
n := batch * nbits / 8
if n <= 0 {
- return 0
+ return 0, nil
}
buffer := bufferPool.Get().(*bytes.Buffer)
defer bufferPool.Put(buffer)
buffer.Reset()
buffer.Grow(n)
- io.CopyN(buffer, in, int64(n))
+ nread, err := io.CopyN(buffer, in, int64(n))
+ if err == io.EOF && int(nread)%(nbits*4) != 0 {
+ err = io.ErrUnexpectedEOF
+ }
+ completeBatch := int(nread) * 8 / nbits / 32 * 32
+ if completeBatch == 0 {
+ return 0, err
+ }
var (
input = unsafe.Pointer(&buffer.Bytes()[0])
output = unsafe.Pointer(&out[0])
)
- return _unpack32_avx2(input, output, len(out), nbits)
+ return _unpack32_avx2(input, output, completeBatch, nbits), err
}
diff --git a/parquet/internal/utils/bit_packing_default.go
b/parquet/internal/utils/bit_packing_default.go
index fa814f63..27f4b275 100644
--- a/parquet/internal/utils/bit_packing_default.go
+++ b/parquet/internal/utils/bit_packing_default.go
@@ -17,11 +17,12 @@
package utils
import (
+ "bytes"
"encoding/binary"
"io"
)
-var unpack32 func(io.Reader, []uint32, int) int = unpack32Default
+var unpack32 func(io.Reader, []uint32, int) (int, error) = unpack32Default
type unpackFunc func(in io.Reader, out []uint32)
@@ -1860,7 +1861,7 @@ func nullunpack32(_ io.Reader, out []uint32) {
}
}
-func unpack32Default(in io.Reader, out []uint32, nbits int) int {
+func unpack32Default(in io.Reader, out []uint32, nbits int) (int, error) {
batch := len(out) / 32 * 32
nloops := batch / 32
@@ -1933,11 +1934,16 @@ func unpack32Default(in io.Reader, out []uint32, nbits
int) int {
case 32:
f = unpack32_32
default:
- return 0
+ return 0, nil
}
+ var packed [128]byte
+ packedBytes := nbits * 4
for i := 0; i < nloops; i++ {
- f(in, out[i*32:])
+ if _, err := io.ReadFull(in, packed[:packedBytes]); err != nil {
+ return i * 32, err
+ }
+ f(bytes.NewReader(packed[:packedBytes]), out[i*32:])
}
- return batch
+ return batch, nil
}
diff --git a/parquet/internal/utils/bit_packing_neon_arm64.go
b/parquet/internal/utils/bit_packing_neon_arm64.go
index 580f9a1f..e474f64c 100644
--- a/parquet/internal/utils/bit_packing_neon_arm64.go
+++ b/parquet/internal/utils/bit_packing_neon_arm64.go
@@ -31,23 +31,30 @@ var bufferPool = sync.Pool{New: func() interface{} { return
&bytes.Buffer{} }}
//go:noescape
func _unpack32_neon(in, out unsafe.Pointer, batchSize, nbits int) (num int)
-func unpack32NEON(in io.Reader, out []uint32, nbits int) int {
+func unpack32NEON(in io.Reader, out []uint32, nbits int) (int, error) {
batch := len(out) / 32 * 32
n := batch * nbits / 8
if n <= 0 {
- return 0
+ return 0, nil
}
buffer := bufferPool.Get().(*bytes.Buffer)
defer bufferPool.Put(buffer)
buffer.Reset()
buffer.Grow(n)
- io.CopyN(buffer, in, int64(n))
+ nread, err := io.CopyN(buffer, in, int64(n))
+ if err == io.EOF && int(nread)%(nbits*4) != 0 {
+ err = io.ErrUnexpectedEOF
+ }
+ completeBatch := int(nread) * 8 / nbits / 32 * 32
+ if completeBatch == 0 {
+ return 0, err
+ }
var (
input = unsafe.Pointer(&buffer.Bytes()[0])
output = unsafe.Pointer(&out[0])
)
- return _unpack32_neon(input, output, len(out), nbits)
+ return _unpack32_neon(input, output, completeBatch, nbits), err
}
diff --git a/parquet/internal/utils/bit_reader.go
b/parquet/internal/utils/bit_reader.go
index 5feaa8fa..c7e2e79c 100644
--- a/parquet/internal/utils/bit_reader.go
+++ b/parquet/internal/utils/bit_reader.go
@@ -287,9 +287,12 @@ func (b *BitReader) GetBatchIndex(bits uint, out
[]IndexType) (i int, err error)
b.reader.Seek(b.byteoffset, io.SeekStart)
// grab as many 32 byte chunks as possible in one shot
if i < length { // IndexType should be a 32 bit value so we can do
quick unpacking right into the output
- numUnpacked := unpack32(b.reader,
(*(*[]uint32)(unsafe.Pointer(&out)))[i:], int(bits))
+ numUnpacked, unpackErr := unpack32(b.reader,
(*(*[]uint32)(unsafe.Pointer(&out)))[i:], int(bits))
i += numUnpacked
b.byteoffset += int64(numUnpacked * int(bits) / 8)
+ if unpackErr != nil {
+ return i, unpackErr
+ }
}
// re-fill our buffer just in case.
@@ -455,16 +458,19 @@ func (b *BitReader) GetBatch(bits uint, out []uint64)
(int, error) {
for i < length {
// unpack groups of 32 bytes at a time into a buffer since it's
more efficient
unpackSize := utils.Min(buflen, length-i)
- numUnpacked := unpack32(b.reader, b.unpackBuf[:unpackSize],
int(bits))
- if numUnpacked == 0 {
- break
- }
+ numUnpacked, err := unpack32(b.reader,
b.unpackBuf[:unpackSize], int(bits))
for k := 0; k < numUnpacked; k++ {
out[i+k] = uint64(b.unpackBuf[k])
}
i += numUnpacked
b.byteoffset += int64(numUnpacked * int(bits) / 8)
+ if err != nil {
+ return i, err
+ }
+ if numUnpacked == 0 {
+ break
+ }
}
b.fillbuffer()
diff --git a/parquet/internal/utils/bit_reader_test.go
b/parquet/internal/utils/bit_reader_test.go
index b0384971..9df2f617 100644
--- a/parquet/internal/utils/bit_reader_test.go
+++ b/parquet/internal/utils/bit_reader_test.go
@@ -252,6 +252,147 @@ func TestBitArrayVals(t *testing.T) {
}
}
+func TestBitReaderRejectsTruncatedPackedBatches(t *testing.T) {
+ for width := 1; width <= 32; width++ {
+ for _, batchSize := range []int{32, 64} {
+ t.Run(fmt.Sprintf("width=%d/batch=%d", width,
batchSize), func(t *testing.T) {
+ input := bytes.Repeat([]byte{0xFF},
width*batchSize/8-1)
+ expected := batchSize - 32
+ want := uint64(1<<uint(width)) - 1
+
+ values := make([]uint64, batchSize)
+ for i := range values {
+ values[i] = math.MaxUint64
+ }
+ reader :=
utils.NewBitReader(bytes.NewReader(input))
+ n, err := reader.GetBatch(uint(width), values)
+ assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+ assert.Equal(t, expected, n)
+ for _, got := range values[:n] {
+ assert.Equal(t, want, got)
+ }
+ for _, got := range values[n:] {
+ assert.Equal(t, uint64(math.MaxUint64),
got)
+ }
+
+ indices := make([]utils.IndexType, batchSize)
+ for i := range indices {
+ indices[i] = math.MinInt32
+ }
+ reader =
utils.NewBitReader(bytes.NewReader(input))
+ n, err = reader.GetBatchIndex(uint(width),
indices)
+ assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+ assert.Equal(t, expected, n)
+ for _, got := range indices[:n] {
+ assert.Equal(t, int32(want), got)
+ }
+ for _, got := range indices[n:] {
+ assert.Equal(t, int32(math.MinInt32),
got)
+ }
+ })
+ }
+ }
+}
+
+func TestBitReaderStopsAtPackedGroupBoundary(t *testing.T) {
+ for width := 1; width <= 32; width++ {
+ t.Run(fmt.Sprintf("width=%d", width), func(t *testing.T) {
+ input := bytes.Repeat([]byte{0xFF}, width*4)
+
+ reader := utils.NewBitReader(bytes.NewReader(input))
+ values := make([]uint64, 64)
+ n, err := reader.GetBatch(uint(width), values)
+ assert.ErrorIs(t, err, io.EOF)
+ assert.Equal(t, 32, n)
+
+ reader = utils.NewBitReader(bytes.NewReader(input))
+ indices := make([]utils.IndexType, 64)
+ n, err = reader.GetBatchIndex(uint(width), indices)
+ assert.ErrorIs(t, err, io.EOF)
+ assert.Equal(t, 32, n)
+ })
+ }
+}
+
+type identityDictionaryConverter struct{}
+
+func (identityDictionaryConverter) Copy(out []uint64, indices
[]utils.IndexType) error {
+ for i, idx := range indices {
+ out[i] = uint64(idx)
+ }
+ return nil
+}
+
+func (identityDictionaryConverter) Fill(out []uint64, index utils.IndexType)
error {
+ for i := range out {
+ out[i] = uint64(index)
+ }
+ return nil
+}
+
+func (identityDictionaryConverter) FillZero(out []uint64) {
+ clear(out)
+}
+
+func (identityDictionaryConverter) IsValid(...utils.IndexType) bool { return
true }
+
+func (identityDictionaryConverter) IsValidSingle(utils.IndexType) bool {
return true }
+
+func TestRleDecoderPropagatesTruncatedPackedRun(t *testing.T) {
+ for _, tc := range []struct {
+ name string
+ payload int
+ err error
+ }{
+ {name: "partial group", payload: 7, err: io.ErrUnexpectedEOF},
+ {name: "group boundary", payload: 4, err: io.EOF},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ input := append([]byte{17}, bytes.Repeat([]byte{0xFF},
tc.payload)...)
+ decoder := utils.NewRleDecoder(bytes.NewReader(input),
1)
+
+ values := make([]uint64, 64)
+ n, err := decoder.GetBatch(values)
+ assert.ErrorIs(t, err, tc.err)
+ assert.Equal(t, 32, n)
+ for _, value := range values[:n] {
+ assert.Equal(t, uint64(1), value)
+ }
+ })
+ }
+}
+
+func TestTypedRleDecoderPropagatesTruncatedPackedRun(t *testing.T) {
+ converter := identityDictionaryConverter{}
+ for _, tc := range []struct {
+ name string
+ payload int
+ err error
+ }{
+ {name: "partial group", payload: 7, err: io.ErrUnexpectedEOF},
+ {name: "group boundary", payload: 4, err: io.EOF},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ input := append([]byte{17}, bytes.Repeat([]byte{0xFF},
tc.payload)...)
+
+ t.Run("unspaced", func(t *testing.T) {
+ decoder :=
utils.NewTypedRleDecoder[uint64](bytes.NewReader(input), 1)
+ _, err := decoder.GetBatchWithDict(converter,
make([]uint64, 64))
+ assert.ErrorIs(t, err, tc.err)
+ })
+
+ t.Run("spaced", func(t *testing.T) {
+ decoder :=
utils.NewTypedRleDecoder[uint64](bytes.NewReader(input), 1)
+ validBits := bytes.Repeat([]byte{0xFF}, 9)
+ bitutil.ClearBit(validBits, 0)
+
+ _, err :=
decoder.GetBatchWithDictSpaced(converter, make([]uint64, 65), 1, validBits, 0)
+ assert.ErrorIs(t, err, tc.err)
+ })
+ })
+ }
+}
+
func TestMixedValues(t *testing.T) {
const buflen = 1024
buf := make([]byte, buflen)
@@ -381,7 +522,9 @@ func (r *RLETestSuite) ValidateRle(vals []uint64, width
int, expected []byte, ex
r.Run("decode batch read", func() {
dec := utils.NewRleDecoder(bytes.NewReader(buf), width)
check := make([]uint64, len(vals))
- r.Equal(len(vals), dec.GetBatch(check))
+ n, err := dec.GetBatch(check)
+ r.NoError(err)
+ r.Equal(len(vals), n)
r.Equal(vals, check)
})
}
@@ -603,7 +746,9 @@ func (r *RLERandomSuite) checkRoundTrip(vals []uint64,
width int) bool {
res = res && r.Run("batch decode", func() {
dec := utils.NewRleDecoder(bytes.NewReader(buf[:encoded]),
width)
read := make([]uint64, len(vals))
- r.Require().Equal(len(vals), dec.GetBatch(read))
+ n, err := dec.GetBatch(read)
+ r.Require().NoError(err)
+ r.Require().Equal(len(vals), n)
r.Equal(vals, read)
})
diff --git a/parquet/internal/utils/rle.go b/parquet/internal/utils/rle.go
index 23a86598..596f2018 100644
--- a/parquet/internal/utils/rle.go
+++ b/parquet/internal/utils/rle.go
@@ -170,7 +170,7 @@ func (r *RleDecoder) Next() bool {
func (r *RleDecoder) GetValue() (uint64, bool) {
vals := make([]uint64, 1)
- n := r.GetBatch(vals)
+ n, _ := r.GetBatch(vals)
return vals[0], n == 1
}
@@ -201,7 +201,7 @@ func (r *RleDecoder) Discard(n int) int {
return read
}
-func (r *RleDecoder) GetBatch(values []uint64) int {
+func (r *RleDecoder) GetBatch(values []uint64) (int, error) {
read := 0
size := len(values)
@@ -220,26 +220,28 @@ func (r *RleDecoder) GetBatch(values []uint64) int {
out = out[repbatch:]
} else if r.litCount > 0 {
litbatch := min(remain, int(r.litCount))
- n, _ := r.r.GetBatch(uint(r.bitWidth), out[:litbatch])
+ n, err := r.r.GetBatch(uint(r.bitWidth), out[:litbatch])
+ r.litCount -= int32(n)
+ read += n
+ out = out[n:]
+ if err != nil {
+ return read, err
+ }
if n != litbatch {
- return read
+ return read, nil
}
-
- r.litCount -= int32(litbatch)
- read += litbatch
- out = out[litbatch:]
} else {
if !r.Next() {
- return read
+ return read, nil
}
}
}
- return read
+ return read, nil
}
func (r *RleDecoder) GetBatchSpaced(vals []uint64, nullcount int, validBits
[]byte, validBitsOffset int64) (int, error) {
if nullcount == 0 {
- return r.GetBatch(vals), nil
+ return r.GetBatch(vals)
}
converter := plainConverter[uint64]{}
@@ -259,7 +261,7 @@ func (r *RleDecoder) GetBatchSpaced(vals []uint64,
nullcount int, validBits []by
}
if block.AllSet() {
- processed = r.GetBatch(vals[:block.Len])
+ processed, err = r.GetBatch(vals[:block.Len])
} else if block.NoneSet() {
converter.FillZero(vals[:block.Len])
processed = int(block.Len)
@@ -271,6 +273,9 @@ func (r *RleDecoder) GetBatchSpaced(vals []uint64,
nullcount int, validBits []by
}
totalProcessed += processed
+ if err != nil {
+ return totalProcessed, err
+ }
vals = vals[int(block.Len):]
validBitsOffset += int64(block.Len)
diff --git a/parquet/internal/utils/typed_rle_dict.go
b/parquet/internal/utils/typed_rle_dict.go
index eafae2b8..07785cc2 100644
--- a/parquet/internal/utils/typed_rle_dict.go
+++ b/parquet/internal/utils/typed_rle_dict.go
@@ -91,7 +91,10 @@ func consumeLiterals[T parquet.ColumnTypes | uint64](r
*RleDecoder, dc Dictionar
batch := min(remain, int(r.litCount), len(buf))
buf = buf[:batch]
- n, _ := r.r.GetBatchIndex(uint(r.bitWidth), buf)
+ n, err := r.r.GetBatchIndex(uint(r.bitWidth), buf)
+ if err != nil {
+ return 0, 0, run, err
+ }
if n != batch {
return 0, 0, run, errors.New("was not able to retrieve correct
number of indexes")
}
@@ -163,7 +166,10 @@ func (r *TypedRleDecoder[T]) GetBatchWithDict(dc
DictionaryConverter[T], vals []
case r.litCount > 0:
litbatch := min(remain, int(r.litCount), 1024)
buf := indexbuffer[:litbatch]
- n, _ := r.r.GetBatchIndex(uint(r.bitWidth), buf)
+ n, err := r.r.GetBatchIndex(uint(r.bitWidth), buf)
+ if err != nil {
+ return read, err
+ }
if n != litbatch {
return read, nil
}
