fix(parquet): reject truncated bit-packed batches

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)
 			}

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")
 	}

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 {

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

parquet/internal/encoding/decoder.go

@@ -167,15 +167,15 @@ 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 {
 		toAppend[i] = int(v)
 	}
 	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) {

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)
 

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,27 +253,28 @@ 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)))
 	for n > 0 {
 		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
 }

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[:])