This is an automated email from the ASF dual-hosted git repository.
laskoviymishka pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new 28cd1baa5 perf(table): prune parquet row groups by dictionary (#2006)
28cd1baa5 is described below
commit 28cd1baa51ad36e082f2e320470592675c4db0a8
Author: Minh Vu <[email protected]>
AuthorDate: Tue Sep 22 23:44:05 2026 +0200
perf(table): prune parquet row groups by dictionary (#2006)
* perf(table): prune parquet row groups by dictionary
* perf(table): tighten dictionary pruning
* fix(table): remove redundant benchmark capture
* fix(table): address dictionary pruning review feedback
* perf(table): make dictionary pruning adaptive
* fix(table): harden dictionary row group pruning
* perf(table): accelerate large dictionary predicates
* fix(table): address dictionary pruning review
---
table/arrow_filter_plan.go | 12 +-
table/arrow_filter_plan_test.go | 5 +
table/arrow_scanner.go | 17 +-
table/arrow_scanner_pruning_internal_test.go | 8 +
table/evaluators.go | 34 ++
table/internal/export_test.go | 16 +-
table/internal/parquet_dictionary_test.go | 321 ++++++++++++++++
table/internal/parquet_files.go | 461 ++++++++++++++++++++++-
table/internal/parquet_files_test.go | 531 ++++++++++++++++++++++++++-
9 files changed, 1384 insertions(+), 21 deletions(-)
diff --git a/table/arrow_filter_plan.go b/table/arrow_filter_plan.go
index 66a619c95..e839ce820 100644
--- a/table/arrow_filter_plan.go
+++ b/table/arrow_filter_plan.go
@@ -37,8 +37,9 @@ import (
// from statsFilter for each file because inclusiveMetricsEval stores mutable
// per-row-group maps.
type compiledFileFilterPlan struct {
- statsFilter iceberg.BooleanExpression
- bloomPreds []tblutils.RowGroupBloomPred
+ statsFilter iceberg.BooleanExpression
+ bloomPreds []tblutils.RowGroupBloomPred
+ dictionaryPreds []tblutils.RowGroupDictionaryPred
recordFilter expr.Expression
extensionRegistry *expr.ExtensionRegistry
@@ -299,7 +300,7 @@ func compileFileFilterPlan(
if includePruning {
// Variant extract terms have no row-group statistics. The spec
prunes them
// at the file level via variant bounds (format/spec.md "Bounds
for Variant"),
- // so they are excluded from the row-group stats/bloom filter
here.
+ // so they are excluded from the row-group
stats/dictionary/bloom filter here.
pruneFilter := boundFilter
if len(extracts) > 0 {
stripped, err := stripExtractPredicates(rowFilter)
@@ -326,9 +327,14 @@ func compileFileFilterPlan(
if err != nil {
return nil, err
}
+ dictionaryPreds, err :=
newDictionaryPredicatesFromRewritten(statsFilter)
+ if err != nil {
+ return nil, err
+ }
plan.statsFilter = statsFilter
plan.bloomPreds = bloomPreds
+ plan.dictionaryPreds = dictionaryPreds
}
if !includeRecordFilter {
return plan, nil
diff --git a/table/arrow_filter_plan_test.go b/table/arrow_filter_plan_test.go
index 5d23e94da..76f3bf87f 100644
--- a/table/arrow_filter_plan_test.go
+++ b/table/arrow_filter_plan_test.go
@@ -83,8 +83,12 @@ func TestCompiledFileFilterPlansSeparatePhysicalTypes(t
*testing.T) {
assert.Len(t, scan.filterPlanCache.plans, 2)
require.Len(t, int32Plans.pruning.bloomPreds, 1)
require.Len(t, int64Plans.pruning.bloomPreds, 1)
+ require.Len(t, int32Plans.pruning.dictionaryPreds, 1)
+ require.Len(t, int64Plans.pruning.dictionaryPreds, 1)
assert.Equal(t, []byte{1, 0, 0, 0},
int32Plans.pruning.bloomPreds[0].PhysBytes[0])
assert.Equal(t, []byte{1, 0, 0, 0, 0, 0, 0, 0},
int64Plans.pruning.bloomPreds[0].PhysBytes[0])
+ assert.Equal(t, []byte{1, 0, 0, 0},
int32Plans.pruning.dictionaryPreds[0].PhysBytes[0])
+ assert.Equal(t, []byte{1, 0, 0, 0, 0, 0, 0, 0},
int64Plans.pruning.dictionaryPreds[0].PhysBytes[0])
}
func TestPhysicalSchemaKeyIncludesNestedFieldIDs(t *testing.T) {
@@ -195,4 +199,5 @@ func
TestCompiledFileFilterPlansDisablePruningForMissingInitialDefault(t *testin
assert.True(t, plans.pruning.statsFilter.Equals(iceberg.AlwaysTrue{}))
assert.Empty(t, plans.pruning.bloomPreds)
+ assert.Empty(t, plans.pruning.dictionaryPreds)
}
diff --git a/table/arrow_scanner.go b/table/arrow_scanner.go
index 129cda66b..64778ab77 100644
--- a/table/arrow_scanner.go
+++ b/table/arrow_scanner.go
@@ -1728,7 +1728,7 @@ func (as *arrowScan) processRecordsWithPlans(
pruningFilter = iceberg.AlwaysTrue{}
}
- // Row-group stats/bloom pruning skips whole groups, so emitted batches
no
+ // Row-group stats/dictionary/bloom pruning skips whole groups, so
emitted batches no
// longer cover contiguous file positions. Steps that key on the
original
// position (row-lineage _row_id, positional/DV deletes, generated
position
// deletes) recover it from posSource, which the tester seeds with each
@@ -1739,8 +1739,9 @@ func (as *arrowScan) processRecordsWithPlans(
var tester *tblutils.ParquetRowGroupTester
if plans != nil && plans.pruning != nil {
tester = &tblutils.ParquetRowGroupTester{
- StatsFn: plans.pruning.statsEvaluator(),
- BloomPreds: plans.pruning.bloomPreds,
+ StatsFn: plans.pruning.statsEvaluator(),
+ BloomPreds: plans.pruning.bloomPreds,
+ DictionaryPreds: plans.pruning.dictionaryPreds,
}
} else {
logicalSchema := as.filterSchema
@@ -1786,9 +1787,15 @@ func (as *arrowScan) processRecordsWithPlans(
return err
}
+ dictionaryPreds, err :=
newDictionaryPredicates(filePruningFilter)
+ if err != nil {
+ return err
+ }
+
tester = &tblutils.ParquetRowGroupTester{
- StatsFn: statsFn,
- BloomPreds: bloomPreds,
+ StatsFn: statsFn,
+ BloomPreds: bloomPreds,
+ DictionaryPreds: dictionaryPreds,
}
}
// A complete task already reads every row group. Leave its
byte range
diff --git a/table/arrow_scanner_pruning_internal_test.go
b/table/arrow_scanner_pruning_internal_test.go
index a9df9f741..668125afe 100644
--- a/table/arrow_scanner_pruning_internal_test.go
+++ b/table/arrow_scanner_pruning_internal_test.go
@@ -121,6 +121,9 @@ func TestProcessRecordsUsesRowGroupFilterForPruning(t
*testing.T) {
require.Len(t, reader.tester.BloomPreds, 1)
assert.Equal(t, 1, reader.tester.BloomPreds[0].FieldID)
assert.Len(t, reader.tester.BloomPreds[0].PhysBytes, 1)
+ require.Len(t, reader.tester.DictionaryPreds, 1)
+ assert.Equal(t, 1, reader.tester.DictionaryPreds[0].FieldID)
+ assert.Equal(t, reader.tester.BloomPreds[0].PhysBytes,
reader.tester.DictionaryPreds[0].PhysBytes)
result := <-out
result.Record.Value.Release()
@@ -183,6 +186,9 @@ func
TestProcessRecordsRebindsRowGroupFilterToPromotedFileSchema(t *testing.T) {
require.Len(t, reader.tester.BloomPreds, 1)
assert.Equal(t, []byte{1, 0, 0, 0},
reader.tester.BloomPreds[0].PhysBytes[0],
"the bloom predicate should use the INT32 file encoding")
+ require.Len(t, reader.tester.DictionaryPreds, 1)
+ assert.Equal(t, []byte{1, 0, 0, 0},
reader.tester.DictionaryPreds[0].PhysBytes[0],
+ "the dictionary predicate should use the INT32 file encoding")
result := <-out
result.Record.Value.Release()
@@ -227,6 +233,8 @@ func TestProcessRecordsDoesNotPruneMissingInitialDefault(t
*testing.T) {
require.NotNil(t, reader.tester)
assert.Empty(t, reader.tester.BloomPreds,
"a missing field with an initial-default must disable bloom
pruning")
+ assert.Empty(t, reader.tester.DictionaryPreds,
+ "a missing field with an initial-default must disable
dictionary pruning")
result := <-out
result.Record.Value.Release()
diff --git a/table/evaluators.go b/table/evaluators.go
index 65b5276e6..660d969ae 100644
--- a/table/evaluators.go
+++ b/table/evaluators.go
@@ -1940,3 +1940,37 @@ func newBloomFilterPredicatesFromRewritten(expr
iceberg.BooleanExpression) ([]in
return iceberg.VisitExpr(expr, &bloomPredicateCollector{})
}
+
+// newDictionaryPredicates reuses the same conservative EqualTo/In collector
+// as Bloom filters. The physical literal bytes are also the representation
+// needed to compare values decoded from a PLAIN dictionary page.
+func newDictionaryPredicates(expr iceberg.BooleanExpression)
([]internal.RowGroupDictionaryPred, error) {
+ if expr == nil {
+ return nil, nil
+ }
+
+ rewritten, err := iceberg.RewriteNotExpr(expr)
+ if err != nil {
+ return nil, err
+ }
+
+ return newDictionaryPredicatesFromRewritten(rewritten)
+}
+
+func newDictionaryPredicatesFromRewritten(expr iceberg.BooleanExpression)
([]internal.RowGroupDictionaryPred, error) {
+ bloomPreds, err := newBloomFilterPredicatesFromRewritten(expr)
+ if err != nil {
+ return nil, err
+ }
+
+ if len(bloomPreds) == 0 {
+ return nil, nil
+ }
+
+ dictionaryPreds := make([]internal.RowGroupDictionaryPred,
len(bloomPreds))
+ for i, pred := range bloomPreds {
+ dictionaryPreds[i] = internal.RowGroupDictionaryPred(pred)
+ }
+
+ return dictionaryPreds, nil
+}
diff --git a/table/internal/export_test.go b/table/internal/export_test.go
index 740b6327f..d9f76601c 100644
--- a/table/internal/export_test.go
+++ b/table/internal/export_test.go
@@ -17,9 +17,23 @@
package internal
-import "github.com/apache/arrow-go/v18/parquet/pqarrow"
+import (
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/pqarrow"
+)
// WrapParquetFileReader exposes wrapPqArrowReader for testing.
func WrapParquetFileReader(fr *pqarrow.FileReader) FileReader {
return wrapPqArrowReader{FileReader: fr}
}
+
+// WrapParquetFileReaderWithDictionarySource exposes the dictionary source for
+// tests that need to exercise the section-read path used by production scans.
+func WrapParquetFileReaderWithDictionarySource(fr *pqarrow.FileReader, source
parquet.ReaderAtSeeker) FileReader {
+ return wrapPqArrowReader{
+ FileReader: fr,
+ dictionarySource: source,
+ dictionaryMem: memory.DefaultAllocator,
+ }
+}
diff --git a/table/internal/parquet_dictionary_test.go
b/table/internal/parquet_dictionary_test.go
new file mode 100644
index 000000000..4da15fcdf
--- /dev/null
+++ b/table/internal/parquet_dictionary_test.go
@@ -0,0 +1,321 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package internal
+
+import (
+ "encoding/binary"
+ "math"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/file"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestDictionaryMatchesPredicatesByteArray(t *testing.T) {
+ data := make([]byte, 0, 32)
+ appendValue := func(value string) {
+ var length [4]byte
+ binary.LittleEndian.PutUint32(length[:], uint32(len(value)))
+ data = append(data, length[:]...)
+ data = append(data, value...)
+ }
+ appendValue("first")
+ appendValue("second")
+
+ page := file.NewDictionaryPage(memory.NewBufferBytes(data), 2,
parquet.Encodings.Plain)
+ matches, known := dictionaryMatchesPredicates(
+ page,
+ parquet.Types.ByteArray,
+ -1,
+ []RowGroupDictionaryPred{{FieldID: 1, PhysBytes:
[][]byte{[]byte("missing"), []byte("second")}}},
+ []int{0},
+ )
+ page.Release()
+
+ require.True(t, known)
+ assert.True(t, matches)
+}
+
+func TestDictionaryMatchesPredicatesUsesLocalPredicateIndexes(t *testing.T) {
+ data := make([]byte, 0, 16)
+ var length [4]byte
+ binary.LittleEndian.PutUint32(length[:], uint32(len("second")))
+ data = append(data, length[:]...)
+ data = append(data, "second"...)
+
+ page := file.NewDictionaryPage(memory.NewBufferBytes(data), 1,
parquet.Encodings.Plain)
+ matches, known := dictionaryMatchesPredicates(
+ page,
+ parquet.Types.ByteArray,
+ -1,
+ []RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes: [][]byte{[]byte("missing")}},
+ {FieldID: 1, PhysBytes: [][]byte{[]byte("second")}},
+ },
+ []int{1},
+ )
+ page.Release()
+
+ require.True(t, known)
+ assert.True(t, matches)
+}
+
+func TestDictionaryMatchesPredicatesRejectsMalformedPage(t *testing.T) {
+ data := []byte{5, 0, 0, 0, 'x'}
+ page := file.NewDictionaryPage(memory.NewBufferBytes(data), 1,
parquet.Encodings.Plain)
+ matches, known := dictionaryMatchesPredicates(
+ page,
+ parquet.Types.ByteArray,
+ -1,
+ []RowGroupDictionaryPred{{FieldID: 1, PhysBytes:
[][]byte{[]byte("x")}}},
+ []int{0},
+ )
+ page.Release()
+
+ assert.False(t, known)
+ assert.False(t, matches)
+}
+
+func TestDictionaryMatchesPredicatesPrunesEmptyDictionary(t *testing.T) {
+ page := file.NewDictionaryPage(memory.NewBufferBytes(nil), 0,
parquet.Encodings.Plain)
+ matches, known := dictionaryMatchesPredicates(
+ page,
+ parquet.Types.Int32,
+ -1,
+ []RowGroupDictionaryPred{{FieldID: 1, PhysBytes:
[][]byte{int32BytesForTest(1)}}},
+ []int{0},
+ )
+ page.Release()
+
+ assert.True(t, known)
+ assert.False(t, matches)
+}
+
+func TestDictionaryMatchesPredicatesRejectsZeroCountWithData(t *testing.T) {
+ page := file.NewDictionaryPage(
+ memory.NewBufferBytes(int32BytesForTest(1)), 0,
parquet.Encodings.Plain)
+ matches, known := dictionaryMatchesPredicates(
+ page,
+ parquet.Types.Int32,
+ -1,
+ []RowGroupDictionaryPred{{FieldID: 1, PhysBytes:
[][]byte{int32BytesForTest(1)}}},
+ []int{0},
+ )
+ page.Release()
+
+ assert.False(t, known)
+ assert.False(t, matches)
+}
+
+func TestDictionaryMatchesPredicatesTreatsSignedZeroAsEqual(t *testing.T) {
+ data := make([]byte, 4)
+ binary.LittleEndian.PutUint32(data,
math.Float32bits(float32(math.Copysign(0, -1))))
+ candidate := make([]byte, 4)
+ binary.LittleEndian.PutUint32(candidate, math.Float32bits(0))
+
+ page := file.NewDictionaryPage(memory.NewBufferBytes(data), 1,
parquet.Encodings.Plain)
+ matches, known := dictionaryMatchesPredicates(
+ page,
+ parquet.Types.Float,
+ -1,
+ []RowGroupDictionaryPred{{FieldID: 1, PhysBytes:
[][]byte{candidate}}},
+ []int{0},
+ )
+ page.Release()
+
+ require.True(t, known)
+ assert.True(t, matches)
+}
+
+func appendFloat32DictionaryValue(data []byte, value float32) []byte {
+ var encoded [4]byte
+ binary.LittleEndian.PutUint32(encoded[:], math.Float32bits(value))
+
+ return append(data, encoded[:]...)
+}
+
+func appendFloat64DictionaryValue(data []byte, value float64) []byte {
+ var encoded [8]byte
+ binary.LittleEndian.PutUint64(encoded[:], math.Float64bits(value))
+
+ return append(data, encoded[:]...)
+}
+
+func float32DictionaryBytes(values ...float32) []byte {
+ data := make([]byte, 0, len(values)*4)
+ for _, value := range values {
+ data = appendFloat32DictionaryValue(data, value)
+ }
+
+ return data
+}
+
+func float64DictionaryBytes(values ...float64) []byte {
+ data := make([]byte, 0, len(values)*8)
+ for _, value := range values {
+ data = appendFloat64DictionaryValue(data, value)
+ }
+
+ return data
+}
+
+func TestDictionaryMatchesPredicatesHandlesNaN(t *testing.T) {
+ tests := []struct {
+ name string
+ physicalType parquet.Type
+ data []byte
+ numValues int32
+ candidate []byte
+ wantMatch bool
+ }{
+ {
+ name: "Float dictionary NaN does not match a
number",
+ physicalType: parquet.Types.Float,
+ data:
float32DictionaryBytes(float32(math.NaN()), 1, 2),
+ numValues: 3,
+ candidate: float32BytesForTest(99),
+ wantMatch: false,
+ },
+ {
+ name: "Double dictionary NaN does not match a
number",
+ physicalType: parquet.Types.Double,
+ data: float64DictionaryBytes(math.NaN(), 1, 2),
+ numValues: 3,
+ candidate: float64BytesForTest(99),
+ wantMatch: false,
+ },
+ {
+ name: "Float NaN candidate keeps the group",
+ physicalType: parquet.Types.Float,
+ data: float32DictionaryBytes(1),
+ numValues: 1,
+ candidate: appendFloat32DictionaryValue(nil,
float32(math.NaN())),
+ wantMatch: true,
+ },
+ {
+ name: "Double NaN candidate keeps the group",
+ physicalType: parquet.Types.Double,
+ data: float64DictionaryBytes(1),
+ numValues: 1,
+ candidate: appendFloat64DictionaryValue(nil,
math.NaN()),
+ wantMatch: true,
+ },
+ }
+
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ page :=
file.NewDictionaryPage(memory.NewBufferBytes(test.data), test.numValues,
parquet.Encodings.Plain)
+
+ matches, known := dictionaryMatchesPredicates(
+ page,
+ test.physicalType,
+ -1,
+ []RowGroupDictionaryPred{{FieldID: 1,
PhysBytes: [][]byte{test.candidate}}},
+ []int{0},
+ )
+ page.Release()
+
+ require.True(t, known)
+ assert.Equal(t, test.wantMatch, matches)
+ })
+ }
+}
+
+func TestDictionaryMatchesPredicatesUsesCandidateSet(t *testing.T) {
+ const numValues = 32
+ data := make([]byte, numValues*4)
+ for i := range numValues {
+ binary.LittleEndian.PutUint32(data[i*4:], uint32(i))
+ }
+
+ candidates := make([][]byte, 8)
+ for i := range candidates {
+ candidates[i] = int32BytesForTest(int32(1000 + i))
+ }
+ candidates[len(candidates)-1] = int32BytesForTest(17)
+
+ page := file.NewDictionaryPage(memory.NewBufferBytes(data), numValues,
parquet.Encodings.Plain)
+ matches, known := dictionaryMatchesPredicates(
+ page,
+ parquet.Types.Int32,
+ -1,
+ []RowGroupDictionaryPred{{FieldID: 1, PhysBytes: candidates}},
+ []int{0},
+ )
+ page.Release()
+
+ assert.True(t, known)
+ assert.True(t, matches)
+}
+
+func BenchmarkDictionaryMatchesPredicatesLargeIn(b *testing.B) {
+ const (
+ numValues = 4096
+ numCandidates = 200
+ physicalWidth = 4
+ )
+
+ data := make([]byte, numValues*physicalWidth)
+ for i := range numValues {
+ binary.LittleEndian.PutUint32(data[i*physicalWidth:], uint32(i))
+ }
+
+ candidates := make([][]byte, numCandidates)
+ for i := range candidates {
+ candidates[i] = int32BytesForTest(int32(-i - 1))
+ }
+
+ page := file.NewDictionaryPage(
+ memory.NewBufferBytes(data), numValues, parquet.Encodings.Plain)
+ defer page.Release()
+ preds := []RowGroupDictionaryPred{{FieldID: 1, PhysBytes: candidates}}
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for range b.N {
+ matches, known := dictionaryMatchesPredicates(
+ page, parquet.Types.Int32, -1, preds, []int{0})
+ if matches || !known {
+ b.Fatal("large IN dictionary match returned an
unexpected result")
+ }
+ }
+}
+
+func float32BytesForTest(value float32) []byte {
+ var encoded [4]byte
+ binary.LittleEndian.PutUint32(encoded[:], math.Float32bits(value))
+
+ return encoded[:]
+}
+
+func int32BytesForTest(value int32) []byte {
+ var encoded [4]byte
+ binary.LittleEndian.PutUint32(encoded[:], uint32(value))
+
+ return encoded[:]
+}
+
+func float64BytesForTest(value float64) []byte {
+ var encoded [8]byte
+ binary.LittleEndian.PutUint64(encoded[:], math.Float64bits(value))
+
+ return encoded[:]
+}
diff --git a/table/internal/parquet_files.go b/table/internal/parquet_files.go
index 94f46e968..c398bd9c7 100644
--- a/table/internal/parquet_files.go
+++ b/table/internal/parquet_files.go
@@ -19,12 +19,14 @@ package internal
import (
"context"
+ "encoding/binary"
"errors"
"fmt"
"io"
"io/fs"
"log/slog"
"maps"
+ "math"
"os"
"slices"
"strconv"
@@ -127,7 +129,11 @@ func (parquetFormat) Open(ctx context.Context, fs
iceio.IO, path string) (_ File
return nil, err
}
- return wrapPqArrowReader{FileReader: arrRdr}, nil
+ return wrapPqArrowReader{
+ FileReader: arrRdr,
+ dictionarySource: inputfile,
+ dictionaryMem: alloc,
+ }, nil
}
func (parquetFormat) PathToIDMapping(sc *iceberg.Schema) (map[string]int,
error) {
@@ -1892,11 +1898,23 @@ type RowGroupBloomPred struct {
PhysBytes [][]byte // one entry for EqualTo; one per value for In
}
-// ParquetRowGroupTester combines stats-based and bloom filter row group
pruning.
+// RowGroupDictionaryPred holds the physical-encoded bytes for each literal in
+// a dictionary-prunable predicate on one field. A row group can be skipped
+// when NONE of the bytes occur in its complete dictionary. It intentionally
+// mirrors RowGroupBloomPred because both predicates use the same physical
+// literal encoding.
+type RowGroupDictionaryPred struct {
+ FieldID int
+ PhysBytes [][]byte // one entry for EqualTo; one per value for In
+}
+
+// ParquetRowGroupTester combines stats-based, dictionary, and bloom filter row
+// group pruning.
// Pass it as the tester argument to wrapPqArrowReader.GetRecords.
type ParquetRowGroupTester struct {
- StatsFn func(*metadata.RowGroupMetaData, []int) (bool, error)
- BloomPreds []RowGroupBloomPred // nil = no bloom filter pass
+ StatsFn func(*metadata.RowGroupMetaData, []int) (bool, error)
+ BloomPreds []RowGroupBloomPred // nil = no bloom filter pass
+ DictionaryPreds []RowGroupDictionaryPred // nil = no dictionary filter
pass
// RangeSet indicates that Start and Length came from an explicit scan
task.
// Without it, the zero value keeps the historical full-file behavior
for
// callers that construct a tester only for row-group pruning.
@@ -1926,6 +1944,10 @@ type RowGroupSpan struct {
type wrapPqArrowReader struct {
*pqarrow.FileReader
rowGroupInfos []parquetRowGroupInfo
+ // dictionarySource is shared with FileReader. The temporary reader
created
+ // from it wraps the source so it cannot close the main reader's input.
+ dictionarySource parquet.ReaderAtSeeker
+ dictionaryMem memory.Allocator
}
func (w wrapPqArrowReader) Metadata() Metadata {
@@ -2025,7 +2047,8 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
var rgList []int
if rowGroupTester != nil && (rangeSet ||
- rowGroupTester.StatsFn != nil || len(rowGroupTester.BloomPreds)
> 0) {
+ rowGroupTester.StatsFn != nil ||
len(rowGroupTester.DictionaryPreds) > 0 ||
+ len(rowGroupTester.BloomPreds) > 0) {
fileMeta := w.ParquetReader().MetaData()
numRg := w.ParquetReader().NumRowGroups()
planningOffsets := rowGroupTester.PlanningSplitOffsets
@@ -2034,12 +2057,35 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
}
var (
- fieldIDToColIdx map[int]int
- bfReader *metadata.BloomFilterReader
+ fieldIDToColIdx map[int]int
+ bfReader *metadata.BloomFilterReader
+ dictionaryReader *file.Reader
)
- if len(rowGroupTester.BloomPreds) > 0 {
+ if len(rowGroupTester.BloomPreds) > 0 ||
len(rowGroupTester.DictionaryPreds) > 0 {
fieldIDToColIdx = buildFieldIDToColIdx(fileMeta)
+ }
+ var dictionaryPredsByColumn map[int][]int
+ if len(rowGroupTester.DictionaryPreds) > 0 {
+ dictionaryPredsByColumn =
groupRowGroupDictionaryPredicates(
+ fieldIDToColIdx, rowGroupTester.DictionaryPreds)
+ if len(dictionaryPredsByColumn) > 0 &&
w.dictionarySource != nil {
+ // Dictionary inspection is an optional
optimisation. If the
+ // section-read reader cannot be created, use
the main reader and
+ // keep the existing conservative behaviour.
+ dictionaryReader, _ =
newParquetDictionaryReader(
+ w.dictionarySource, fileMeta,
w.dictionaryMem)
+ if dictionaryReader != nil {
+ // The temporary reader owns only its
buffers. Its source is
+ // wrapped so this cannot close the
main reader's input.
+ defer dictionaryReader.Close()
+ }
+ }
+ if len(dictionaryPredsByColumn) > 0 && dictionaryReader
== nil {
+ dictionaryReader = w.ParquetReader()
+ }
+ }
+ if len(rowGroupTester.BloomPreds) > 0 {
bfReader = w.ParquetReader().GetBloomFilterReader()
}
@@ -2088,6 +2134,15 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
}
}
+ if use && len(dictionaryPredsByColumn) > 0 {
+ if rgMeta == nil {
+ rgMeta = fileMeta.RowGroup(rg)
+ }
+ use = checkRowGroupDictionaries(
+ dictionaryReader, fileMeta, rgMeta, rg,
+ dictionaryPredsByColumn,
rowGroupTester.DictionaryPreds)
+ }
+
if use && bfReader != nil {
var err error
use, err = checkRowGroupBloomFilters(bfReader,
rg, fieldIDToColIdx, rowGroupTester.BloomPreds)
@@ -2117,6 +2172,389 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
return w.GetRecordReader(ctx, cols, rgList)
}
+// nonClosingReaderAtSeeker gives a temporary Parquet reader access to a shared
+// source without allowing that reader's Close method to close the source owned
+// by the main Arrow reader.
+type nonClosingReaderAtSeeker struct {
+ parquet.ReaderAtSeeker
+}
+
+func (nonClosingReaderAtSeeker) Close() error {
+ return nil
+}
+
+func newParquetDictionaryReader(
+ source parquet.ReaderAtSeeker, fileMeta *metadata.FileMetaData, mem
memory.Allocator,
+) (*file.Reader, error) {
+ readProps := parquet.NewReaderProperties(mem)
+ readProps.BufferedStreamEnabled = true
+
+ return file.NewParquetReader(nonClosingReaderAtSeeker{ReaderAtSeeker:
source},
+ file.WithMetadata(fileMeta), file.WithReadProps(readProps))
+}
+
+func groupRowGroupDictionaryPredicates(
+ fieldIDToColIdx map[int]int,
+ preds []RowGroupDictionaryPred,
+) map[int][]int {
+ predsByColumn := make(map[int][]int)
+ for i, pred := range preds {
+ if len(pred.PhysBytes) == 0 {
+ continue
+ }
+
+ colIdx, ok := fieldIDToColIdx[pred.FieldID]
+ if ok {
+ predsByColumn[colIdx] = append(predsByColumn[colIdx], i)
+ }
+ }
+
+ return predsByColumn
+}
+
+// checkRowGroupDictionaries checks each dictionary predicate against the
+// complete dictionary for its column in row group rg. It returns false only
+// when the column metadata proves that every data page is dictionary encoded
+// and none of the predicate values occur in the dictionary. Missing encoding
+// metadata, dictionary pages, unsupported physical types, malformed pages, or
+// reader errors keep the row group so this optimisation cannot drop data.
+func checkRowGroupDictionaries(
+ rdr *file.Reader,
+ fileMeta *metadata.FileMetaData,
+ rgMeta *metadata.RowGroupMetaData,
+ rg int,
+ predsByColumn map[int][]int,
+ preds []RowGroupDictionaryPred,
+) bool {
+ if len(preds) == 0 || len(predsByColumn) == 0 {
+ return true
+ }
+
+ rgReader := rdr.RowGroup(rg)
+ for colIdx, predIndexes := range predsByColumn {
+ chunk, err := rgMeta.ColumnChunk(colIdx)
+ if err != nil || chunk == nil ||
!parquetColumnUsesOnlyDictionaryData(chunk) {
+ continue
+ }
+
+ column := fileMeta.Schema.Column(colIdx)
+ pageRdr, err := rgReader.GetColumnPageReader(colIdx)
+ if err != nil {
+ continue
+ }
+
+ dictPage, dictErr := pageRdr.GetDictionaryPage()
+ if dictErr != nil || dictPage == nil {
+ if dictPage != nil {
+ dictPage.Release()
+ }
+ _ = pageRdr.Close()
+
+ continue
+ }
+
+ matches, known := dictionaryMatchesPredicates(
+ dictPage, column.PhysicalType(), column.TypeLength(),
preds, predIndexes)
+ dictPage.Release()
+ closeErr := pageRdr.Close()
+ if !known {
+ continue
+ }
+ if closeErr != nil {
+ // A close error makes an otherwise conclusive
dictionary result
+ // unusable, so keep the row group conservatively.
+ slog.Warn("dictionary row-group pruning skipped after
page reader close error",
+ "rowGroup", rg, "err", closeErr)
+
+ continue
+ }
+
+ if !matches {
+ return false
+ }
+ }
+
+ return true
+}
+
+// parquetColumnUsesOnlyDictionaryData verifies that the optional encoding
+// statistics identify a dictionary page and dictionary-encoded data pages,
+// with no PLAIN fallback pages. EncodingStats is deliberately required: the
+// column encoding list cannot distinguish the PLAIN dictionary page itself
+// from a later PLAIN data page in older files.
+func parquetColumnUsesOnlyDictionaryData(chunk *metadata.ColumnChunkMetaData)
bool {
+ if !chunk.HasDictionaryPage() {
+ return false
+ }
+
+ stats := chunk.EncodingStats()
+ if len(stats) == 0 {
+ return false
+ }
+
+ var hasDictionaryPage, hasDictionaryData bool
+ for _, stat := range stats {
+ switch stat.PageType {
+ case file.PageTypeDictionaryPage:
+ if stat.Encoding != parquet.Encodings.Plain &&
stat.Encoding != parquet.Encodings.PlainDict {
+ return false
+ }
+ hasDictionaryPage = true
+ case file.PageTypeDataPage, file.PageTypeDataPageV2:
+ if stat.Encoding != parquet.Encodings.RLEDict &&
stat.Encoding != parquet.Encodings.PlainDict {
+ return false
+ }
+ hasDictionaryData = true
+ default:
+ return false
+ }
+ }
+
+ return hasDictionaryPage && hasDictionaryData
+}
+
+// dictionaryMatchesPredicates decodes a PLAIN dictionary page and reports
+// whether every requested predicate has at least one matching value. The
+// returned known flag is false for unsupported or malformed input that
+// prevents a pruning decision, which means the caller must retain the row
+// group.
+func dictionaryMatchesPredicates(
+ page *file.DictionaryPage,
+ physicalType parquet.Type,
+ typeLen int,
+ preds []RowGroupDictionaryPred,
+ predIndexes []int,
+) (bool, bool) {
+ pageEncoding := parquet.Encoding(page.Encoding())
+ if pageEncoding != parquet.Encodings.Plain && pageEncoding !=
parquet.Encodings.PlainDict {
+ return false, false
+ }
+
+ width, variableWidth, ok := parquetDictionaryValueLayout(physicalType,
typeLen)
+ if !ok {
+ return false, false
+ }
+
+ for _, predIndex := range predIndexes {
+ usable := false
+ for _, candidate := range preds[predIndex].PhysBytes {
+ if variableWidth || len(candidate) == width {
+ usable = true
+
+ break
+ }
+ }
+ if !usable {
+ return false, false
+ }
+ }
+
+ data := page.Data()
+ offset := 0
+ numValues := page.NumValues()
+ if numValues < 0 {
+ return false, false
+ }
+ if numValues == 0 {
+ return false, len(data) == 0
+ }
+
+ if len(predIndexes) == 1 {
+ pred := preds[predIndexes[0]]
+ candidateSet := parquetDictionaryCandidateSet(
+ physicalType, width, variableWidth, numValues,
pred.PhysBytes)
+ for range numValues {
+ value, next, ok := nextParquetDictionaryValue(data,
offset, physicalType, width)
+ if !ok {
+ return false, false
+ }
+ offset = next
+
+ if candidateSet != nil {
+ if _, ok := candidateSet[string(value)]; ok {
+ return true, true
+ }
+ } else {
+ for _, candidate := range pred.PhysBytes {
+ if
parquetDictionaryValueEqual(physicalType, value, candidate) {
+ return true, true
+ }
+ }
+ }
+ }
+
+ if offset != len(data) {
+ return false, false
+ }
+
+ return false, true
+ }
+
+ // predIndexes contains indexes into the full preds slice. Keep matches
+ // aligned with predIndexes so columns with a small predicate subset do
not
+ // allocate a slice for predicates belonging to other columns.
+ matches := make([]bool, len(predIndexes))
+ candidateSets := make([]map[string]struct{}, len(predIndexes))
+ for matchIndex, predIndex := range predIndexes {
+ candidateSets[matchIndex] = parquetDictionaryCandidateSet(
+ physicalType, width, variableWidth, numValues,
preds[predIndex].PhysBytes)
+ }
+ remaining := len(predIndexes)
+ for range numValues {
+ if remaining == 0 {
+ break
+ }
+
+ value, next, ok := nextParquetDictionaryValue(data, offset,
physicalType, width)
+ if !ok {
+ return false, false
+ }
+ offset = next
+
+ for matchIndex, predIndex := range predIndexes {
+ if matches[matchIndex] {
+ continue
+ }
+
+ if candidateSets[matchIndex] != nil {
+ if _, ok :=
candidateSets[matchIndex][string(value)]; ok {
+ matches[matchIndex] = true
+ remaining--
+ }
+ } else {
+ for _, candidate := range
preds[predIndex].PhysBytes {
+ if
parquetDictionaryValueEqual(physicalType, value, candidate) {
+ matches[matchIndex] = true
+ remaining--
+
+ break
+ }
+ }
+ }
+ }
+ }
+
+ if remaining == 0 {
+ // All predicates already match, so trailing dictionary entries
cannot
+ // change the decision and do not need to be decoded or
validated.
+ return true, true
+ }
+
+ if offset != len(data) {
+ return false, false
+ }
+
+ return false, true
+}
+
+const (
+ parquetDictionaryCandidateSetLimit = 8
+ parquetDictionaryCandidateSetMinValues = 16
+)
+
+func parquetDictionaryCandidateSet(
+ physicalType parquet.Type,
+ width int,
+ variableWidth bool,
+ numValues int32,
+ candidates [][]byte,
+) map[string]struct{} {
+ if len(candidates) < parquetDictionaryCandidateSetLimit ||
+ numValues < parquetDictionaryCandidateSetMinValues {
+ return nil
+ }
+
+ switch physicalType {
+ case parquet.Types.Int32, parquet.Types.Int64,
+ parquet.Types.ByteArray, parquet.Types.FixedLenByteArray:
+ default:
+ // Float and Double need numeric comparison for signed zero and
NaN.
+ return nil
+ }
+
+ result := make(map[string]struct{}, len(candidates))
+ for _, candidate := range candidates {
+ if variableWidth || len(candidate) == width {
+ result[string(candidate)] = struct{}{}
+ }
+ }
+
+ return result
+}
+
+func parquetDictionaryValueLayout(physicalType parquet.Type, typeLen int)
(width int, variableWidth, ok bool) {
+ switch physicalType {
+ case parquet.Types.Int32, parquet.Types.Float:
+ return 4, false, true
+ case parquet.Types.Int64, parquet.Types.Double:
+ return 8, false, true
+ case parquet.Types.ByteArray:
+ return 0, true, true
+ case parquet.Types.FixedLenByteArray:
+ if typeLen > 0 {
+ return typeLen, false, true
+ }
+ }
+
+ return 0, false, false
+}
+
+func nextParquetDictionaryValue(data []byte, offset int, physicalType
parquet.Type, width int) ([]byte, int, bool) {
+ if physicalType == parquet.Types.ByteArray {
+ if offset < 0 || offset > len(data) || len(data)-offset < 4 {
+ return nil, 0, false
+ }
+
+ valueLen := int(binary.LittleEndian.Uint32(data[offset:]))
+ valueStart := offset + 4
+ if valueLen > len(data)-valueStart {
+ return nil, 0, false
+ }
+
+ return data[valueStart : valueStart+valueLen], valueStart +
valueLen, true
+ }
+
+ if offset < 0 || width < 0 || width > len(data)-offset {
+ return nil, 0, false
+ }
+
+ return data[offset : offset+width], offset + width, true
+}
+
+func parquetDictionaryValueEqual(physicalType parquet.Type, value, candidate
[]byte) bool {
+ if physicalType == parquet.Types.Float && len(value) == 4 &&
len(candidate) == 4 {
+ left := math.Float32frombits(binary.LittleEndian.Uint32(value))
+ right :=
math.Float32frombits(binary.LittleEndian.Uint32(candidate))
+ if math.IsNaN(float64(right)) {
+ // Keep groups for NaN literals because reader equality
semantics
+ // can differ.
+ return true
+ }
+ if math.IsNaN(float64(left)) {
+ return false
+ }
+
+ return left == right
+ }
+
+ if physicalType == parquet.Types.Double && len(value) == 8 &&
len(candidate) == 8 {
+ left := math.Float64frombits(binary.LittleEndian.Uint64(value))
+ right :=
math.Float64frombits(binary.LittleEndian.Uint64(candidate))
+ if math.IsNaN(right) {
+ // Keep groups for NaN literals because reader equality
semantics
+ // can differ.
+ return true
+ }
+ if math.IsNaN(left) {
+ return false
+ }
+
+ return left == right
+ }
+
+ return slices.Equal(value, candidate)
+}
+
// buildFieldIDToColIdx maps each Iceberg field ID to its 0-based column index
in
// the Parquet file schema. Used to translate bloom filter predicates (which
carry
// field IDs) into the column positions required by BloomFilterReader.
@@ -2248,7 +2686,12 @@ func (pfs *ParquetFileSource) GetReader(ctx
context.Context) (result FileReader,
return nil, err
}
- result = wrapPqArrowReader{FileReader: fr, rowGroupInfos: rowGroupInfos}
+ result = wrapPqArrowReader{
+ FileReader: fr,
+ rowGroupInfos: rowGroupInfos,
+ dictionarySource: pf,
+ dictionaryMem: pfs.mem,
+ }
return result, nil
}
diff --git a/table/internal/parquet_files_test.go
b/table/internal/parquet_files_test.go
index 7a30bc06a..663beeeaf 100644
--- a/table/internal/parquet_files_test.go
+++ b/table/internal/parquet_files_test.go
@@ -2022,7 +2022,7 @@ func TestParquetBatchSizeFromTableProperties(t
*testing.T) {
// buildBloomTestParquet writes a Parquet file with two row groups into buf.
// The single required INT32 column "id" has Iceberg field_id=1 and bloom
// filters enabled. RG0 holds values [1..rgSize], RG1 holds
[rgSize+1..2*rgSize].
-func buildBloomTestParquet(t *testing.T, rgSize int) []byte {
+func buildBloomTestParquet(t testing.TB, rgSize int) []byte {
t.Helper()
idNode := schema.NewInt32Node("id", parquet.Repetitions.Required, 1)
@@ -2066,7 +2066,20 @@ func buildBloomTestParquet(t *testing.T, rgSize int)
[]byte {
// openBloomTestReader creates a fresh FileReader from raw Parquet bytes for
// each subtest call, so row-group state is not shared between subtests.
-func openBloomTestReader(t *testing.T, data []byte) internal.FileReader {
+func openBloomTestReader(t testing.TB, data []byte) internal.FileReader {
+ t.Helper()
+
+ source := bytes.NewReader(data)
+ pqRdr, err := file.NewParquetReader(source)
+ require.NoError(t, err)
+
+ arrRdr, err := pqarrow.NewFileReader(pqRdr,
pqarrow.ArrowReadProperties{}, memory.DefaultAllocator)
+ require.NoError(t, err)
+
+ return internal.WrapParquetFileReaderWithDictionarySource(arrRdr,
source)
+}
+
+func openTestReaderWithoutDictionarySource(t testing.TB, data []byte)
internal.FileReader {
t.Helper()
pqRdr, err := file.NewParquetReader(bytes.NewReader(data))
@@ -2085,7 +2098,7 @@ func int32PhysBytes(v int32) []byte {
return b
}
-func countRecords(t *testing.T, rr array.RecordReader) int64 {
+func countRecords(t testing.TB, rr array.RecordReader) int64 {
t.Helper()
defer rr.Release()
@@ -2213,6 +2226,518 @@ func TestBloomFilterRowGroupPruning(t *testing.T) {
})
}
+// buildDictionaryTestParquet writes a Parquet file with two low-cardinality
+// dictionary-encoded row groups. The single required INT32 column "id" has
+// Iceberg field_id=1.
+func buildDictionaryTestParquet(t testing.TB, rowGroups ...[]int32) []byte {
+ t.Helper()
+
+ idNode := schema.NewInt32Node("id", parquet.Repetitions.Required, 1)
+ rootNode, err := schema.NewGroupNode("schema",
parquet.Repetitions.Required,
+ schema.FieldList{idNode}, -1)
+ require.NoError(t, err)
+
+ writerProps := parquet.NewWriterProperties(
+ parquet.WithStats(true),
+ parquet.WithDictionaryDefault(true),
+ )
+
+ var buf bytes.Buffer
+ pw := file.NewParquetWriter(&buf, rootNode,
file.WithWriterProps(writerProps))
+ for _, values := range rowGroups {
+ rgw, werr := pw.AppendRowGroupChecked()
+ require.NoError(t, werr)
+ cw, werr := rgw.NextColumn()
+ require.NoError(t, werr)
+ _, werr = cw.(*file.Int32ColumnChunkWriter).WriteBatch(values,
nil, nil)
+ require.NoError(t, werr)
+ require.NoError(t, cw.Close())
+ require.NoError(t, rgw.Close())
+ }
+ require.NoError(t, pw.Close())
+
+ return buf.Bytes()
+}
+
+type multiColumnDictionaryTestRowGroup struct {
+ ids []int32
+ categories []string
+}
+
+func buildMultiColumnDictionaryTestParquet(
+ t testing.TB, rowGroups ...multiColumnDictionaryTestRowGroup,
+) []byte {
+ t.Helper()
+
+ idNode := schema.NewInt32Node("id", parquet.Repetitions.Required, 1)
+ categoryNode := schema.NewByteArrayNode("category",
parquet.Repetitions.Required, 2)
+ rootNode, err := schema.NewGroupNode("schema",
parquet.Repetitions.Required,
+ schema.FieldList{idNode, categoryNode}, -1)
+ require.NoError(t, err)
+
+ writerProps := parquet.NewWriterProperties(
+ parquet.WithStats(true),
+ parquet.WithDictionaryDefault(true),
+ )
+
+ var buf bytes.Buffer
+ pw := file.NewParquetWriter(&buf, rootNode,
file.WithWriterProps(writerProps))
+ for _, values := range rowGroups {
+ require.Len(t, values.categories, len(values.ids))
+
+ rgw, werr := pw.AppendRowGroupChecked()
+ require.NoError(t, werr)
+
+ idColumn, werr := rgw.NextColumn()
+ require.NoError(t, werr)
+ _, werr =
idColumn.(*file.Int32ColumnChunkWriter).WriteBatch(values.ids, nil, nil)
+ require.NoError(t, werr)
+ require.NoError(t, idColumn.Close())
+
+ categoryColumn, werr := rgw.NextColumn()
+ require.NoError(t, werr)
+ categories := make([]parquet.ByteArray, len(values.categories))
+ for i, category := range values.categories {
+ categories[i] = parquet.ByteArray(category)
+ }
+ _, werr =
categoryColumn.(*file.ByteArrayColumnChunkWriter).WriteBatch(categories, nil,
nil)
+ require.NoError(t, werr)
+ require.NoError(t, categoryColumn.Close())
+ require.NoError(t, rgw.Close())
+ }
+ require.NoError(t, pw.Close())
+
+ return buf.Bytes()
+}
+
+// TestDictionaryRowGroupPruning verifies that dictionary-only row groups are
+// skipped when an EqualTo/In predicate has no value in the complete
dictionary.
+func TestDictionaryRowGroupPruning(t *testing.T) {
+ const rgSize = 1024
+ rg0 := make([]int32, rgSize)
+ rg1 := make([]int32, rgSize)
+ for i := range rg0 {
+ rg0[i] = 1 + int32(i%2)*2
+ rg1[i] = 5 + int32(i%2)*2
+ }
+
+ data := buildDictionaryTestParquet(t, rg0, rg1)
+ alwaysKeep := func(_ *metadata.RowGroupMetaData, _ []int) (bool, error)
{
+ return true, nil
+ }
+
+ ctx := context.Background()
+ cols := []int{0}
+
+ t.Run("EqualTo absent from both dictionaries", func(t *testing.T) {
+ rdr := openBloomTestReader(t, data)
+ var survivors []internal.RowGroupSpan
+ tester := &internal.ParquetRowGroupTester{
+ DictionaryPreds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes:
[][]byte{int32PhysBytes(2)}},
+ },
+ Survivors: &survivors,
+ }
+ rr, err := rdr.GetRecords(ctx, cols, tester)
+ require.NoError(t, err)
+
+ assert.Equal(t, int64(0), countRecords(t, rr))
+ assert.Empty(t, survivors)
+ require.NoError(t, rdr.Close())
+ })
+
+ t.Run("EqualTo keeps the matching dictionary", func(t *testing.T) {
+ rdr := openBloomTestReader(t, data)
+ var survivors []internal.RowGroupSpan
+ tester := &internal.ParquetRowGroupTester{
+ DictionaryPreds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes:
[][]byte{int32PhysBytes(7)}},
+ },
+ StatsFn: alwaysKeep,
+ Survivors: &survivors,
+ }
+ rr, err := rdr.GetRecords(ctx, cols, tester)
+ require.NoError(t, err)
+
+ assert.Equal(t, int64(rgSize), countRecords(t, rr))
+ assert.Equal(t, []internal.RowGroupSpan{{FirstRowPos: rgSize,
NumRows: rgSize}}, survivors)
+ require.NoError(t, rdr.Close())
+ })
+
+ t.Run("In keeps a row group when any value is present", func(t
*testing.T) {
+ rdr := openBloomTestReader(t, data)
+ var survivors []internal.RowGroupSpan
+ tester := &internal.ParquetRowGroupTester{
+ DictionaryPreds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes:
[][]byte{int32PhysBytes(2), int32PhysBytes(5)}},
+ },
+ StatsFn: alwaysKeep,
+ Survivors: &survivors,
+ }
+ rr, err := rdr.GetRecords(ctx, cols, tester)
+ require.NoError(t, err)
+
+ assert.Equal(t, int64(rgSize), countRecords(t, rr))
+ assert.Equal(t, []internal.RowGroupSpan{{FirstRowPos: rgSize,
NumRows: rgSize}}, survivors)
+ require.NoError(t, rdr.Close())
+ })
+}
+
+func TestDictionaryRowGroupPruningMultiColumnAnd(t *testing.T) {
+ const rgSize = 1024
+
+ alwaysKeep := func(_ *metadata.RowGroupMetaData, _ []int) (bool, error)
{
+ return true, nil
+ }
+ repeatedIDs := func(first, second int32) []int32 {
+ values := make([]int32, rgSize)
+ for i := range values {
+ if i%2 == 0 {
+ values[i] = first
+ } else {
+ values[i] = second
+ }
+ }
+
+ return values
+ }
+ repeatedCategory := func(category string) []string {
+ values := make([]string, rgSize)
+ for i := range values {
+ values[i] = category
+ }
+
+ return values
+ }
+
+ tests := []struct {
+ name string
+ rowGroups []multiColumnDictionaryTestRowGroup
+ preds []internal.RowGroupDictionaryPred
+ wantRows int64
+ survivors []internal.RowGroupSpan
+ }{
+ {
+ name: "id predicate prunes a group while category
matches",
+ rowGroups: []multiColumnDictionaryTestRowGroup{
+ {ids: repeatedIDs(1, 3), categories:
repeatedCategory("a")},
+ {ids: repeatedIDs(5, 7), categories:
repeatedCategory("a")},
+ },
+ preds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes:
[][]byte{int32PhysBytes(5)}},
+ {FieldID: 2, PhysBytes: [][]byte{[]byte("a")}},
+ },
+ wantRows: rgSize,
+ survivors: []internal.RowGroupSpan{{FirstRowPos:
rgSize, NumRows: rgSize}},
+ },
+ {
+ name: "category predicate prunes a group while id
matches",
+ rowGroups: []multiColumnDictionaryTestRowGroup{
+ {ids: repeatedIDs(1, 3), categories:
repeatedCategory("a")},
+ {ids: repeatedIDs(1, 3), categories:
repeatedCategory("b")},
+ },
+ preds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes:
[][]byte{int32PhysBytes(1)}},
+ {FieldID: 2, PhysBytes: [][]byte{[]byte("a")}},
+ },
+ wantRows: rgSize,
+ survivors: []internal.RowGroupSpan{{FirstRowPos: 0,
NumRows: rgSize}},
+ },
+ }
+
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ data := buildMultiColumnDictionaryTestParquet(t,
test.rowGroups...)
+ rdr := openBloomTestReader(t, data)
+ defer rdr.Close()
+
+ var survivors []internal.RowGroupSpan
+ tester := &internal.ParquetRowGroupTester{
+ StatsFn: alwaysKeep,
+ DictionaryPreds: test.preds,
+ Survivors: &survivors,
+ }
+ rr, err := rdr.GetRecords(context.Background(),
[]int{0, 1}, tester)
+ require.NoError(t, err)
+
+ assert.Equal(t, test.wantRows, countRecords(t, rr))
+ assert.Equal(t, test.survivors, survivors)
+ })
+ }
+}
+
+func TestDictionaryRowGroupPruningChecksTailAfterMatchingPrefix(t *testing.T) {
+ const (
+ rgSize = 1024
+ matchingGroups = 16
+ )
+
+ present := make([]int32, rgSize)
+ absent := make([]int32, rgSize)
+ for i := range present {
+ present[i] = 7
+ absent[i] = 1
+ }
+
+ rowGroups := make([][]int32, 0, matchingGroups+2)
+ for range matchingGroups {
+ rowGroups = append(rowGroups, present)
+ }
+ rowGroups = append(rowGroups, absent, absent)
+
+ data := buildDictionaryTestParquet(t, rowGroups...)
+ rdr := openBloomTestReader(t, data)
+ defer rdr.Close()
+
+ var survivors []internal.RowGroupSpan
+ tester := &internal.ParquetRowGroupTester{
+ StatsFn: func(_ *metadata.RowGroupMetaData, _ []int) (bool,
error) {
+ return true, nil
+ },
+ DictionaryPreds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes: [][]byte{int32PhysBytes(7)}},
+ },
+ Survivors: &survivors,
+ }
+ rr, err := rdr.GetRecords(context.Background(), []int{0}, tester)
+ require.NoError(t, err)
+
+ assert.Equal(t, int64(matchingGroups*rgSize), countRecords(t, rr))
+ expected := make([]internal.RowGroupSpan, matchingGroups)
+ for i := range matchingGroups {
+ expected[i] = internal.RowGroupSpan{
+ FirstRowPos: int64(i * rgSize),
+ NumRows: rgSize,
+ }
+ }
+ assert.Equal(t, expected, survivors)
+}
+
+func TestDictionaryRowGroupPruningFallbackUsesMainReader(t *testing.T) {
+ const rgSize = 1024
+ rg0 := make([]int32, rgSize)
+ rg1 := make([]int32, rgSize)
+ for i := range rg0 {
+ rg0[i] = 1
+ rg1[i] = 7
+ }
+
+ rdr := openTestReaderWithoutDictionarySource(t,
buildDictionaryTestParquet(t, rg0, rg1))
+ defer rdr.Close()
+
+ tester := &internal.ParquetRowGroupTester{
+ StatsFn: func(_ *metadata.RowGroupMetaData, _ []int) (bool,
error) {
+ return true, nil
+ },
+ DictionaryPreds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes: [][]byte{int32PhysBytes(7)}},
+ },
+ }
+ rr, err := rdr.GetRecords(context.Background(), []int{0}, tester)
+ require.NoError(t, err)
+
+ assert.Equal(t, int64(rgSize), countRecords(t, rr))
+}
+
+func TestDictionaryRowGroupPruningClosesTemporaryReaderWithoutClosingSource(t
*testing.T) {
+ const rgSize = 1024
+ rg0 := make([]int32, rgSize)
+ rg1 := make([]int32, rgSize)
+ for i := range rg0 {
+ rg0[i] = 1
+ rg1[i] = 7
+ }
+ data := buildDictionaryTestParquet(t, rg0, rg1)
+
+ source := &trackingOpenFile{Reader: bytes.NewReader(data)}
+ pqRdr, err := file.NewParquetReader(source)
+ require.NoError(t, err)
+ arrRdr, err := pqarrow.NewFileReader(pqRdr,
pqarrow.ArrowReadProperties{}, memory.DefaultAllocator)
+ require.NoError(t, err)
+ rdr := internal.WrapParquetFileReaderWithDictionarySource(arrRdr,
source)
+
+ tester := &internal.ParquetRowGroupTester{
+ StatsFn: func(_ *metadata.RowGroupMetaData, _ []int) (bool,
error) {
+ return true, nil
+ },
+ DictionaryPreds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes: [][]byte{int32PhysBytes(7)}},
+ },
+ }
+ rr, err := rdr.GetRecords(context.Background(), []int{0}, tester)
+ require.NoError(t, err)
+ assert.Equal(t, int64(rgSize), countRecords(t, rr))
+ assert.False(t, source.closed, "temporary dictionary reader must not
close the shared source")
+
+ require.NoError(t, rdr.Close())
+ assert.True(t, source.closed)
+}
+
+func BenchmarkDictionaryRowGroupPruning(b *testing.B) {
+ const (
+ numRowGroups = 16
+ rowsPerGroup = 4096
+ )
+
+ const commonValue int32 = 1 << 30
+ rowGroups := make([][]int32, numRowGroups)
+ for group := range rowGroups {
+ values := make([]int32, rowsPerGroup)
+ for i := range values {
+ if i == 0 {
+ values[i] = commonValue
+ } else {
+ values[i] = int32(group*2 + i%2)
+ }
+ }
+ rowGroups[group] = values
+ }
+
+ data := buildDictionaryTestParquet(b, rowGroups...)
+ alwaysKeep := func(_ *metadata.RowGroupMetaData, _ []int) (bool, error)
{
+ return true, nil
+ }
+ cols := []int{0}
+ benchmarks := []struct {
+ name string
+ physBytes [][]byte
+ wantRows int64
+ }{
+ {name: "without dictionary", wantRows: int64(numRowGroups *
rowsPerGroup)},
+ {name: "dictionary target absent", physBytes:
[][]byte{int32PhysBytes(-1)}},
+ {
+ name: "dictionary target in one group",
+ physBytes: [][]byte{int32PhysBytes(0)},
+ wantRows: int64(rowsPerGroup),
+ },
+ {
+ name: "dictionary target in every group",
+ physBytes: [][]byte{int32PhysBytes(commonValue)},
+ wantRows: int64(numRowGroups * rowsPerGroup),
+ },
+ }
+
+ for _, benchmark := range benchmarks {
+ b.Run(benchmark.name, func(b *testing.B) {
+ b.ReportAllocs()
+ b.SetBytes(int64(len(data)))
+ b.ResetTimer()
+
+ for range b.N {
+ rdr := openBloomTestReader(b, data)
+ tester :=
&internal.ParquetRowGroupTester{StatsFn: alwaysKeep}
+ if len(benchmark.physBytes) > 0 {
+ tester.DictionaryPreds =
[]internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes:
benchmark.physBytes},
+ }
+ }
+
+ rr, err := rdr.GetRecords(context.Background(),
cols, tester)
+ if err != nil {
+ b.Fatal(err)
+ }
+ rows := countRecords(b, rr)
+ if rows != benchmark.wantRows {
+ b.Fatalf("returned %d rows, want %d",
rows, benchmark.wantRows)
+ }
+ if err := rdr.Close(); err != nil {
+ b.Fatal(err)
+ }
+ }
+ })
+ }
+}
+
+// buildDictionaryFallbackTestParquet writes one row group with dictionary
+// pages followed by plain data pages. The returned target exists only in the
+// plain portion of the column.
+func buildDictionaryFallbackTestParquet(t *testing.T) ([]byte, int) {
+ t.Helper()
+
+ valueWidth := 32
+ lowCardinality := make([]parquet.ByteArray, 2048)
+ for i := range lowCardinality {
+ value := make([]byte, valueWidth)
+ copy(value, fmt.Sprintf("category-%02d", i%16))
+ lowCardinality[i] = parquet.ByteArray(value)
+ }
+
+ target := "plain-only-target"
+ plainValues := make([]parquet.ByteArray, 16)
+ for i := range plainValues {
+ value := fmt.Sprintf("plain-value-%02d", i)
+ if i == len(plainValues)-1 {
+ value = target
+ }
+ plainValues[i] = parquet.ByteArray(value)
+ }
+
+ valueNode := schema.NewByteArrayNode("value",
parquet.Repetitions.Required, 1)
+ rootNode, err := schema.NewGroupNode("schema",
parquet.Repetitions.Required,
+ schema.FieldList{valueNode}, -1)
+ require.NoError(t, err)
+
+ writerProps := parquet.NewWriterProperties(
+ parquet.WithStats(true),
+ parquet.WithDictionaryDefault(true),
+ parquet.WithDataPageSize(1024),
+ )
+
+ var buf bytes.Buffer
+ pw := file.NewParquetWriter(&buf, rootNode,
file.WithWriterProps(writerProps))
+ rgw, err := pw.AppendRowGroupChecked()
+ require.NoError(t, err)
+ cw, err := rgw.NextColumn()
+ require.NoError(t, err)
+ byteWriter := cw.(*file.ByteArrayColumnChunkWriter)
+ _, err = byteWriter.WriteBatch(lowCardinality, nil, nil)
+ require.NoError(t, err)
+ byteWriter.FallbackToPlain()
+ _, err = byteWriter.WriteBatch(plainValues, nil, nil)
+ require.NoError(t, err)
+ require.NoError(t, byteWriter.Close())
+ require.NoError(t, rgw.Close())
+ require.NoError(t, pw.Close())
+
+ return buf.Bytes(), len(lowCardinality) + len(plainValues)
+}
+
+func TestDictionaryRowGroupPruningKeepsFallbackColumns(t *testing.T) {
+ data, numRows := buildDictionaryFallbackTestParquet(t)
+
+ pqReader, err := file.NewParquetReader(bytes.NewReader(data))
+ require.NoError(t, err)
+ defer pqReader.Close()
+ chunk, err := pqReader.MetaData().RowGroup(0).ColumnChunk(0)
+ require.NoError(t, err)
+ assert.True(t, chunk.HasDictionaryPage())
+ assert.Contains(t, chunk.Encodings(), parquet.Encodings.RLEDict)
+ assert.Contains(t, chunk.Encodings(), parquet.Encodings.Plain)
+ var hasPlainData bool
+ for _, stat := range chunk.EncodingStats() {
+ if stat.PageType == file.PageTypeDataPage && stat.Encoding ==
parquet.Encodings.Plain {
+ hasPlainData = true
+ }
+ }
+ require.True(t, hasPlainData, "expected a plain fallback data page")
+
+ rdr := openBloomTestReader(t, data)
+ defer rdr.Close()
+ tester := &internal.ParquetRowGroupTester{
+ DictionaryPreds: []internal.RowGroupDictionaryPred{
+ {FieldID: 1, PhysBytes:
[][]byte{[]byte("plain-only-target")}},
+ },
+ StatsFn: func(_ *metadata.RowGroupMetaData, _ []int) (bool,
error) {
+ return true, nil
+ },
+ }
+ rr, err := rdr.GetRecords(context.Background(), []int{0}, tester)
+ require.NoError(t, err)
+
+ assert.Equal(t, int64(numRows), countRecords(t, rr),
+ "fallback data pages must keep the row group even when the
target is absent from the dictionary")
+}
+
func TestParquetRowGroupRangeSelection(t *testing.T) {
const rgSize = 100
