diff --git a/parquet/metadata/bloom_filter.go b/parquet/metadata/bloom_filter.go index b607c5f1..562d0d44 100644 --- a/parquet/metadata/bloom_filter.go +++ b/parquet/metadata/bloom_filter.go @@ -493,6 +493,9 @@ func (r *RowGroupBloomFilterReader) GetColumnBloomFilter(i int) (BloomFilter, er headerBuf := r.bufferPool.Get().(*memory.Buffer) headerBuf.ResizeNoShrink(int(bloomFilterReadSize)) + if sectionSize := sectionRdr.Size(); sectionSize < int64(headerBuf.Len()) { + headerBuf.ResizeNoShrink(int(sectionSize)) + } defer func() { if headerBuf != nil { headerBuf.ResizeNoShrink(0) @@ -500,7 +503,7 @@ func (r *RowGroupBloomFilterReader) GetColumnBloomFilter(i int) (BloomFilter, er } }() - if _, err = sectionRdr.Read(headerBuf.Bytes()); err != nil { + if _, err = io.ReadFull(sectionRdr, headerBuf.Bytes()); err != nil { return nil, err } @@ -527,7 +530,7 @@ func (r *RowGroupBloomFilterReader) GetColumnBloomFilter(i int) (BloomFilter, er copy(buf.Bytes(), headerBuf.Bytes()[headerSize:]) } - if _, err = sectionRdr.Read(buf.Bytes()[filterBytesInHeader:]); err != nil { + if _, err = io.ReadFull(sectionRdr, buf.Bytes()[filterBytesInHeader:]); err != nil { return nil, err } bitset = buf.Bytes() diff --git a/parquet/metadata/bloom_filter_reader_test.go b/parquet/metadata/bloom_filter_reader_test.go index 7ae7c39b..e2c4111c 100644 --- a/parquet/metadata/bloom_filter_reader_test.go +++ b/parquet/metadata/bloom_filter_reader_test.go @@ -18,6 +18,7 @@ package metadata_test import ( "bytes" + "io" "os" "runtime" "sync" @@ -35,6 +36,57 @@ import ( "github.com/stretchr/testify/suite" ) +// shortReadAtSeeker deliberately returns short successful ReadAt calls to +// exercise callers that must keep reading until their requested buffer is full. +type shortReadAtSeeker struct { + data []byte + offset int64 + max int +} + +func (r *shortReadAtSeeker) ReadAt(p []byte, off int64) (int, error) { + return r.readAt(p, off) +} + +func (r *shortReadAtSeeker) Seek(offset int64, whence int) (int64, error) { + var next int64 + switch whence { + case io.SeekStart: + next = offset + case io.SeekCurrent: + next = r.offset + offset + case io.SeekEnd: + next = int64(len(r.data)) + offset + default: + return 0, os.ErrInvalid + } + if next < 0 { + return 0, os.ErrInvalid + } + r.offset = next + return next, nil +} + +func (r *shortReadAtSeeker) readAt(p []byte, off int64) (int, error) { + if off >= int64(len(r.data)) { + return 0, io.EOF + } + + n := len(p) + if r.max > 0 && n > r.max { + n = r.max + } + if remaining := len(r.data) - int(off); n > remaining { + n = remaining + } + start := int(off) + copy(p, r.data[start:start+n]) + if n < len(p) && int(off)+n >= len(r.data) { + return n, io.EOF + } + return n, nil +} + type BloomFilterBuilderSuite struct { suite.Suite @@ -283,6 +335,91 @@ func TestBloomFilterRoundTrip(t *testing.T) { suite.Run(t, new(EncryptedBloomFilterBuilderSuite)) } +func TestBloomFilterReaderHandlesShortReads(t *testing.T) { + data, fileMeta, insertedHash := makeBloomFilterReaderTestData(t, 1024) + + t.Run("with bloom filter length", func(t *testing.T) { + bf := readBloomFilterWithShortReads(t, data, fileMeta, 7) + assert.True(t, bf.CheckHash(insertedHash)) + }) + + t.Run("without bloom filter length", func(t *testing.T) { + data, fileMeta, insertedHash := makeBloomFilterReaderTestData(t, 1024) + fileMeta.RowGroups[0].Columns[0].MetaData.BloomFilterLength = nil + + bf := readBloomFilterWithShortReads(t, data, fileMeta, 256) + assert.True(t, bf.CheckHash(insertedHash)) + }) + + t.Run("without bloom filter length and small section", func(t *testing.T) { + data, fileMeta, insertedHash := makeBloomFilterReaderTestData(t, 32) + fileMeta.RowGroups[0].Columns[0].MetaData.BloomFilterLength = nil + + bf := readBloomFilterWithShortReads(t, data, fileMeta, 7) + assert.True(t, bf.CheckHash(insertedHash)) + }) +} + +func makeBloomFilterReaderTestData(t *testing.T, filterSize uint32) ([]byte, *metadata.FileMetaData, uint64) { + t.Helper() + + sc := schema.NewSchema(schema.MustGroup( + schema.NewGroupNode("schema", parquet.Repetitions.Repeated, + schema.FieldList{ + schema.NewByteArrayNode("c1", parquet.Repetitions.Optional, -1), + }, -1))) + + props := parquet.NewWriterProperties() + bldr := metadata.FileBloomFilterBuilder{Schema: sc} + metaBldr := metadata.NewFileMetadataBuilder(sc, props, nil) + rgMeta := metaBldr.AppendRowGroup() + filterMap := make(map[string]metadata.BloomFilterBuilder) + bldr.AppendRowGroup(rgMeta, filterMap) + + bf := metadata.NewBloomFilter(filterSize, filterSize, memory.NewGoAllocator()) + insertedHash := (uint64(^uint32(0)) << 32) | 1 + bf.InsertHash(insertedHash) + filterMap["c1"] = bf + + rgMeta.NextColumnChunk() + require.NoError(t, rgMeta.Finish(0, 0)) + + var buf bytes.Buffer + wr := &utils.TellWrapper{Writer: &buf} + _, err := wr.Write([]byte("PAR1")) + require.NoError(t, err) + require.NoError(t, bldr.WriteTo(wr)) + + fileMeta, err := metaBldr.Finish() + require.NoError(t, err) + fileMeta.SetSourceFileSize(int64(buf.Len())) + + return buf.Bytes(), fileMeta, insertedHash +} + +func readBloomFilterWithShortReads(t *testing.T, data []byte, fileMeta *metadata.FileMetaData, maxRead int) metadata.BloomFilter { + t.Helper() + + rdr := metadata.BloomFilterReader{ + Input: &shortReadAtSeeker{data: data, max: maxRead}, + FileMetadata: fileMeta, + BufferPool: &sync.Pool{ + New: func() interface{} { + return memory.NewResizableBuffer(memory.DefaultAllocator) + }, + }, + } + + rg, err := rdr.RowGroup(0) + require.NoError(t, err) + require.NotNil(t, rg) + + bf, err := rg.GetColumnBloomFilter(0) + require.NoError(t, err) + require.NotNil(t, bf) + return bf +} + func TestReadBloomFilter(t *testing.T) { dir := os.Getenv("PARQUET_TEST_DATA") if dir == "" {