laskoviymishka commented on code in PR #2006:
URL: https://github.com/apache/iceberg-go/pull/2006#discussion_r4046674822
##########
table/evaluators.go:
##########
@@ -1901,3 +1901,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)
Review Comment:
This bare conversion is what makes the "can evolve independently" line on
`RowGroupDictionaryPred` untrue in practice. It only compiles because both
structs have identical fields in the same order today, and if
`RowGroupBloomPred` ever gains a bloom-specific field this silently
zero-initializes it in the dict copy with no compile error and no failing test.
I'd either drop the "evolve independently" claim and note that the two are
intentionally coupled, or collect the dict predicates with their own visitor.
Not blocking, but worth picking one before the comment and the code drift apart.
##########
table/internal/parquet_files.go:
##########
@@ -1936,12 +1959,30 @@ 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(
Review Comment:
I'd make the ownership of this reader explicit before merge. When
`newParquetDictionaryReader` succeeds we get an owned `*file.Reader` that's
never closed, and we can't just add a `defer dictionaryReader.Close()` to fix
that: `file.Reader.Close()` closes the underlying `io.Closer`, and
`dictionarySource` is the same handle the main reader is still using, so
closing it corrupts every subsequent read.
So the code is only correct because it leaks, and the fallback branch
(`dictionaryReader = w.ParquetReader()`) is a borrowed reader that must never
be closed at all. That's a fragile contract for the next person.
I'd wrap `dictionarySource` in a non-closeable shim before handing it to
`newParquetDictionaryReader`, then `defer Close()` on the owned reader only
(frees its buffer pool, leaves the shared source alone). While we're here, the
fallback-to-main-reader path has no test and it shares reader state with the
subsequent data read, so it's worth one.
##########
table/internal/parquet_files.go:
##########
@@ -2019,6 +2075,302 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
return w.GetRecordReader(ctx, cols, rgList)
}
+// Once a few consecutive row groups survive dictionary checks, further checks
+// are unlikely to pay for themselves. Stopping is fail-open: it can only leave
+// extra row groups to the normal reader, never drop matching data.
+const parquetDictionaryKeepStreakLimit = 2
+
+func newParquetDictionaryReader(
+ source parquet.ReaderAtSeeker, fileMeta *metadata.FileMetaData, mem
memory.Allocator,
+) (*file.Reader, error) {
+ readProps := parquet.NewReaderProperties(mem)
+ readProps.BufferedStreamEnabled = true
+
+ return file.NewParquetReader(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,
+ rg int,
+ predsByColumn map[int][]int,
+ preds []RowGroupDictionaryPred,
+) bool {
+ if len(preds) == 0 || len(predsByColumn) == 0 {
+ return true
+ }
+
+ rgMeta := fileMeta.RowGroup(rg)
+ 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
+ }
+
+ for _, predIndex := range predIndexes {
+ if !matches[predIndex] {
+ 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 records
+// which requested predicates have 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 nil, false
+ }
+
+ width, variableWidth, ok := parquetDictionaryValueLayout(physicalType,
typeLen)
+ if !ok {
+ return nil, false
+ }
+
+ for _, predIndex := range predIndexes {
+ usable := false
+ for _, candidate := range preds[predIndex].PhysBytes {
+ if variableWidth || len(candidate) == width {
+ usable = true
+
+ break
+ }
+ }
+ if !usable {
+ return nil, false
+ }
+ }
+
+ // predIndexes contains indexes into the full preds slice, so matches
keeps
+ // the global predicate indexes instead of remapping them per column.
+ matches := make([]bool, len(preds))
+ data := page.Data()
+ offset := 0
+ numValues := page.NumValues()
Review Comment:
One edge here I'd guard: if a page reports `NumValues() == 0` with empty
data, the decode loop is a no-op, `remaining` stays above 0, and `offset !=
len(data)` is `0 != 0` (false), so we return `(matches, true)` with every match
false and the caller prunes the group.
`parquetColumnUsesOnlyDictionaryData` makes that unreachable for a valid
file, but a malformed page slips past it because the completeness check runs
before we read the page. This is the one path where a degenerate page causes an
incorrect prune instead of a conservative keep, which is exactly the invariant
this rests on.
`if numValues == 0 { return nil, false }` right after the `< 0` guard closes
it.
##########
table/internal/parquet_files.go:
##########
@@ -1990,6 +2033,19 @@ func (w wrapPqArrowReader) GetRecords(ctx
context.Context, cols []int, tester an
}
}
+ if use && dictionaryPruningActive {
+ use = checkRowGroupDictionaries(
+ dictionaryReader, fileMeta, rg,
dictionaryPredsByColumn, rowGroupTester.DictionaryPreds)
+ if use {
+ dictionaryKeepStreak++
+ if dictionaryKeepStreak >=
parquetDictionaryKeepStreakLimit {
+ dictionaryPruningActive = false
Review Comment:
This is the piece I'd hold on. Once `dictionaryKeepStreak` hits 2,
`dictionaryPruningActive` flips to false and stays false for the rest of the
file. The `dictionaryKeepStreak = 0` reset only runs inside the `use &&
dictionaryPruningActive` branch, so a later prunable group never turns pruning
back on.
For the common shape (a value clustered in a few groups, then absent across
a long tail) that means we stop pruning right where the savings are: groups 0-4
pruned, 5 and 6 kept, then 7-99 read without a dict check even though none of
them contain the value. The `dictionary target in every group` benchmark only
exercises the case where nothing was prunable anyway, so nothing flags the
regression.
I'd either raise the limit (Java-side heuristics sit around 8-16) or tie the
cutoff to the fraction of groups pruned rather than an absolute streak, and add
a unit test with a 4+ group file where the target is in groups 0-1 and absent
after, asserting on Survivors that the tail still gets pruned. wdyt?
--
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]
