Skip to content

[Data] Shared code changes for Iceberg [1/6] - #65908

Open
AarryaSaraf wants to merge 1 commit into
ray-project:masterfrom
AarryaSaraf:iceberg-dsv2-1-shared-seams
Open

[Data] Shared code changes for Iceberg [1/6]#65908
AarryaSaraf wants to merge 1 commit into
ray-project:masterfrom
AarryaSaraf:iceberg-dsv2-1-shared-seams

Conversation

@AarryaSaraf

Copy link
Copy Markdown
Contributor

Description

First of a 6-PR stack porting read_iceberg to Datasource V2. This one
generalizes 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 DataSourceV2 in the tree today.

  • DataSourceV2.schema_needs_file_sample lets a source whose schema comes
    from a catalog skip file sampling at plan time, rather than failing with
    no files found under ... on an empty listing. infer_schema and
    resolve_partitioning widen to Optional to accept the missing sample.

  • ReadFiles.apply_predicate re-emits the residual from
    Scanner.push_filters as a Filter instead of dropping it, so a scanner that
    can translate only part of a predicate no longer silently returns rows the
    predicate rejects.

  • PushdownCountFiles calls Scanner.metadata_row_count_is_exact() and
    FileIndexer.as_whole_file_indexer() instead of type-testing for
    ArrowFileScanner and NonSamplingFileIndexer. Both hooks default to "not
    eligible", 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.

@AarryaSaraf AarryaSaraf added the data Ray Data-related issues label Sep 3, 2026
@AarryaSaraf
AarryaSaraf marked this pull request as ready for review September 3, 2026 23:51
@AarryaSaraf
AarryaSaraf requested a review from a team as a code owner September 3, 2026 23:51
@AarryaSaraf AarryaSaraf added the go add ONLY when ready to merge, run all tests label Sep 3, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +39 to +51
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Comment on lines +71 to +75
return (
self.predicate is None
and self.partition_predicate is None
and self.limit is None
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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
        )

@AarryaSaraf
AarryaSaraf force-pushed the iceberg-dsv2-1-shared-seams branch from 93bb4fc to 299fb45 Compare September 3, 2026 23:57
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>
@AarryaSaraf
AarryaSaraf force-pushed the iceberg-dsv2-1-shared-seams branch from 299fb45 to 7371a96 Compare September 4, 2026 01:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

data Ray Data-related issues go add ONLY when ready to merge, run all tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant