laskoviymishka commented on code in PR #2006:
URL: https://github.com/apache/iceberg-go/pull/2006#discussion_r4064217621
##########
table/internal/parquet_files.go:
##########
@@ -1990,6 +2038,23 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
}
}
+ if use && dictionaryPruningActive {
+ if rgMeta == nil {
+ rgMeta = fileMeta.RowGroup(rg)
+ }
+ use = checkRowGroupDictionaries(
+ dictionaryReader, fileMeta, rgMeta, rg,
+ dictionaryPredsByColumn,
rowGroupTester.DictionaryPreds)
+ if use {
+ dictionaryKeepStreak++
+ if dictionaryKeepStreak >=
parquetDictionaryKeepStreakLimit {
+ dictionaryPruningActive = false
Review Comment:
This is the streak breaker from last round, and it's still the thing I'd
hold on. Once eight groups in a row survive, `dictionaryPruningActive` flips
false and never comes back on for the rest of the file, because the `else`
reset here only runs while pruning is still active so it can't re-arm it. For
the common shape (a value clustered in the first groups, then absent down the
long tail) we stop pruning exactly where the wins are, and for older files with
no bloom filter this is the only sub-stats pass, so we lose it for most of the
file. Java's `ParquetDictionaryRowGroupFilter` just decodes every row group
unconditionally.
Given the buffered `dictionarySource` already covers the per-page read cost,
I'd drop the streak mechanism rather than tune the limit. If we want to keep a
guard it needs to key on the prune fraction so a genuinely-pruning tail keeps
getting checked, but I think removing it is the cleaner call. wdyt?
##########
table/internal/parquet_files.go:
##########
@@ -2019,6 +2084,396 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
return w.GetRecordReader(ctx, cols, rgList)
}
+// Once several consecutive row groups survive dictionary checks, further
+// checks are unlikely to pay for themselves. The limit is deliberately high
+// enough to keep checking after a short cluster of matching groups, since a
+// later group may still be prunable. Stopping is fail-open: it can only leave
+// extra row groups to the normal reader, never drop matching data.
+const parquetDictionaryKeepStreakLimit = 8
+
+// 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 || !parquetColumnUsesOnlyDictionaryData(chunk) {
Review Comment:
If `ColumnChunk` ever returns a nil chunk with a nil error,
`parquetColumnUsesOnlyDictionaryData` dereferences it right away via
`chunk.HasDictionaryPage()` and we panic in the pruning hot path instead of
failing open. A `chunk == nil` in this guard keeps it conservative:
```go
if err != nil || chunk == nil || !parquetColumnUsesOnlyDictionaryData(chunk)
{
```
##########
table/evaluators.go:
##########
@@ -1901,3 +1901,40 @@ 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{
Review Comment:
This is the red check now. Staticcheck S1016 wants a direct conversion
instead of the field-by-field literal, since the two pred types are
structurally identical:
```go
for i, pred := range bloomPreds {
dictionaryPreds[i] = internal.RowGroupDictionaryPred(pred)
}
```
That shares the `PhysBytes` slice between the bloom and dictionary preds,
which is fine today since neither path writes it, but a `slices.Clone` there is
cheap insurance if we ever sort or dedup in place. Green CI is the one hard
blocker left before I approve.
##########
table/internal/parquet_files_test.go:
##########
@@ -2211,6 +2224,505 @@ 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 TestDictionaryRowGroupPruningKeepsCheckingAfterKeepStreak(t *testing.T) {
Review Comment:
This test is named for the keep-streak, but with the limit at 8 and only
four groups it never actually trips the breaker. `dictionaryPruningActive`
stays true the whole way, so the disable path has no coverage, and that's how
the permanent-disable behavior slipped through.
If we keep any streak logic I'd want a case with exactly
`parquetDictionaryKeepStreakLimit` matching groups followed by two prunable
ones, asserting the tail groups survive so the fail-open trade-off is pinned.
If we drop the streak, this test can just go.
##########
table/internal/parquet_files.go:
##########
@@ -2019,6 +2084,396 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
return w.GetRecordReader(ctx, cols, rgList)
}
+// Once several consecutive row groups survive dictionary checks, further
+// checks are unlikely to pay for themselves. The limit is deliberately high
+// enough to keep checking after a short cluster of matching groups, since a
+// later group may still be prunable. Stopping is fail-open: it can only leave
+// extra row groups to the normal reader, never drop matching data.
+const parquetDictionaryKeepStreakLimit = 8
+
+// 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 || !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 {
Review Comment:
Last round this path pruned when it shouldn't; now it fails open, which is
the safe direction, thanks. It's slightly over-corrected though: a zero-entry
dictionary page means the column has no non-null values, so an EqualTo/In
genuinely can't match, and this could return `false, true` (known, prunable)
rather than handing the group to the reader. Minor, and only worth it if it's
easy, since the safe answer is already correct.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
