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 ced0f265..d46f0325 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. @@ -404,16 +407,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 366ae408..f34c323d 100644 --- a/parquet/internal/utils/bit_reader_test.go +++ b/parquet/internal/utils/bit_reader_test.go @@ -20,6 +20,7 @@ import ( "bytes" "encoding/binary" "fmt" + "io" "math" "math/bits" "math/rand/v2" @@ -151,6 +152,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< 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 }