Cache reusable Parquet pruning setup across files with the same schema#21566
Open
kosiew wants to merge 12 commits into
Open
Cache reusable Parquet pruning setup across files with the same schema#21566kosiew wants to merge 12 commits into
kosiew wants to merge 12 commits into
Conversation
Contributor
|
This is really awesome! I was planning on trying to add something like that once the morsel work is more mature, but I was wondering if it'll be possible to make it more format-independent? Ran into similar issue on our |
Contributor
Author
|
Thanks! I agree this is the direction we should take. This PR keeps the cache Parquet-local on purpose because the reusable setup currently stores Parquet-specific artifacts: the adapted Parquet projection/predicate, the physical schema after Parquet file-schema coercions / INT96 handling, and the row-group PruningPredicate. It also leaves page-index work, reader metadata, file metrics, and access-plan execution per file. That said, the shape is a good stepping stone toward the more format-independent problem in #20078. The parts that seem general are: scan-local reuse keyed by logical schema, physical schema, projection, predicate, and adapter cache-safety; I would prefer to land this narrowly for Parquet first, with the cache-safety contract and tests in place, then follow up by extracting the format-neutral expression adaptation / pruning setup cache into a datasource-level helper once another FileSource can exercise it. Vortex or another custom FileSource would be a good second consumer to make sure the abstraction is not overfit to Parquet row-group pruning. |
Contributor
|
That makes perfect sense, would be happy to help with anything here! This is really awesome stuff |
Contributor
|
👋 Is this still being worked on? |
Contributor
Author
|
@JSOD11 |
Add a scan-local ParquetPruningSetupCache in opener.rs to reuse adapted projection, adapted predicate, and row-group pruning predicate setup for same-schema files. Implement PhysicalExprAdapterFactory::supports_reusable_rewrites() in schema_rewriter.rs to allow default adapters while defaulting custom ones to opt-out. Connect the cache to ParquetMorselizer in source.rs and introduce a regression test to ensure the setup is reused correctly across same-schema files while keeping it conservative to avoid per-file literal replacements.
Simplify the pruning setup cache in opener.rs by removing the wrapper entry struct, deriving PartialEq/Eq, and returning cloned setup values instead of an extra Arc. Extract the cache-or-build branch into build_or_get_pruning_setup and avoid repeated literal_columns.is_empty() checks. Move returned pruning setup fields directly into prepared and simplify the test-only counting adapter and cache regression test loop. Tighten the supports_reusable_rewrites doc comment in schema_rewriter.rs without changing the public interface.
Refactor cache in opener.rs to utilize a HashMap, computing cold misses outside the mutex for better performance, with a re-check on insert. Add cache-boundary tests for non-reusable adapters and diverse physical schemas. Clarify documentation in schema_rewriter.rs for supports_reusable_rewrites() to specify the same logical/physical schema rewrite inputs.
Enhance opener.rs with cache-key rationale comments for clarity. Update ParquetPruningSetupCache to handle concurrent cold misses by implementing a per-entry pending/ready state. Clarify the public contract of supports_reusable_rewrites() in schema_rewriter.rs.
- Added missing local imports for TimeUnit and ParquetAccessPlan. - Removed stale DataFusionError import. - Updated fully qualified RecordBatch types to avoid lint issues. - Rewired three cache tests to use the local ParquetMorselizerBuilder helper instead of an outdated API. - Ensured all tests passed by running `cargo test -p datafusion-datasource-parquet --lib -- opener::test`, with all 23 opener tests successful.
kosiew
force-pushed
the
schema_caching-01-21554
branch
from
1a44061 to
31af7fc
Compare
- Disabled pruning setup cache when predicate contains dynamic filters. - Disabled cache when projection contains input_file_name(). Added tests for: - Dynamic filter stale snapshot regression. - input_file_name() projection not populating reusable cache.
- Added cache entries type alias for improved readability. - Deduplicated lock-poison error construction for consistency. - Extracted cache eligibility helper to streamline code. - Implemented shared test helper for input_file_name_expr(). - Added a two-file cache test helper and removed repeated loops for efficiency.
kosiew
changed the title
- Removed `ParquetPruningSetupCacheEntry`, `Pending`, `Ready`, `Failed` states, and associated per-entry synchronization logic (Mutex, Condvar). - Introduced a new cache structure using `Mutex<HashMap<ParquetPruningSetupCacheKey, ParquetPruningSetup>>`. - Updated cache retrieval logic: - On hit: return a cloned entry. - On miss: build a new entry, insert it, and return. - On error: no insertion is performed.
Contributor
|
run benchmark clickbench_partitioned |
Contributor
|
I merged up from main and kicked off some benchmark runs If it shows improvements I'll take a closer look |
Contributor
|
I think in general for performance related PRs, I try to make sure we have benchmark results that justify the new code (and also to help figure out what to review in what order) |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing schema_caching-01-21554 (94b3e6f) to 754e113 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
Contributor
|
run benchmark clickbench_partitioned |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing schema_caching-01-21554 (94b3e6f) to 754e113 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
Parquet scans currently rewrite, simplify, and construct row-group pruning predicates independently for each file, even when multiple files share the same logical and physical schema. This repeated CPU-only preparation work is identical across those files and can be reused.
This change introduces a scan-scoped cache for reusable pruning setup while preserving all file-specific work such as metadata loading, page-index handling, reader state, and execution planning.
What changes are included in this PR?
Add a scan-scoped
ParquetPruningSetupCacheowned byParquetMorselizer.Cache adapted projection, adapted predicate, and the row-group pruning predicate for reusable schema adaptation results.
Refactor pruning setup construction into reusable helper functions (
build_pruning_setup/build_or_get_pruning_setup).Reuse cached setup only when it is safe to do so, including:
Skip caching for cases that depend on per-file or mutable state, including:
input_file_name()projections,Extend
PhysicalExprAdapterFactorywith asupports_reusable_rewrites()capability, with the default adapter opting in and custom adapters remaining non-reusable by default.Preserve the existing fallback path for non-cacheable scenarios.
Are these changes tested?
Yes. This PR adds targeted tests covering:
input_file_name()projections,Are there any user-facing changes?
No. This change is an internal optimization for Parquet scan preparation and is not intended to change query results or user-facing behavior.
LLM-generated code disclosure
This PR includes LLM-generated code and comments. All LLM-generated content has been manually reviewed.