From 10ae79e1da6603cc1c07c9ade91902706488dd61 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sat, 29 Aug 2026 12:10:27 +0200 Subject: [PATCH] perf(table): reuse partition projections across planning phases Signed-off-by: Minh Vu --- table/incremental_append_scan.go | 6 +- table/scanner.go | 29 ++++++-- table/scanner_internal_test.go | 70 +++++++++++++++++- table/scanner_partition_bench_test.go | 102 +++++++++++++++++++++++++- 4 files changed, 195 insertions(+), 12 deletions(-) diff --git a/table/incremental_append_scan.go b/table/incremental_append_scan.go index 5cc449ff6..98a1f6415 100644 --- a/table/incremental_append_scan.go +++ b/table/incremental_append_scan.go @@ -197,7 +197,9 @@ func (s *IncrementalAppendScan) PlanFiles(ctx context.Context) ([]FileScanTask, manifestList = append(manifestList, manifestsByPath[path]) } - manifestList, err = planningScan.filterManifestsWithSchema(manifestList, schema, &acc) + // Use one projection cache for manifest-summary and data-file pruning. + partitionFilters := planningScan.partitionFiltersForSchema(schema) + manifestList, err = planningScan.filterManifestsWithSchema(manifestList, schema, &acc, partitionFilters) if err != nil { return nil, err } @@ -208,7 +210,7 @@ func (s *IncrementalAppendScan) PlanFiles(ctx context.Context) ([]FileScanTask, // one factory result per concurrent batch, then reacquire through the factory // so long-running incremental plans can renew vended credentials between // batches. - entries, err := planningScan.collectManifestEntriesWithSchema(ctx, manifestList, schema) + entries, err := planningScan.collectManifestEntriesWithSchema(ctx, manifestList, schema, partitionFilters) if err != nil { return nil, err } diff --git a/table/scanner.go b/table/scanner.go index cd177838a..00a0e4ef4 100644 --- a/table/scanner.go +++ b/table/scanner.go @@ -897,20 +897,27 @@ func (scan *Scan) fetchPartitionSpecFilteredManifests(ctx context.Context) ([]ic // this accumulator are intentionally discarded. A future caller that needs // those counts should use fetchPartitionSpecFilteredManifestsWithSchema and // pass in an accumulator it actually reads. - return scan.fetchPartitionSpecFilteredManifestsWithSchema(snap, fs, schema, &scanMetricsAccumulator{}) + return scan.fetchPartitionSpecFilteredManifestsWithSchema( + snap, fs, schema, &scanMetricsAccumulator{}, scan.partitionFiltersForSchema(schema)) } // fetchPartitionSpecFilteredManifestsWithSchema loads the snapshot's manifests // with fs and filters them using the given schema. It records // total/scanned/skipped manifest counts (split by data vs delete content) into acc. -func (scan *Scan) fetchPartitionSpecFilteredManifestsWithSchema(snap *Snapshot, fs io.IO, schema *iceberg.Schema, acc *scanMetricsAccumulator) ([]iceberg.ManifestFile, error) { +func (scan *Scan) fetchPartitionSpecFilteredManifestsWithSchema( + snap *Snapshot, + fs io.IO, + schema *iceberg.Schema, + acc *scanMetricsAccumulator, + partitionFilters *keyDefaultMapErr[int, iceberg.BooleanExpression], +) ([]iceberg.ManifestFile, error) { // Fetch all manifests for the current snapshot. manifestList, err := snap.Manifests(fs) if err != nil { return nil, err } - return scan.filterManifestsWithSchema(manifestList, schema, acc) + return scan.filterManifestsWithSchema(manifestList, schema, acc, partitionFilters) } // filterManifestsWithSchema applies partition-summary pruning to an existing @@ -920,9 +927,9 @@ func (scan *Scan) filterManifestsWithSchema( manifestList []iceberg.ManifestFile, schema *iceberg.Schema, acc *scanMetricsAccumulator, + partitionFilters *keyDefaultMapErr[int, iceberg.BooleanExpression], ) ([]iceberg.ManifestFile, error) { // Build per-spec manifest evaluators and filter out irrelevant manifests. - partitionFilters := scan.partitionFiltersForSchema(schema) manifestEvaluators := newKeyDefaultMapWrapErr(func(specID int) (func(iceberg.ManifestFile) (bool, error), error) { return buildManifestEvaluator(specID, scan.metadata, schema, partitionFilters, scan.caseSensitive) }) @@ -1022,13 +1029,15 @@ func (scan *Scan) collectManifestEntries( return nil, err } - return scan.collectManifestEntriesWithSchema(ctx, manifestList, schema) + return scan.collectManifestEntriesWithSchema( + ctx, manifestList, schema, scan.partitionFiltersForSchema(schema)) } func (scan *Scan) collectManifestEntriesWithSchema( ctx context.Context, manifestList []iceberg.ManifestFile, schema *iceberg.Schema, + partitionFilters *keyDefaultMapErr[int, iceberg.BooleanExpression], ) (*manifestEntries, error) { metricsEval, err := newInclusiveMetricsEvaluator( schema, @@ -1048,7 +1057,6 @@ func (scan *Scan) collectManifestEntriesWithSchema( g, gctx := errgroup.WithContext(ctx) g.SetLimit(concurrencyLimit) - partitionFilters := scan.partitionFiltersForSchema(schema) partitionEvaluators := newKeyDefaultMapWrapErr(func(specID int) (func(iceberg.DataFile) (bool, error), error) { return buildPartitionEvaluator(specID, scan.metadata, schema, partitionFilters, scan.caseSensitive) }) @@ -1187,8 +1195,13 @@ func (scan *Scan) planFilesLocal(ctx context.Context, acc *scanMetricsAccumulato // one FileIO within each concurrent batch, while the next batch loads again // so credential-renewing factories retain their checkpoints. + // Keep the projection cache alive across both local planning phases. The + // manifest and data-file evaluators need the same per-spec projections. + partitionFilters := scan.partitionFiltersForSchema(schema) + // Step 1: Retrieve filtered manifests based on snapshot and partition specs. - manifestList, err := scan.fetchPartitionSpecFilteredManifestsWithSchema(snap, fs, schema, acc) + manifestList, err := scan.fetchPartitionSpecFilteredManifestsWithSchema( + snap, fs, schema, acc, partitionFilters) if err != nil || len(manifestList) == 0 { return nil, err } @@ -1201,7 +1214,7 @@ func (scan *Scan) planFilesLocal(ctx context.Context, acc *scanMetricsAccumulato } // Step 2: Read manifest entries concurrently, accumulating data and positional deletes. - entries, err := scan.collectManifestEntriesWithSchema(ctx, manifestList, schema) + entries, err := scan.collectManifestEntriesWithSchema(ctx, manifestList, schema, partitionFilters) if err != nil { return nil, err } diff --git a/table/scanner_internal_test.go b/table/scanner_internal_test.go index 5ce22de40..206642109 100644 --- a/table/scanner_internal_test.go +++ b/table/scanner_internal_test.go @@ -1176,7 +1176,8 @@ func TestFetchManifestCountersWithRealSnapshot(t *testing.T) { var acc scanMetricsAccumulator snapshot, err := scan.ResolveSnapshot() require.NoError(t, err) - filtered, err := scan.fetchPartitionSpecFilteredManifestsWithSchema(snapshot, memIO, schema, &acc) + filtered, err := scan.fetchPartitionSpecFilteredManifestsWithSchema( + snapshot, memIO, schema, &acc, scan.partitionFiltersForSchema(schema)) require.NoError(t, err) // Two data manifests, one delete manifest. @@ -1238,6 +1239,7 @@ func TestFilterManifestsWithSchemaSkipsKnownEmptyManifests(t *testing.T) { []iceberg.ManifestFile{knownEmptyData, knownEmptyDelete, unknownCounts, live}, schema, &acc, + scan.partitionFiltersForSchema(schema), ) require.NoError(t, err) require.Len(t, filtered, 2) @@ -1313,6 +1315,72 @@ func TestPlanFilesSkipsKnownEmptyManifestsBeforeOpening(t *testing.T) { assert.Zero(t, fs.openCount[deletePath]) } +func TestScanReusesPartitionFiltersAcrossPlanningPhases(t *testing.T) { + schema := iceberg.NewSchema(1, iceberg.NestedField{ + ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32, Required: true, + }) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, + FieldID: 1000, + Name: "id", + Transform: iceberg.IdentityTransform{}, + }) + metadata, err := NewMetadata( + schema, &spec, UnsortedSortOrder, "mem://default/table", iceberg.Properties{}, + ) + require.NoError(t, err) + + const manifestPath = "mem://default/table/metadata/manifest.avro" + snapshotID := int64(1) + dataFile, err := iceberg.NewDataFileBuilder( + spec, + iceberg.EntryContentData, + "mem://default/table/data.parquet", + iceberg.ParquetFile, + map[int]any{1000: int32(5)}, + nil, + nil, + 1, + 1, + ) + require.NoError(t, err) + entry := iceberg.NewManifestEntryBuilder( + iceberg.EntryStatusADDED, &snapshotID, dataFile.Build(), + ).SequenceNum(1).Build() + var manifestBytes bytes.Buffer + manifest, err := iceberg.WriteManifest( + manifestPath, &manifestBytes, 2, spec, schema, snapshotID, []iceberg.ManifestEntry{entry}, + ) + require.NoError(t, err) + + memIO := iceio.NewMemFS() + require.NoError(t, memIO.WriteFile(manifestPath, manifestBytes.Bytes())) + + scan := &Scan{ + metadata: metadata, + ioF: func(context.Context) (iceio.IO, error) { return memIO, nil }, + rowFilter: iceberg.EqualTo(iceberg.Reference("id"), int32(5)), + caseSensitive: true, + concurrency: 1, + } + var projectionCalls atomic.Int32 + partitionFilters := newKeyDefaultMapWrapErr(func(specID int) (iceberg.BooleanExpression, error) { + projectionCalls.Add(1) + + return buildPartitionProjection(specID, metadata, schema, scan.rowFilter, scan.caseSensitive) + }) + var acc scanMetricsAccumulator + + _, err = scan.filterManifestsWithSchema([]iceberg.ManifestFile{manifest}, schema, &acc, partitionFilters) + require.NoError(t, err) + _, err = scan.collectManifestEntriesWithSchema( + context.Background(), []iceberg.ManifestFile{manifest}, schema, partitionFilters) + require.NoError(t, err) + + assert.Len(t, partitionFilters.data, 1) + assert.Equal(t, int32(1), projectionCalls.Load()) +} + func TestBuildManifestEvaluatorWithInvalidSpecID(t *testing.T) { schema := iceberg.NewSchema( 1, diff --git a/table/scanner_partition_bench_test.go b/table/scanner_partition_bench_test.go index 192ee0e30..0d9eee161 100644 --- a/table/scanner_partition_bench_test.go +++ b/table/scanner_partition_bench_test.go @@ -24,7 +24,107 @@ import ( "github.com/apache/iceberg-go" ) -var partitionEvaluatorBenchmarkSink int +var ( + partitionEvaluatorBenchmarkSink int + partitionProjectionBenchmarkSink int +) + +func BenchmarkPartitionProjectionPlanning(b *testing.B) { + for _, specCount := range []int{8, 64, 256} { + scan, schema := benchmarkPartitionProjectionScan(b, specCount) + b.Run(fmt.Sprintf("specs=%d", specCount), func(b *testing.B) { + b.Run("separate_caches", func(b *testing.B) { + benchmarkPartitionProjectionPhases(b, scan, schema, specCount, false) + }) + b.Run("shared_cache", func(b *testing.B) { + benchmarkPartitionProjectionPhases(b, scan, schema, specCount, true) + }) + }) + } +} + +func benchmarkPartitionProjectionPhases( + b *testing.B, + scan *Scan, + schema *iceberg.Schema, + specCount int, + shared bool, +) { + b.ReportAllocs() + b.ResetTimer() + var projectionBuilds int64 + newPartitionFilters := func() *keyDefaultMapErr[int, iceberg.BooleanExpression] { + return newKeyDefaultMapWrapErr(func(specID int) (iceberg.BooleanExpression, error) { + projectionBuilds++ + + return buildPartitionProjection(specID, scan.metadata, schema, scan.rowFilter, scan.caseSensitive) + }) + } + + for b.Loop() { + manifestFilters := newPartitionFilters() + partitionFilters := manifestFilters + if !shared { + partitionFilters = newPartitionFilters() + } + + manifestEvaluators := newKeyDefaultMapWrapErr(func(specID int) (func(iceberg.ManifestFile) (bool, error), error) { + return buildManifestEvaluator(specID, scan.metadata, schema, manifestFilters, scan.caseSensitive) + }) + for specID := range specCount { + if _, err := manifestEvaluators.Get(specID); err != nil { + b.Fatal(err) + } + } + + partitionEvaluators := newKeyDefaultMapWrapErr(func(specID int) (func(iceberg.DataFile) (bool, error), error) { + return buildPartitionEvaluator(specID, scan.metadata, schema, partitionFilters, scan.caseSensitive) + }) + for specID := range specCount { + if _, err := partitionEvaluators.Get(specID); err != nil { + b.Fatal(err) + } + } + + partitionProjectionBenchmarkSink = len(manifestFilters.data) + len(partitionFilters.data) + } + + b.StopTimer() + b.ReportMetric(float64(projectionBuilds)/float64(b.N), "projection-builds/op") +} + +func benchmarkPartitionProjectionScan(b *testing.B, specCount int) (*Scan, *iceberg.Schema) { + b.Helper() + + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32, Required: true}, + iceberg.NestedField{ID: 2, Name: "payload", Type: iceberg.PrimitiveTypes.String, Required: true}, + ) + specs := make([]iceberg.PartitionSpec, specCount) + for specID := range specCount { + specs[specID] = iceberg.NewPartitionSpecID(specID, iceberg.PartitionField{ + SourceIDs: []int{1}, + FieldID: 1000 + specID, + Name: fmt.Sprintf("id_%d", specID), + Transform: iceberg.IdentityTransform{}, + }) + } + + metadata := &metadataV2{commonMetadata: commonMetadata{ + SchemaList: []*iceberg.Schema{schema}, + CurrentSchemaID: schema.ID, + Specs: specs, + }} + + return &Scan{ + metadata: metadata, + rowFilter: iceberg.NewAnd( + iceberg.EqualTo(iceberg.Reference("id"), int32(7)), + iceberg.GreaterThanEqual(iceberg.Reference("payload"), "a"), + ), + caseSensitive: true, + }, schema +} func BenchmarkPartitionEvaluator(b *testing.B) { for _, fieldCount := range []int{1, 8, 32} {