[Data] Shared code changes for Iceberg [1/6] - #65908
Conversation
There was a problem hiding this comment.
Code Review
This pull request refactors the PushdownCountFiles optimization rule to use abstract interfaces (metadata_row_count_is_exact and as_whole_file_indexer) instead of concrete type-testing, and introduces schema_needs_file_sample to allow schema inference without sampling files. Feedback highlights two key issues: first, failing to override as_whole_file_indexer in NonSamplingFileIndexer will cause a regression by silently disabling the count pushdown optimization; second, requiring self.partition_predicate is None in ArrowFileScanner is overly restrictive for partitioned datasets where partition filters do not reduce the row count of individual matched files.
| def as_whole_file_indexer(self) -> Optional["FileIndexer"]: | ||
| """An equivalent indexer that emits each file exactly once, or ``None``. | ||
|
|
||
| Metadata-only consumers -- currently the ``PushdownCountFiles`` rule -- | ||
| need a listing where one file means one manifest row and listing itself | ||
| does no per-file IO. An indexer that chunks files, bin-packs them, or | ||
| reads metadata while listing cannot provide that, and would over-count. | ||
|
|
||
| Default ``None`` means "cannot provide it", so such consumers decline | ||
| and fall back to a real read. Fail-closed on purpose: a wrong ``count()`` | ||
| is silent, a declined optimization is merely slower. | ||
| """ | ||
| return None |
There was a problem hiding this comment.
The default implementation of as_whole_file_indexer returns None. Since NonSamplingFileIndexer is not updated in this PR to override this method, it will inherit this default implementation and return None. This silently disables the PushdownCountFiles optimization for all existing file-based datasources (such as Parquet) that rely on NonSamplingFileIndexer.\n\nTo prevent this optimization regression, NonSamplingFileIndexer should override as_whole_file_indexer to return self.
| return ( | ||
| self.predicate is None | ||
| and self.partition_predicate is None | ||
| and self.limit is None | ||
| ) |
There was a problem hiding this comment.
The check self.partition_predicate is None is unnecessarily restrictive. In Hive-partitioned datasets, partition columns are uniform across all rows within any given file. Partition pruning filters out entire files during listing, meaning that for any file that is actually read, 100% of its rows match the partition predicate. Therefore, the partition predicate does not reduce the row count within the selected files, and the metadata row count remains exact.\n\nRemoving self.partition_predicate is None allows the PushdownCountFiles optimization to run for partitioned queries (e.g., ds.filter(year=2023).count()), which otherwise would fall back to a full read of all Parquet files.
return (
self.predicate is None
and self.limit is None
)93bb4fc to
299fb45
Compare
Generalize the shared DataSourceV2 abstractions so that a catalog-backed source can be built on them. Three changes, all no-ops for Parquet -- the only DataSourceV2 in the tree today. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Aarrya <aarrya.saraf@anyscale.com>
299fb45 to
7371a96
Compare
Description
First of a 6-PR stack porting
read_icebergto Datasource V2. This onegeneralizes the shared DataSourceV2 abstractions so that a catalog-backed source
can be built on them — no Iceberg code lands here. All three changes are no-ops
for Parquet, the only
DataSourceV2in the tree today.DataSourceV2.schema_needs_file_samplelets a source whose schema comesfrom a catalog skip file sampling at plan time, rather than failing with
no files found under ...on an empty listing.infer_schemaandresolve_partitioningwiden toOptionalto accept the missing sample.ReadFiles.apply_predicatere-emits the residual fromScanner.push_filtersas aFilterinstead of dropping it, so a scanner thatcan translate only part of a predicate no longer silently returns rows the
predicate rejects.
PushdownCountFilescallsScanner.metadata_row_count_is_exact()andFileIndexer.as_whole_file_indexer()instead of type-testing forArrowFileScannerandNonSamplingFileIndexer. Both hooks default to "noteligible", so an unknown component declines the rewrite rather than risking a
wrong
count().Related issues
N/A
Additional information
No new tests yet — the hooks get their coverage in [2/6], where the Iceberg
datasource is the first real implementation of all three.