From 11f3d74b26f9eeca396265665a287f0d23e59ba6 Mon Sep 17 00:00:00 2001 From: Aarrya Date: Fri, 31 Jul 2026 17:13:40 -0700 Subject: [PATCH 1/6] [Data] Fix Iceberg Dataset.count() returning wrong row counts `Dataset.count()` returned a wrong number for three common Iceberg queries, in both directions: table has 100 rows; 10 of them have id < 10 query count() truth filter(col('id') < 10).count() 0 10 filter(col('id')<10).select_columns(['country']).count() 0 10 read_iceberg(row_filter=id<10).count() 100 10 Both are cases of the datasource claiming a row count it does not deliver. Empty projections lost every row. `count()` asks the read for zero columns -- it needs a row count, not values -- but PyIceberg builds its output by reconstructing a table against the projected schema, and a table reconstructed from zero columns reports zero rows however many the scan matched. Read the cheapest single column instead and hand the result to Ray's zero-column projection helper, which slices the column away (preserving the row count) and appends the existing `__bsp_stub` placeholder. Reaching this needs two optimizer passes: predicate pushdown deletes the `Filter`, which is what puts the `Project` directly above the read for projection pushdown to fold in. Read-task metadata ignored pushed-down filters. `num_rows` came from the manifests' `record_count`, written when the data was written and therefore filter-blind, so a selective filter still reported every row in the surviving files. Report the count only when it is provably exact: PyIceberg exposes what is left of the filter after partition pruning as `FileScanTask.residual`, and `AlwaysTrue` means pruning resolved the filter completely. This keeps the free count for partition-column filters (which drop whole files) and reports `None` -- unknown -- otherwise. Also stop mutating caller-owned dicts. `catalog_kwargs` was stored by reference and then `pop("name")`-ed, and `snapshot_id` was written into `scan_kwargs`, so a second `read_iceberg` reusing the same dict raised `NoSuchTableError` or silently inherited a stale snapshot. Copy both. Tests: the Iceberg suite had no `count()` or `num_rows` assertions at all, which is how both bugs survived. Adds 12, asserting an invariant rather than constants -- the count the datasource claims must equal the rows it produces -- so future pushdowns are covered too. Two run at the user level (`count()` vs `take_all()`) and two at the datasource level (metadata vs rows read), and `partition_filter` pins the case where the shortcut is sound so the fix cannot regress it into a full read. python/ray/data/tests/datasource/test_iceberg.py: before: 6 failed, 60 passed after: 66 passed AI assistance was used for this change. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Aarrya --- .../datasource/iceberg_datasource.py | 125 +++++++++++- .../ray/data/tests/datasource/test_iceberg.py | 179 ++++++++++++++++++ 2 files changed, 294 insertions(+), 10 deletions(-) diff --git a/python/ray/data/_internal/datasource/iceberg_datasource.py b/python/ray/data/_internal/datasource/iceberg_datasource.py index 41f8afe6d031..48543d22980e 100644 --- a/python/ray/data/_internal/datasource/iceberg_datasource.py +++ b/python/ray/data/_internal/datasource/iceberg_datasource.py @@ -11,9 +11,10 @@ import pyarrow as pa from packaging import version +from ray.data._internal.arrow_block import _BATCH_SIZE_PRESERVING_STUB_COL_NAME from ray.data._internal.planner.plan_expression.expression_visitors import _ExprVisitor from ray.data._internal.util import _check_import -from ray.data.block import Block, BlockMetadata +from ray.data.block import Block, BlockAccessor, BlockMetadata from ray.data.datasource.datasource import Datasource, ReadTask from ray.data.expressions import ( AliasExpr, @@ -84,6 +85,26 @@ logger = logging.getLogger(__name__) +# Approximate width, in bytes, of the Iceberg primitive types that are cheap to +# read. Used to choose a stand-in column for an empty projection, see +# ``IcebergDatasource._get_sentinel_field``. Keys are Iceberg type names as +# returned by ``str(field_type)``. Types missing from this map -- strings, +# binary, decimals and nested types -- have unbounded width, and are only used +# as a sentinel if the table has nothing cheaper to offer. +_SENTINEL_FIELD_WIDTHS = { + "boolean": 1, + "int": 4, + "float": 4, + "date": 4, + "long": 8, + "double": 8, + "time": 8, + "timestamp": 8, + "timestamptz": 8, + "uuid": 16, +} +_UNKNOWN_FIELD_WIDTH = 1 << 20 + class _IcebergExpressionVisitor( _ExprVisitor["BooleanExpression | UnboundTerm | Literal"] @@ -210,6 +231,7 @@ def _get_read_task( case_sensitive: bool, limit: Optional[int], schema: "Schema", + drop_all_columns: bool = False, ) -> Iterable[Block]: # Determine the PyIceberg version to handle backward compatibility import pyiceberg @@ -255,7 +277,16 @@ def _generate_tables() -> Iterable[pa.Table]: ) yield table - yield from _generate_tables() + for table in _generate_tables(): + if drop_all_columns: + # ``schema`` holds a sentinel column the caller never asked for, see + # ``IcebergDatasource._get_sentinel_field``. Hand it to Ray's + # zero-column projection helper, which slices every column away + # (preserving the row count, unlike rebuilding a table from zero + # columns) and appends the ``__bsp_stub`` placeholder that keeps the + # count alive through downstream ``pa.concat_tables`` calls. + table = BlockAccessor.for_block(table).select([]) + yield table @DeveloperAPI @@ -297,8 +328,14 @@ def __init__( _check_import(self, module="pyiceberg", package="pyiceberg") from pyiceberg.expressions import AlwaysTrue - self._scan_kwargs = scan_kwargs if scan_kwargs is not None else {} - self._catalog_kwargs = catalog_kwargs if catalog_kwargs is not None else {} + # Copy both dicts: below we pop ``name`` out of the catalog kwargs and + # write ``snapshot_id`` into the scan kwargs. Doing that in place would + # mutate dicts the caller still holds, so a second read reusing the same + # dict would fail to find its catalog or silently inherit a snapshot ID. + self._scan_kwargs = dict(scan_kwargs) if scan_kwargs is not None else {} + self._catalog_kwargs = ( + dict(catalog_kwargs) if catalog_kwargs is not None else {} + ) if "name" in self._catalog_kwargs: self._catalog_name = self._catalog_kwargs.pop("name") @@ -367,13 +404,50 @@ def _get_combined_filter(self) -> "BooleanExpression": return combined_filter + def _get_sentinel_field(self) -> str: + """Return the cheapest column to read in place of an empty projection. + + PyIceberg builds its output by reconstructing a table against the + projected schema, and a table reconstructed from zero columns reports + zero rows however many rows the scan actually matched. So an empty + projection -- which ``Dataset.count()`` legitimately requests, since it + needs a row count but no values -- would silently count zero rows. We + read a single column instead and slice it away in ``_get_read_task``, + which preserves the row count. Picking the narrowest column keeps the + wasted I/O small: reading a string column here instead of a boolean one + can cost several orders of magnitude more bytes. + """ + fields = self.table.schema().fields + if not fields: + raise ValueError( + f"Cannot read table '{self.table_identifier}': it has no columns." + ) + # ``min`` is stable, so an all-equal-width schema deterministically + # yields the first field. + cheapest_field = min( + fields, + key=lambda field: _SENTINEL_FIELD_WIDTHS.get( + str(field.field_type), _UNKNOWN_FIELD_WIDTH + ), + ) + return cheapest_field.name + + def _is_empty_projection(self) -> bool: + """Whether the pushed-down projection selects no columns at all.""" + return self._get_data_columns() == [] + def _get_data_scan(self) -> "DataScan": # Get the combined filter combined_filter = self._get_combined_filter() # Convert back to tuple for PyIceberg API (None -> ("*",)) data_columns = self._get_data_columns() - selected_fields = ("*",) if data_columns is None else tuple(data_columns) + if data_columns is None: + selected_fields = ("*",) + elif not data_columns: + selected_fields = (self._get_sentinel_field(),) + else: + selected_fields = tuple(data_columns) data_scan = self.table.scan( row_filter=combined_filter, @@ -433,13 +507,14 @@ def get_read_tasks( per_task_row_limit: Optional[int] = None, data_context: Optional["DataContext"] = None, ) -> List[ReadTask]: + from pyiceberg.expressions import AlwaysTrue from pyiceberg.io import pyarrow as pyi_pa_io from pyiceberg.manifest import DataFileContent # Get the PyIceberg scan data_scan = self._get_data_scan() # Get the plan files in this query - plan_files = self.plan_files + plan_files = list(self.plan_files) # Get the projected schema for this scan, given all the row filters, # snapshot ID, etc. @@ -447,10 +522,20 @@ def get_read_tasks( # Get the arrow schema, to set in the metadata pya_schema = pyi_pa_io.schema_to_pyarrow(projected_schema) + # An empty projection is read as a single sentinel column, see + # ``_get_sentinel_field``, which the read tasks then replace with Ray's + # ``__bsp_stub`` placeholder. Declare that placeholder so the reported + # schema matches the blocks; it is hidden from the user-visible schema. + drop_all_columns = self._is_empty_projection() + if drop_all_columns: + pya_schema = pa.schema( + [pa.field(_BATCH_SIZE_PRESERVING_STUB_COL_NAME, pa.null())] + ) + # Set the n_chunks to the min of the number of plan files and the actual # requested n_chunks, so that there are no empty tasks - if parallelism > len(list(plan_files)): - parallelism = len(list(plan_files)) + if parallelism > len(plan_files): + parallelism = len(plan_files) logger.warning( f"Reducing the parallelism to {parallelism}, as that is the number of files" ) @@ -467,6 +552,21 @@ def get_read_tasks( case_sensitive = self._scan_kwargs.get("case_sensitive", True) limit = self._scan_kwargs.get("limit") + # Whether the manifests' ``record_count`` is still an exact row count. + # It is written when the data is written, so it only holds if every row + # of every surviving file matches the filter. PyIceberg reports whatever + # is left of the filter after partition pruning as + # ``FileScanTask.residual``: ``AlwaysTrue`` means pruning resolved the + # filter completely (e.g. a filter on a partition column dropped whole + # files), so the counts still hold. Anything else has to be evaluated row + # by row, and the counts would over-report. Old PyIceberg versions do not + # populate ``residual``, in which case we cannot prove anything and fall + # back to reporting an unknown row count. + counts_are_exact = isinstance(row_filter, AlwaysTrue) or all( + isinstance(getattr(task, "residual", None), AlwaysTrue) + for task in plan_files + ) + get_read_task = partial( _get_read_task, table_io=table_io, @@ -475,6 +575,7 @@ def get_read_tasks( case_sensitive=case_sensitive, limit=limit, schema=projected_schema, + drop_all_columns=drop_all_columns, ) read_tasks = [] @@ -496,8 +597,12 @@ def get_read_tasks( if delete.content == DataFileContent.POSITION_DELETES ) metadata = BlockMetadata( - num_rows=sum(task.file.record_count for task in chunk_tasks) - - position_delete_count, + num_rows=( + sum(task.file.record_count for task in chunk_tasks) + - position_delete_count + if counts_are_exact + else None + ), size_bytes=sum(task.file.file_size_in_bytes for task in chunk_tasks), input_files=[task.file.file_path for task in chunk_tasks], exec_stats=None, diff --git a/python/ray/data/tests/datasource/test_iceberg.py b/python/ray/data/tests/datasource/test_iceberg.py index a7473c71262b..0636ea1e513f 100644 --- a/python/ray/data/tests/datasource/test_iceberg.py +++ b/python/ray/data/tests/datasource/test_iceberg.py @@ -223,6 +223,185 @@ def test_filtered_read(): assert all(len(rt.metadata.input_files) == 1 for rt in read_tasks) +def _iceberg_scan_row_count(**scan_kwargs) -> int: + """Rows PyIceberg itself returns for a scan -- the ground truth for counts.""" + sql_catalog = pyi_catalog.load_catalog(**_CATALOG_KWARGS.copy()) + table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}") + return table.scan(**scan_kwargs).to_arrow().num_rows + + +# Every entry builds a Dataset whose row count is unambiguous, so ``count()`` can +# be checked against the rows the same query actually yields. ``count()`` has two +# strategies -- read the count from plan metadata, or project to zero columns and +# count what comes back -- and these cases cover both, including the cases that +# defeat the metadata short-circuit. +_COUNT_CASES = { + "plain": lambda: read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + ), + "select_columns": lambda: read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + ).select_columns(["col_b"]), + # An expression filter is pushed into the datasource, which removes the + # ``Filter`` from the plan and leaves the zero-column projection sitting + # directly on the read. + "expr_filter": lambda: read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + ).filter(expr=col("col_a") < 10), + # Same query as a Python UDF: the predicate cannot be pushed down, so the + # ``Filter`` stays in the plan. This is the control case -- it must keep + # working, so a fix cannot simply disable predicate pushdown. + "udf_filter": lambda: read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + ).filter(lambda row: row["col_a"] < 10), + "expr_filter_then_select": lambda: read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + ) + .filter(expr=col("col_a") < 10) + .select_columns(["col_b"]), + # A filter supplied at read time leaves no ``Filter`` in the plan at all, so + # the count is answered from the Iceberg manifest. + "read_time_row_filter": lambda: read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + row_filter=pyi_expr.LessThan("col_a", 10), + catalog_kwargs=_CATALOG_KWARGS.copy(), + ), + "select_columns_materialized": lambda: read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + ) + .select_columns(["col_b"]) + .materialize(), +} + + +@pytest.mark.skipif( + get_pyarrow_version() < parse_version("14.0.0"), + reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", +) +@pytest.mark.parametrize("case", list(_COUNT_CASES), ids=list(_COUNT_CASES)) +def test_count_matches_rows_actually_produced(case): + """``count()`` must agree with the number of rows the query yields. + + A projection or a pushed-down filter must not change the answer: the count is + a property of the query, not of how much of it was pushed into the reader. + """ + make_ds = _COUNT_CASES[case] + expected = len(make_ds().take_all()) + assert make_ds().count() == expected + + +@pytest.mark.skipif( + get_pyarrow_version() < parse_version("14.0.0"), + reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", +) +def test_empty_projection_preserves_row_count(): + """Selecting zero columns must yield N rows of no columns, not zero rows. + + ``Dataset.count()`` projects to zero columns to avoid reading column data, so + an empty projection that drops rows silently corrupts every count built that + way. + """ + iceberg_ds = IcebergDatasource( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + selected_fields=(), + catalog_kwargs=_CATALOG_KWARGS.copy(), + ) + expected = _iceberg_scan_row_count() + assert expected > 0, "fixture table should not be empty" + + rows_read = sum( + block.num_rows + for read_task in iceberg_ds.get_read_tasks(1) + for block in read_task() + ) + assert rows_read == expected + + +@pytest.mark.skipif( + get_pyarrow_version() < parse_version("14.0.0"), + reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", +) +@pytest.mark.parametrize( + ("row_filter", "count_must_be_exact"), + [ + (None, True), + # Prunes whole files, so the manifest row counts stay exact and must + # keep being reported -- this is the free ``count()`` we do not want to + # regress while fixing the case below. + (pyi_expr.In("col_c", {1, 2}), True), + # Selective *within* a file: file-level pruning cannot resolve it, so a + # manifest row count would overcount and must be reported as unknown. + (pyi_expr.LessThan("col_a", 10), False), + ], + ids=["no_filter", "partition_filter", "row_level_filter"], +) +def test_reported_num_rows_matches_rows_read(row_filter, count_must_be_exact): + """Read-task metadata must never claim a row count the read does not deliver. + + ``Dataset.count()`` returns this number directly when the plan has nothing + that could change the row count, so an overcount here reaches the user as + the answer. ``None`` means "unknown" and makes ``count()`` do the read, so it + is always safe -- but it also gives up a free count, hence + ``count_must_be_exact`` pinning the cases where the shortcut is sound. + """ + kwargs = {} if row_filter is None else {"row_filter": row_filter} + iceberg_ds = IcebergDatasource( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + **kwargs, + ) + read_tasks = iceberg_ds.get_read_tasks(2) + + claimed = [read_task.metadata.num_rows for read_task in read_tasks] + actual = sum(block.num_rows for read_task in read_tasks for block in read_task()) + + if count_must_be_exact: + assert all(count is not None for count in claimed), ( + "row counts are exact for this filter and must still be reported, " + "otherwise count() pays for a read it does not need" + ) + if all(count is not None for count in claimed): + assert sum(claimed) == actual + + +@pytest.mark.skipif( + get_pyarrow_version() < parse_version("14.0.0"), + reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", +) +def test_read_iceberg_does_not_mutate_caller_kwargs(): + """The caller's dicts belong to the caller. + + ``catalog_kwargs`` is stored by reference and then ``pop("name")``-ed, so the + caller loses the catalog name and a second read of the same table fails. + """ + catalog_kwargs = _CATALOG_KWARGS.copy() + scan_kwargs = {} + expected_catalog_kwargs = catalog_kwargs.copy() + + first = read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=catalog_kwargs, + scan_kwargs=scan_kwargs, + snapshot_id=None, + ) + assert catalog_kwargs == expected_catalog_kwargs + assert scan_kwargs == {} + + # Reusing the same dicts must work exactly like the first read. + second = read_iceberg( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=catalog_kwargs, + scan_kwargs=scan_kwargs, + ) + assert second.count() == first.count() + + @pytest.mark.skipif( get_pyarrow_version() < parse_version("14.0.0"), reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", From 3146ceddea0afa2d59632c4d584421c19c64bb83 Mon Sep 17 00:00:00 2001 From: Aarrya Date: Mon, 3 Aug 2026 09:16:30 -0700 Subject: [PATCH 2/6] [Data] Iceberg: document the row-count exactness test, harden its test Record why ``FileScanTask.residual`` is the exactness test for the manifests' ``record_count``: it is the same test PyIceberg's own ``DataScan.count()`` applies, so the two agree on when a manifest count still holds. Also strengthen ``test_read_iceberg_does_not_mutate_caller_kwargs``. It passed ``snapshot_id=None``, but the ``scan_kwargs`` write it guards against is behind ``if snapshot_id``, so the ``scan_kwargs == {}`` assertion held whether or not the bug was present -- leaving the silent stale-snapshot variant uncovered. Load the table's current snapshot ID and pass it, so the assertion is reached. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Aarrya --- .../_internal/datasource/iceberg_datasource.py | 3 +++ python/ray/data/tests/datasource/test_iceberg.py | 16 +++++++++++++--- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/python/ray/data/_internal/datasource/iceberg_datasource.py b/python/ray/data/_internal/datasource/iceberg_datasource.py index 48543d22980e..710484b2c1f9 100644 --- a/python/ray/data/_internal/datasource/iceberg_datasource.py +++ b/python/ray/data/_internal/datasource/iceberg_datasource.py @@ -562,6 +562,9 @@ def get_read_tasks( # by row, and the counts would over-report. Old PyIceberg versions do not # populate ``residual``, in which case we cannot prove anything and fall # back to reporting an unknown row count. + # + # This is deliberately the same test PyIceberg's own ``DataScan.count()`` + # applies, so the two agree on when a manifest count still holds. counts_are_exact = isinstance(row_filter, AlwaysTrue) or all( isinstance(getattr(task, "residual", None), AlwaysTrue) for task in plan_files diff --git a/python/ray/data/tests/datasource/test_iceberg.py b/python/ray/data/tests/datasource/test_iceberg.py index 0636ea1e513f..8e4edc62bc2c 100644 --- a/python/ray/data/tests/datasource/test_iceberg.py +++ b/python/ray/data/tests/datasource/test_iceberg.py @@ -377,18 +377,28 @@ def test_reported_num_rows_matches_rows_read(row_filter, count_must_be_exact): def test_read_iceberg_does_not_mutate_caller_kwargs(): """The caller's dicts belong to the caller. - ``catalog_kwargs`` is stored by reference and then ``pop("name")``-ed, so the - caller loses the catalog name and a second read of the same table fails. + Both dicts were stored by reference. ``catalog_kwargs`` is then + ``pop("name")``-ed, so the caller loses the catalog name and a second read of + the same table raises ``NoSuchTableError``. ``scan_kwargs`` has + ``snapshot_id`` written into it, which fails silently instead: a later read + reusing the dict inherits the pin and returns rows from a stale snapshot. """ catalog_kwargs = _CATALOG_KWARGS.copy() scan_kwargs = {} expected_catalog_kwargs = catalog_kwargs.copy() + # An explicit ``snapshot_id`` is required to reach the ``scan_kwargs`` write: + # the buggy line is guarded by ``if snapshot_id``, so passing ``None`` would + # leave the dict untouched whether or not the bug is present. + sql_catalog = pyi_catalog.load_catalog(**_CATALOG_KWARGS.copy()) + table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}") + snapshot_id = table.current_snapshot().snapshot_id + first = read_iceberg( table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", catalog_kwargs=catalog_kwargs, scan_kwargs=scan_kwargs, - snapshot_id=None, + snapshot_id=snapshot_id, ) assert catalog_kwargs == expected_catalog_kwargs assert scan_kwargs == {} From cdec0eb622743c2a147cc1fc260a488ee6307433 Mon Sep 17 00:00:00 2001 From: Aarrya Date: Mon, 3 Aug 2026 10:25:10 -0700 Subject: [PATCH 3/6] [Data] Iceberg: pushdown clones must re-plan their scan ``FileScanTask.residual`` only proves the manifests' ``record_count`` is exact if the plan files were planned *with* the predicate. ``apply_predicate`` and ``apply_projection`` shallow-copy the datasource, so a clone inherited ``_plan_files``; if that cache was populated first -- ``estimate_inmemory_data_size`` reads it, which is how Ray autodetects parallelism -- every residual in it is ``AlwaysTrue``, since nothing was filtered at plan time. ``counts_are_exact`` then reported the unfiltered manifest total, reintroducing the overcount by another route. Override both methods to drop the clone's inherited cache. Not reachable through the ``Dataset`` API today, since each pushdown rule clones before anything warms the cache, but ``IcebergDatasource`` is a ``@DeveloperAPI`` and it is reachable directly. Also materialize the cache as a list. PyIceberg's ``DataScan.plan_files`` returns a list today, but is annotated ``Iterable[FileScanTask]`` and this cache is read more than once, so a future generator would silently yield zero read tasks. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Aarrya --- .../datasource/iceberg_datasource.py | 38 ++++++++++++- .../ray/data/tests/datasource/test_iceberg.py | 54 +++++++++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) diff --git a/python/ray/data/_internal/datasource/iceberg_datasource.py b/python/ray/data/_internal/datasource/iceberg_datasource.py index 710484b2c1f9..4b63a3a9a376 100644 --- a/python/ray/data/_internal/datasource/iceberg_datasource.py +++ b/python/ray/data/_internal/datasource/iceberg_datasource.py @@ -21,6 +21,7 @@ BinaryExpr, ColumnExpr, DownloadExpr, + Expr, LiteralExpr, MonotonicallyIncreasingIdExpr, Operation, @@ -381,10 +382,43 @@ def plan_files(self) -> List["FileScanTask"]: # Calculate and cache the plan_files if they don't already exist if self._plan_files is None: data_scan = self._get_data_scan() - self._plan_files = data_scan.plan_files() + # ``DataScan.plan_files`` is annotated ``Iterable[FileScanTask]``. + # PyIceberg returns a list today, but callers here iterate the cache + # more than once, so materialize it rather than depend on that. + self._plan_files = list(data_scan.plan_files()) return self._plan_files + def apply_predicate(self, predicate_expr: Expr) -> "IcebergDatasource": + """Push a predicate down, discarding plan files planned without it. + + The base implementation shallow-copies, which would leave the clone + sharing a plan-file cache computed before this predicate existed. Those + tasks carry ``AlwaysTrue`` residuals -- nothing was filtered at plan + time -- which ``get_read_tasks`` would read as "the filter is fully + resolved" and report an unfiltered row count for a filtered read. + """ + return self._invalidate_plan_files(super().apply_predicate(predicate_expr)) + + def apply_projection( + self, projection_map: Optional[Dict[str, str]] + ) -> "IcebergDatasource": + """Push a projection down, discarding plan files planned without it. + + The projected columns are part of the scan, so a cache built for a + different projection does not belong to the clone. See + ``apply_predicate``. + """ + return self._invalidate_plan_files(super().apply_projection(projection_map)) + + def _invalidate_plan_files(self, clone: "Datasource") -> "Datasource": + """Drop ``clone``'s inherited plan-file cache, so it re-plans its scan.""" + # A datasource returns ``self`` to signal "no pushdown applied", in which + # case there is no clone and the cache is still valid for this scan. + if clone is not self: + clone._plan_files = None + return clone + def _get_combined_filter(self) -> "BooleanExpression": """Get the combined filter including both row_filter and pushed-down predicates.""" combined_filter = self._row_filter @@ -514,7 +548,7 @@ def get_read_tasks( # Get the PyIceberg scan data_scan = self._get_data_scan() # Get the plan files in this query - plan_files = list(self.plan_files) + plan_files = self.plan_files # Get the projected schema for this scan, given all the row filters, # snapshot ID, etc. diff --git a/python/ray/data/tests/datasource/test_iceberg.py b/python/ray/data/tests/datasource/test_iceberg.py index 8e4edc62bc2c..b685c0d11937 100644 --- a/python/ray/data/tests/datasource/test_iceberg.py +++ b/python/ray/data/tests/datasource/test_iceberg.py @@ -370,6 +370,60 @@ def test_reported_num_rows_matches_rows_read(row_filter, count_must_be_exact): assert sum(claimed) == actual +@pytest.mark.skipif( + get_pyarrow_version() < parse_version("14.0.0"), + reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", +) +@pytest.mark.parametrize( + "push_down", + [ + lambda ds: ds.apply_predicate(col("col_a") < 10), + lambda ds: ds.apply_predicate(col("col_a") < 10).apply_projection( + {"col_b": "col_b"} + ), + ], + ids=["predicate", "predicate_then_projection"], +) +def test_pushdown_does_not_inherit_stale_plan_files(push_down): + """A pushdown clone must re-plan its scan, not inherit the cache. + + ``apply_predicate`` and ``apply_projection`` shallow-copy the datasource, so + a plan-file cache populated beforehand -- by ``estimate_inmemory_data_size``, + say, which Ray calls to autodetect parallelism -- would be shared with the + clone. Those tasks were planned without the predicate, so every residual is + ``AlwaysTrue`` and ``get_read_tasks`` would report the unfiltered manifest + total: exactly the overcount this file's other tests pin against. + + Reading the cache twice also has to keep working. ``DataScan.plan_files`` is + annotated ``Iterable[FileScanTask]``, so a future PyIceberg returning a + generator would otherwise leave the second read empty -- silently yielding + zero read tasks and an empty dataset. + """ + iceberg_ds = IcebergDatasource( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + ) + + # Warm the cache before pushdown, and read it twice. + assert iceberg_ds.estimate_inmemory_data_size() > 0 + assert iceberg_ds.estimate_inmemory_data_size() > 0, "cache must be re-readable" + unfiltered_files = len(iceberg_ds.plan_files) + assert unfiltered_files > 0, "fixture should plan at least one file" + + filtered_ds = push_down(iceberg_ds) + read_tasks = filtered_ds.get_read_tasks(2) + assert read_tasks, "pushdown must not leave the clone with an exhausted cache" + + claimed = [read_task.metadata.num_rows for read_task in read_tasks] + actual = sum(block.num_rows for read_task in read_tasks for block in read_task()) + + assert actual == 10, f"filter should match 10 of the fixture's rows, got {actual}" + if all(count is not None for count in claimed): + assert ( + sum(claimed) == actual + ), "reported row count came from plan files that predate the predicate" + + @pytest.mark.skipif( get_pyarrow_version() < parse_version("14.0.0"), reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", From 3e99152019149f865afac99247ac7f3ac3633819 Mon Sep 17 00:00:00 2001 From: Aarrya Date: Mon, 3 Aug 2026 10:41:56 -0700 Subject: [PATCH 4/6] [Data] Iceberg: a scan limit is not exact, and pin the sentinel to the snapshot Two more ways the row count could be claimed but not delivered. A ``scan_kwargs`` ``limit`` makes the read stop early, but ``counts_are_exact`` never consulted it: with no row filter every residual is ``AlwaysTrue``, so the manifests' full total was reported for a read that returns ``limit`` rows. Measured 100 claimed against 5 delivered. Report an unknown count instead. (``limit`` is handed to each read task rather than applied across the scan, so the rows returned also depend on how files were chunked -- another reason not to try to predict it. That per-task behaviour is pre-existing and left alone.) ``_get_sentinel_field`` chose the stand-in column for an empty projection from ``Table.schema()``, the current schema, even when ``snapshot_id`` pins an older snapshot. After a column is added, that name does not exist in the pinned snapshot and PyIceberg raises ``ValueError: Could not find column``, so ``count()`` on a historical snapshot failed outright. Choose from ``DataScan.projection()`` instead, which resolves the snapshot's schema itself, so we agree with the scan we are about to build. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Aarrya --- .../datasource/iceberg_datasource.py | 37 ++++++++-- .../ray/data/tests/datasource/test_iceberg.py | 69 +++++++++++++++++++ 2 files changed, 102 insertions(+), 4 deletions(-) diff --git a/python/ray/data/_internal/datasource/iceberg_datasource.py b/python/ray/data/_internal/datasource/iceberg_datasource.py index 4b63a3a9a376..b55276aae343 100644 --- a/python/ray/data/_internal/datasource/iceberg_datasource.py +++ b/python/ray/data/_internal/datasource/iceberg_datasource.py @@ -450,8 +450,14 @@ def _get_sentinel_field(self) -> str: which preserves the row count. Picking the narrowest column keeps the wasted I/O small: reading a string column here instead of a boolean one can cost several orders of magnitude more bytes. + + The column is chosen from the schema the scan will actually read, not + from ``self.table.schema()``. Those differ when ``snapshot_id`` pins a + historical snapshot: after a column is added, the current schema holds + names that do not exist in the pinned snapshot, and selecting one raises + ``ValueError: Could not find column``. """ - fields = self.table.schema().fields + fields = self._get_scan_schema().fields if not fields: raise ValueError( f"Cannot read table '{self.table_identifier}': it has no columns." @@ -466,6 +472,21 @@ def _get_sentinel_field(self) -> str: ) return cheapest_field.name + def _get_scan_schema(self) -> "Schema": + """Return the full schema this scan reads against. + + Delegates snapshot resolution to PyIceberg -- ``DataScan.projection()`` + returns the pinned snapshot's schema rather than the table's current one, + and returns it whole when every field is selected -- so that we agree + with the scan built in ``_get_data_scan``. Builds a scan but does not + plan it, which touches no data files. + """ + return self.table.scan( + row_filter=self._get_combined_filter(), + selected_fields=("*",), + **self._scan_kwargs, + ).projection() + def _is_empty_projection(self) -> bool: """Whether the pushed-down projection selects no columns at all.""" return self._get_data_columns() == [] @@ -599,9 +620,17 @@ def get_read_tasks( # # This is deliberately the same test PyIceberg's own ``DataScan.count()`` # applies, so the two agree on when a manifest count still holds. - counts_are_exact = isinstance(row_filter, AlwaysTrue) or all( - isinstance(getattr(task, "residual", None), AlwaysTrue) - for task in plan_files + # + # A ``limit`` makes the read stop early, so no manifest arithmetic + # describes what comes back and the count cannot be exact. (It is handed + # to each read task, so the rows returned depend on how files were + # chunked -- another reason not to try to predict it.) + counts_are_exact = limit is None and ( + isinstance(row_filter, AlwaysTrue) + or all( + isinstance(getattr(task, "residual", None), AlwaysTrue) + for task in plan_files + ) ) get_read_task = partial( diff --git a/python/ray/data/tests/datasource/test_iceberg.py b/python/ray/data/tests/datasource/test_iceberg.py index b685c0d11937..c8087c33466f 100644 --- a/python/ray/data/tests/datasource/test_iceberg.py +++ b/python/ray/data/tests/datasource/test_iceberg.py @@ -370,6 +370,75 @@ def test_reported_num_rows_matches_rows_read(row_filter, count_must_be_exact): assert sum(claimed) == actual +@pytest.mark.skipif( + get_pyarrow_version() < parse_version("14.0.0"), + reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", +) +def test_reported_num_rows_is_unknown_when_scan_stops_early(): + """A ``limit`` makes the read stop early, so no manifest count describes it. + + The manifests still say how many rows each surviving file holds, and the + filter may well be fully resolved, so the exactness test would otherwise pass + and report the full total for a read that returns ``limit`` rows. + """ + iceberg_ds = IcebergDatasource( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + # No row filter, so residuals are all ``AlwaysTrue``: the limit is the + # only reason the count cannot be trusted. + scan_kwargs={"limit": 5}, + ) + read_tasks = iceberg_ds.get_read_tasks(1) + + claimed = [read_task.metadata.num_rows for read_task in read_tasks] + actual = sum(block.num_rows for read_task in read_tasks for block in read_task()) + + assert actual == 5, f"limit should cap a single task at 5 rows, got {actual}" + assert all( + count is None for count in claimed + ), f"row counts cannot be exact under a limit, got {claimed}" + + +@pytest.mark.skipif( + get_pyarrow_version() < parse_version("14.0.0"), + reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", +) +def test_empty_projection_respects_pinned_snapshot_schema(): + """The stand-in column must exist in the snapshot being read. + + An empty projection reads one cheap column and slices it away. Choosing that + column from the table's *current* schema breaks a read pinned to an older + snapshot: a column added since does not exist there, and PyIceberg raises + ``ValueError: Could not find column``. So ``count()`` on a historical + snapshot would fail outright rather than count its rows. + """ + sql_catalog = pyi_catalog.load_catalog(**_CATALOG_KWARGS.copy()) + table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}") + old_snapshot_id = table.current_snapshot().snapshot_id + expected_rows = 101 # the fixture appends 120 rows, then deletes col_a >= 101 + + # Evolve the schema by adding a boolean, which is the *cheapest* column in + # the table and so the one the sentinel picks -- and which does not exist in + # the snapshot above. + with table.update_schema() as update: + update.add_column("col_d", pyi_types.BooleanType()) + table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}") + assert "col_d" in table.schema().column_names, "schema evolution should apply" + + iceberg_ds = IcebergDatasource( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + catalog_kwargs=_CATALOG_KWARGS.copy(), + snapshot_id=old_snapshot_id, + ) + assert ( + iceberg_ds._get_sentinel_field() != "col_d" + ), "the stand-in column must come from the pinned snapshot's schema" + + read_tasks = iceberg_ds.apply_projection({}).get_read_tasks(2) + actual = sum(block.num_rows for read_task in read_tasks for block in read_task()) + assert actual == expected_rows + + @pytest.mark.skipif( get_pyarrow_version() < parse_version("14.0.0"), reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", From 1bfe52ce303dbc2534bbf3fa61f3ebd4319d93ab Mon Sep 17 00:00:00 2001 From: Aarrya Date: Mon, 3 Aug 2026 14:22:53 -0700 Subject: [PATCH 5/6] [Data] Iceberg: trim comments, rename sentinel helper, tighten types Follow-up cleanup on the count() fixes, no behaviour change: - Trim the added comments and docstrings (96 -> 71 prose lines in the datasource), keeping the "why the obvious thing is wrong" reasoning and dropping the restatements. - Rename _get_sentinel_field to _get_cheapest_sentinel_field and update its references. - Annotate _invalidate_plan_files as taking and returning IcebergDatasource. It assigns _plan_files, which only exists on this subclass, and both callers declare IcebergDatasource returns. - Use float("inf") for _UNKNOWN_FIELD_WIDTH instead of 1 << 20; the value is only ever a sort key and the comment already says "unbounded". - Hoist the num_rows conditional out of the BlockMetadata(...) call. Tested: python/ray/data/tests/datasource/test_iceberg.py, 70 passed (macOS, Python 3.12, ray 2.56.1 with the changed files applied, pyiceberg 0.11.1, pyarrow 25.0.0). ruff and black clean. AI assistance was used for this change. Signed-off-by: Aarrya --- .../datasource/iceberg_datasource.py | 132 +++++++----------- .../ray/data/tests/datasource/test_iceberg.py | 70 ++++------ 2 files changed, 85 insertions(+), 117 deletions(-) diff --git a/python/ray/data/_internal/datasource/iceberg_datasource.py b/python/ray/data/_internal/datasource/iceberg_datasource.py index b55276aae343..037fadce3360 100644 --- a/python/ray/data/_internal/datasource/iceberg_datasource.py +++ b/python/ray/data/_internal/datasource/iceberg_datasource.py @@ -86,12 +86,11 @@ logger = logging.getLogger(__name__) -# Approximate width, in bytes, of the Iceberg primitive types that are cheap to -# read. Used to choose a stand-in column for an empty projection, see -# ``IcebergDatasource._get_sentinel_field``. Keys are Iceberg type names as -# returned by ``str(field_type)``. Types missing from this map -- strings, -# binary, decimals and nested types -- have unbounded width, and are only used -# as a sentinel if the table has nothing cheaper to offer. +# Approximate width in bytes of each Iceberg primitive type, keyed by +# ``str(field_type)``. Used to pick the cheapest stand-in column for an empty +# projection, see ``IcebergDatasource._get_cheapest_sentinel_field``. Types +# absent here -- strings, binary, decimals, nested -- are unbounded, so they +# lose to any type listed. _SENTINEL_FIELD_WIDTHS = { "boolean": 1, "int": 4, @@ -280,12 +279,8 @@ def _generate_tables() -> Iterable[pa.Table]: for table in _generate_tables(): if drop_all_columns: - # ``schema`` holds a sentinel column the caller never asked for, see - # ``IcebergDatasource._get_sentinel_field``. Hand it to Ray's - # zero-column projection helper, which slices every column away - # (preserving the row count, unlike rebuilding a table from zero - # columns) and appends the ``__bsp_stub`` placeholder that keeps the - # count alive through downstream ``pa.concat_tables`` calls. + # We over-read one column so the row count survives; drop it here. + # See ``IcebergDatasource._get_cheapest_sentinel_field``. table = BlockAccessor.for_block(table).select([]) yield table @@ -329,10 +324,9 @@ def __init__( _check_import(self, module="pyiceberg", package="pyiceberg") from pyiceberg.expressions import AlwaysTrue - # Copy both dicts: below we pop ``name`` out of the catalog kwargs and - # write ``snapshot_id`` into the scan kwargs. Doing that in place would - # mutate dicts the caller still holds, so a second read reusing the same - # dict would fail to find its catalog or silently inherit a snapshot ID. + # Copy both dicts: we pop ``name`` from the catalog kwargs and write + # ``snapshot_id`` into the scan kwargs, and doing that in place would + # break a second read that reuses the caller's dict. self._scan_kwargs = dict(scan_kwargs) if scan_kwargs is not None else {} self._catalog_kwargs = ( dict(catalog_kwargs) if catalog_kwargs is not None else {} @@ -382,9 +376,8 @@ def plan_files(self) -> List["FileScanTask"]: # Calculate and cache the plan_files if they don't already exist if self._plan_files is None: data_scan = self._get_data_scan() - # ``DataScan.plan_files`` is annotated ``Iterable[FileScanTask]``. - # PyIceberg returns a list today, but callers here iterate the cache - # more than once, so materialize it rather than depend on that. + # Annotated ``Iterable[FileScanTask]``; a list today, but this cache + # is iterated more than once, so don't depend on that. self._plan_files = list(data_scan.plan_files()) return self._plan_files @@ -392,11 +385,10 @@ def plan_files(self) -> List["FileScanTask"]: def apply_predicate(self, predicate_expr: Expr) -> "IcebergDatasource": """Push a predicate down, discarding plan files planned without it. - The base implementation shallow-copies, which would leave the clone - sharing a plan-file cache computed before this predicate existed. Those - tasks carry ``AlwaysTrue`` residuals -- nothing was filtered at plan - time -- which ``get_read_tasks`` would read as "the filter is fully - resolved" and report an unfiltered row count for a filtered read. + The base implementation shallow-copies, so the clone would inherit a + cache planned before this predicate existed: all-``AlwaysTrue`` + residuals, which ``get_read_tasks`` turns into an unfiltered row count + for a filtered read. """ return self._invalidate_plan_files(super().apply_predicate(predicate_expr)) @@ -406,15 +398,13 @@ def apply_projection( """Push a projection down, discarding plan files planned without it. The projected columns are part of the scan, so a cache built for a - different projection does not belong to the clone. See - ``apply_predicate``. + different projection does not belong to the clone. See ``apply_predicate``. """ return self._invalidate_plan_files(super().apply_projection(projection_map)) - def _invalidate_plan_files(self, clone: "Datasource") -> "Datasource": + def _invalidate_plan_files(self, clone: "IcebergDatasource") -> "IcebergDatasource": """Drop ``clone``'s inherited plan-file cache, so it re-plans its scan.""" - # A datasource returns ``self`` to signal "no pushdown applied", in which - # case there is no clone and the cache is still valid for this scan. + # ``self`` signals "no pushdown applied": no clone, cache still valid. if clone is not self: clone._plan_files = None return clone @@ -438,24 +428,20 @@ def _get_combined_filter(self) -> "BooleanExpression": return combined_filter - def _get_sentinel_field(self) -> str: + def _get_cheapest_sentinel_field(self) -> str: """Return the cheapest column to read in place of an empty projection. - PyIceberg builds its output by reconstructing a table against the - projected schema, and a table reconstructed from zero columns reports - zero rows however many rows the scan actually matched. So an empty - projection -- which ``Dataset.count()`` legitimately requests, since it - needs a row count but no values -- would silently count zero rows. We - read a single column instead and slice it away in ``_get_read_task``, - which preserves the row count. Picking the narrowest column keeps the - wasted I/O small: reading a string column here instead of a boolean one - can cost several orders of magnitude more bytes. - - The column is chosen from the schema the scan will actually read, not - from ``self.table.schema()``. Those differ when ``snapshot_id`` pins a - historical snapshot: after a column is added, the current schema holds - names that do not exist in the pinned snapshot, and selecting one raises - ``ValueError: Could not find column``. + PyIceberg reconstructs its output against the projected schema, and a + table reconstructed from zero columns reports zero rows however many the + scan matched -- so an empty projection, which ``Dataset.count()`` + legitimately requests, would silently count zero. Read one column + instead and slice it away in ``_get_read_task``, which preserves the + count. Pick the narrowest: a string column can cost orders of magnitude + more bytes than a boolean one. + + The column comes from the schema the scan reads, not + ``self.table.schema()``, which with ``snapshot_id`` pinned can name + columns the snapshot lacks -- ``ValueError: Could not find column``. """ fields = self._get_scan_schema().fields if not fields: @@ -475,11 +461,9 @@ def _get_sentinel_field(self) -> str: def _get_scan_schema(self) -> "Schema": """Return the full schema this scan reads against. - Delegates snapshot resolution to PyIceberg -- ``DataScan.projection()`` - returns the pinned snapshot's schema rather than the table's current one, - and returns it whole when every field is selected -- so that we agree - with the scan built in ``_get_data_scan``. Builds a scan but does not - plan it, which touches no data files. + ``DataScan.projection()`` resolves the pinned snapshot's schema itself, + and returns it whole when every field is selected, so this agrees with + ``_get_data_scan`` by construction. Building a scan reads no data files. """ return self.table.scan( row_filter=self._get_combined_filter(), @@ -500,7 +484,7 @@ def _get_data_scan(self) -> "DataScan": if data_columns is None: selected_fields = ("*",) elif not data_columns: - selected_fields = (self._get_sentinel_field(),) + selected_fields = (self._get_cheapest_sentinel_field(),) else: selected_fields = tuple(data_columns) @@ -577,10 +561,10 @@ def get_read_tasks( # Get the arrow schema, to set in the metadata pya_schema = pyi_pa_io.schema_to_pyarrow(projected_schema) - # An empty projection is read as a single sentinel column, see - # ``_get_sentinel_field``, which the read tasks then replace with Ray's - # ``__bsp_stub`` placeholder. Declare that placeholder so the reported - # schema matches the blocks; it is hidden from the user-visible schema. + # An empty projection is read as one sentinel column, see + # ``_get_cheapest_sentinel_field``, which the read tasks replace with Ray's + # ``__bsp_stub``. Declare the placeholder so the reported schema matches + # the blocks; it is hidden from the user-visible schema. drop_all_columns = self._is_empty_projection() if drop_all_columns: pya_schema = pa.schema( @@ -607,24 +591,13 @@ def get_read_tasks( case_sensitive = self._scan_kwargs.get("case_sensitive", True) limit = self._scan_kwargs.get("limit") - # Whether the manifests' ``record_count`` is still an exact row count. - # It is written when the data is written, so it only holds if every row - # of every surviving file matches the filter. PyIceberg reports whatever - # is left of the filter after partition pruning as - # ``FileScanTask.residual``: ``AlwaysTrue`` means pruning resolved the - # filter completely (e.g. a filter on a partition column dropped whole - # files), so the counts still hold. Anything else has to be evaluated row - # by row, and the counts would over-report. Old PyIceberg versions do not - # populate ``residual``, in which case we cannot prove anything and fall - # back to reporting an unknown row count. - # - # This is deliberately the same test PyIceberg's own ``DataScan.count()`` - # applies, so the two agree on when a manifest count still holds. - # - # A ``limit`` makes the read stop early, so no manifest arithmetic - # describes what comes back and the count cannot be exact. (It is handed - # to each read task, so the rows returned depend on how files were - # chunked -- another reason not to try to predict it.) + # Manifests record how many rows each file holds, which is only still an + # exact count if every surviving row matches the filter. PyIceberg puts + # whatever the filter leaves after file pruning in + # ``FileScanTask.residual``, so ``AlwaysTrue`` means nothing is left to + # check and the counts hold. Same test PyIceberg's own + # ``DataScan.count()`` uses. A ``limit`` stops the read early, so no + # count taken from the manifests describes what comes back. counts_are_exact = limit is None and ( isinstance(row_filter, AlwaysTrue) or all( @@ -662,13 +635,16 @@ def get_read_tasks( for delete in unique_deletes if delete.content == DataFileContent.POSITION_DELETES ) - metadata = BlockMetadata( - num_rows=( + # ``None`` means "unknown", which makes ``Dataset.count()`` do the + # read instead of trusting these manifest counts. + num_rows = None + if counts_are_exact: + num_rows = ( sum(task.file.record_count for task in chunk_tasks) - position_delete_count - if counts_are_exact - else None - ), + ) + metadata = BlockMetadata( + num_rows=num_rows, size_bytes=sum(task.file.file_size_in_bytes for task in chunk_tasks), input_files=[task.file.file_path for task in chunk_tasks], exec_stats=None, diff --git a/python/ray/data/tests/datasource/test_iceberg.py b/python/ray/data/tests/datasource/test_iceberg.py index c8087c33466f..25f56de23c4f 100644 --- a/python/ray/data/tests/datasource/test_iceberg.py +++ b/python/ray/data/tests/datasource/test_iceberg.py @@ -230,11 +230,9 @@ def _iceberg_scan_row_count(**scan_kwargs) -> int: return table.scan(**scan_kwargs).to_arrow().num_rows -# Every entry builds a Dataset whose row count is unambiguous, so ``count()`` can -# be checked against the rows the same query actually yields. ``count()`` has two -# strategies -- read the count from plan metadata, or project to zero columns and -# count what comes back -- and these cases cover both, including the cases that -# defeat the metadata short-circuit. +# ``count()`` has two strategies -- read the count from plan metadata, or project +# to zero columns and count what comes back. These cases cover both, including +# the ones that defeat the metadata short-circuit. _COUNT_CASES = { "plain": lambda: read_iceberg( table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", @@ -288,8 +286,8 @@ def _iceberg_scan_row_count(**scan_kwargs) -> int: def test_count_matches_rows_actually_produced(case): """``count()`` must agree with the number of rows the query yields. - A projection or a pushed-down filter must not change the answer: the count is - a property of the query, not of how much of it was pushed into the reader. + The count is a property of the query, not of how much of it was pushed into + the reader. """ make_ds = _COUNT_CASES[case] expected = len(make_ds().take_all()) @@ -304,8 +302,7 @@ def test_empty_projection_preserves_row_count(): """Selecting zero columns must yield N rows of no columns, not zero rows. ``Dataset.count()`` projects to zero columns to avoid reading column data, so - an empty projection that drops rows silently corrupts every count built that - way. + an empty projection that drops rows corrupts every count built that way. """ iceberg_ds = IcebergDatasource( table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", @@ -344,11 +341,10 @@ def test_empty_projection_preserves_row_count(): def test_reported_num_rows_matches_rows_read(row_filter, count_must_be_exact): """Read-task metadata must never claim a row count the read does not deliver. - ``Dataset.count()`` returns this number directly when the plan has nothing - that could change the row count, so an overcount here reaches the user as - the answer. ``None`` means "unknown" and makes ``count()`` do the read, so it - is always safe -- but it also gives up a free count, hence - ``count_must_be_exact`` pinning the cases where the shortcut is sound. + ``Dataset.count()`` returns this number directly when nothing in the plan can + change the row count, so an overcount reaches the user as the answer. ``None`` + is always safe but gives up a free count, hence ``count_must_be_exact`` + pinning the cases where the shortcut is sound. """ kwargs = {} if row_filter is None else {"row_filter": row_filter} iceberg_ds = IcebergDatasource( @@ -377,9 +373,9 @@ def test_reported_num_rows_matches_rows_read(row_filter, count_must_be_exact): def test_reported_num_rows_is_unknown_when_scan_stops_early(): """A ``limit`` makes the read stop early, so no manifest count describes it. - The manifests still say how many rows each surviving file holds, and the - filter may well be fully resolved, so the exactness test would otherwise pass - and report the full total for a read that returns ``limit`` rows. + The manifests still say how many rows each file holds and the filter may be + fully resolved, so the exactness test would otherwise pass and report the + full total for a read that returns ``limit`` rows. """ iceberg_ds = IcebergDatasource( table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", @@ -406,11 +402,10 @@ def test_reported_num_rows_is_unknown_when_scan_stops_early(): def test_empty_projection_respects_pinned_snapshot_schema(): """The stand-in column must exist in the snapshot being read. - An empty projection reads one cheap column and slices it away. Choosing that - column from the table's *current* schema breaks a read pinned to an older - snapshot: a column added since does not exist there, and PyIceberg raises - ``ValueError: Could not find column``. So ``count()`` on a historical - snapshot would fail outright rather than count its rows. + An empty projection reads one cheap column and slices it away. Choosing it + from the table's *current* schema breaks a read pinned to an older snapshot: + a column added since is absent there, PyIceberg raises ``ValueError: Could + not find column``, and ``count()`` fails instead of counting. """ sql_catalog = pyi_catalog.load_catalog(**_CATALOG_KWARGS.copy()) table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}") @@ -431,7 +426,7 @@ def test_empty_projection_respects_pinned_snapshot_schema(): snapshot_id=old_snapshot_id, ) assert ( - iceberg_ds._get_sentinel_field() != "col_d" + iceberg_ds._get_cheapest_sentinel_field() != "col_d" ), "the stand-in column must come from the pinned snapshot's schema" read_tasks = iceberg_ds.apply_projection({}).get_read_tasks(2) @@ -456,17 +451,15 @@ def test_empty_projection_respects_pinned_snapshot_schema(): def test_pushdown_does_not_inherit_stale_plan_files(push_down): """A pushdown clone must re-plan its scan, not inherit the cache. - ``apply_predicate`` and ``apply_projection`` shallow-copy the datasource, so - a plan-file cache populated beforehand -- by ``estimate_inmemory_data_size``, - say, which Ray calls to autodetect parallelism -- would be shared with the - clone. Those tasks were planned without the predicate, so every residual is - ``AlwaysTrue`` and ``get_read_tasks`` would report the unfiltered manifest - total: exactly the overcount this file's other tests pin against. - - Reading the cache twice also has to keep working. ``DataScan.plan_files`` is - annotated ``Iterable[FileScanTask]``, so a future PyIceberg returning a - generator would otherwise leave the second read empty -- silently yielding - zero read tasks and an empty dataset. + ``apply_predicate`` and ``apply_projection`` shallow-copy the datasource, so a + cache populated beforehand -- by ``estimate_inmemory_data_size``, say, which + Ray calls to autodetect parallelism -- is shared with the clone. Those tasks + were planned without the predicate, so every residual is ``AlwaysTrue`` and + ``get_read_tasks`` would report the unfiltered manifest total. + + Reading the cache twice also has to keep working: ``plan_files`` is annotated + ``Iterable``, so a future PyIceberg returning a generator would otherwise + leave the second read empty -- zero read tasks, empty dataset. """ iceberg_ds = IcebergDatasource( table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", @@ -500,11 +493,10 @@ def test_pushdown_does_not_inherit_stale_plan_files(push_down): def test_read_iceberg_does_not_mutate_caller_kwargs(): """The caller's dicts belong to the caller. - Both dicts were stored by reference. ``catalog_kwargs`` is then - ``pop("name")``-ed, so the caller loses the catalog name and a second read of - the same table raises ``NoSuchTableError``. ``scan_kwargs`` has - ``snapshot_id`` written into it, which fails silently instead: a later read - reusing the dict inherits the pin and returns rows from a stale snapshot. + Both were stored by reference. ``catalog_kwargs`` is then ``pop("name")``-ed, + so a second read of the same table raises ``NoSuchTableError``. + ``scan_kwargs`` gets ``snapshot_id`` written into it, which fails silently + instead: a later read reusing the dict inherits the pin and returns stale rows. """ catalog_kwargs = _CATALOG_KWARGS.copy() scan_kwargs = {} From 6e311e04a2ec1311293347f21aa122420e5523dc Mon Sep 17 00:00:00 2001 From: Aarrya Date: Thu, 6 Aug 2026 16:29:18 -0700 Subject: [PATCH 6/6] [Data] Iceberg: read a fabricated stub column for an empty projection An empty projection previously read the cheapest real column and dropped it. That needed a type-width table to score columns, a helper to pick the winner, and another to fetch the schema to pick from -- plus a fix for the chosen column being absent from a pinned snapshot. Project a single optional field that no data file has instead. Iceberg null-fills a missing optional field, the same way it reads a file written before a column was added, so the row count arrives without reading any real column data. Naming the field ``__bsp_stub`` makes the resulting block match what every other reader already produces for an empty projection, so Iceberg now satisfies the same contract as the Parquet readers. This removes the width table, both helpers, one branch of _get_data_scan, and the schema-dependence itself: the fabricated field belongs to no schema version, so no schema change can reach it. The former pinned-snapshot test now guards that directly. Iceberg matches columns by field ID, so the fabricated field's ID must not be a real column's. A collision is not silent: measured on pyiceberg 0.11.1, one with a non-boolean column raises ResolveError ("Cannot promote long to boolean"), and only one with a boolean column would read real data. Verified against pyiceberg 0.8.1 (legacy project_table) and 0.11.1 (ArrowScan), with row filters, position deletes and empty results. Tests: python/ray/data/tests/datasource/test_iceberg.py -- 71 passed. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Aarrya --- .../datasource/iceberg_datasource.py | 136 ++++++++---------- .../ray/data/tests/datasource/test_iceberg.py | 51 +++++-- 2 files changed, 97 insertions(+), 90 deletions(-) diff --git a/python/ray/data/_internal/datasource/iceberg_datasource.py b/python/ray/data/_internal/datasource/iceberg_datasource.py index 037fadce3360..68a19de346a4 100644 --- a/python/ray/data/_internal/datasource/iceberg_datasource.py +++ b/python/ray/data/_internal/datasource/iceberg_datasource.py @@ -86,24 +86,43 @@ logger = logging.getLogger(__name__) -# Approximate width in bytes of each Iceberg primitive type, keyed by -# ``str(field_type)``. Used to pick the cheapest stand-in column for an empty -# projection, see ``IcebergDatasource._get_cheapest_sentinel_field``. Types -# absent here -- strings, binary, decimals, nested -- are unbounded, so they -# lose to any type listed. -_SENTINEL_FIELD_WIDTHS = { - "boolean": 1, - "int": 4, - "float": 4, - "date": 4, - "long": 8, - "double": 8, - "time": 8, - "timestamp": 8, - "timestamptz": 8, - "uuid": 16, -} -_UNKNOWN_FIELD_WIDTH = 1 << 20 +# Field ID for the fabricated stub field, see ``_get_empty_projection_schema``. +# Iceberg matches columns by ID, so this must not be a real column's: colliding +# with a non-boolean column raises ``ResolveError``, and with a boolean column +# reads that column's data instead (the row count is still right, but the read +# is no longer free). Catalogs assign IDs sequentially from 1, so a value this +# large will not collide; Iceberg reserves 2147483546 and above for metadata +# columns, so stay below that. +_STUB_FIELD_ID = 2147483000 + + +def _get_empty_projection_schema() -> "Schema": + """Return a projected schema that reads no data but keeps the row count. + + PyIceberg rebuilds its output to match the projected schema, and a table + rebuilt from zero fields reports zero rows however many the scan matched -- + so an empty projection, which ``Dataset.count()`` legitimately requests, + would silently count zero. Project a single optional field that no data file + has instead: Iceberg fills a missing optional field with nulls, the same way + it reads a file written before a column was added, so the row count arrives + in a column that costs nothing to read. + + The field is named after Ray's stub column so the resulting block matches + what every other reader produces for an empty projection. ``boolean`` is + used because Iceberg gained a null-valued type only after 0.9.0; + ``_get_read_task`` retypes the column to ``null`` on the way out. + """ + from pyiceberg.schema import Schema + from pyiceberg.types import BooleanType, NestedField + + return Schema( + NestedField( + field_id=_STUB_FIELD_ID, + name=_BATCH_SIZE_PRESERVING_STUB_COL_NAME, + field_type=BooleanType(), + required=False, + ) + ) class _IcebergExpressionVisitor( @@ -231,7 +250,7 @@ def _get_read_task( case_sensitive: bool, limit: Optional[int], schema: "Schema", - drop_all_columns: bool = False, + empty_projection: bool = False, ) -> Iterable[Block]: # Determine the PyIceberg version to handle backward compatibility import pyiceberg @@ -278,9 +297,11 @@ def _generate_tables() -> Iterable[pa.Table]: yield table for table in _generate_tables(): - if drop_all_columns: - # We over-read one column so the row count survives; drop it here. - # See ``IcebergDatasource._get_cheapest_sentinel_field``. + if empty_projection: + # The scan read a boolean-typed stub, see + # ``_get_empty_projection_schema``. Re-derive it through Ray's own + # empty projection so the column is ``null``-typed, matching the + # schema we report and every other reader's stub. table = BlockAccessor.for_block(table).select([]) yield table @@ -428,49 +449,6 @@ def _get_combined_filter(self) -> "BooleanExpression": return combined_filter - def _get_cheapest_sentinel_field(self) -> str: - """Return the cheapest column to read in place of an empty projection. - - PyIceberg reconstructs its output against the projected schema, and a - table reconstructed from zero columns reports zero rows however many the - scan matched -- so an empty projection, which ``Dataset.count()`` - legitimately requests, would silently count zero. Read one column - instead and slice it away in ``_get_read_task``, which preserves the - count. Pick the narrowest: a string column can cost orders of magnitude - more bytes than a boolean one. - - The column comes from the schema the scan reads, not - ``self.table.schema()``, which with ``snapshot_id`` pinned can name - columns the snapshot lacks -- ``ValueError: Could not find column``. - """ - fields = self._get_scan_schema().fields - if not fields: - raise ValueError( - f"Cannot read table '{self.table_identifier}': it has no columns." - ) - # ``min`` is stable, so an all-equal-width schema deterministically - # yields the first field. - cheapest_field = min( - fields, - key=lambda field: _SENTINEL_FIELD_WIDTHS.get( - str(field.field_type), _UNKNOWN_FIELD_WIDTH - ), - ) - return cheapest_field.name - - def _get_scan_schema(self) -> "Schema": - """Return the full schema this scan reads against. - - ``DataScan.projection()`` resolves the pinned snapshot's schema itself, - and returns it whole when every field is selected, so this agrees with - ``_get_data_scan`` by construction. Building a scan reads no data files. - """ - return self.table.scan( - row_filter=self._get_combined_filter(), - selected_fields=("*",), - **self._scan_kwargs, - ).projection() - def _is_empty_projection(self) -> bool: """Whether the pushed-down projection selects no columns at all.""" return self._get_data_columns() == [] @@ -479,12 +457,12 @@ def _get_data_scan(self) -> "DataScan": # Get the combined filter combined_filter = self._get_combined_filter() - # Convert back to tuple for PyIceberg API (None -> ("*",)) + # Convert back to tuple for PyIceberg API (None -> ("*",)). An empty + # projection also scans everything: it selects its own columns through + # the projected schema instead, see ``_get_empty_projection_schema``. data_columns = self._get_data_columns() - if data_columns is None: + if not data_columns: selected_fields = ("*",) - elif not data_columns: - selected_fields = (self._get_cheapest_sentinel_field(),) else: selected_fields = tuple(data_columns) @@ -561,12 +539,13 @@ def get_read_tasks( # Get the arrow schema, to set in the metadata pya_schema = pyi_pa_io.schema_to_pyarrow(projected_schema) - # An empty projection is read as one sentinel column, see - # ``_get_cheapest_sentinel_field``, which the read tasks replace with Ray's - # ``__bsp_stub``. Declare the placeholder so the reported schema matches - # the blocks; it is hidden from the user-visible schema. - drop_all_columns = self._is_empty_projection() - if drop_all_columns: + # An empty projection reads a fabricated stub column instead of none at + # all, see ``_get_empty_projection_schema``. Declare the placeholder so + # the reported schema matches the blocks; it is hidden from the + # user-visible schema. + empty_projection = self._is_empty_projection() + if empty_projection: + projected_schema = _get_empty_projection_schema() pya_schema = pa.schema( [pa.field(_BATCH_SIZE_PRESERVING_STUB_COL_NAME, pa.null())] ) @@ -598,6 +577,11 @@ def get_read_tasks( # check and the counts hold. Same test PyIceberg's own # ``DataScan.count()`` uses. A ``limit`` stops the read early, so no # count taken from the manifests describes what comes back. + # + # The subtraction below only accounts for position deletes, but an + # equality delete cannot reach here: ``plan_files`` rejects them in both + # of its paths, locally in ``_plan_files_local`` and over REST in + # ``FileScanTask.from_rest_response``. counts_are_exact = limit is None and ( isinstance(row_filter, AlwaysTrue) or all( @@ -614,7 +598,7 @@ def get_read_tasks( case_sensitive=case_sensitive, limit=limit, schema=projected_schema, - drop_all_columns=drop_all_columns, + empty_projection=empty_projection, ) read_tasks = [] diff --git a/python/ray/data/tests/datasource/test_iceberg.py b/python/ray/data/tests/datasource/test_iceberg.py index 25f56de23c4f..cd86a5ef9596 100644 --- a/python/ray/data/tests/datasource/test_iceberg.py +++ b/python/ray/data/tests/datasource/test_iceberg.py @@ -21,6 +21,7 @@ import ray from ray.data import read_iceberg +from ray.data._internal.arrow_block import _BATCH_SIZE_PRESERVING_STUB_COL_NAME from ray.data._internal.datasource.iceberg_datasource import IcebergDatasource from ray.data._internal.logical.operators import Filter, Project from ray.data._internal.logical.optimizers import LogicalOptimizer @@ -320,6 +321,34 @@ def test_empty_projection_preserves_row_count(): assert rows_read == expected +@pytest.mark.skipif( + get_pyarrow_version() < parse_version("14.0.0"), + reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", +) +def test_empty_projection_reads_only_the_stub_column(): + """An empty projection must read no real column, only the stub. + + It projects a field present in no schema version, relying on PyIceberg + null-filling a missing optional field. Should an upgrade stop doing that, + the row count collapses to zero silently, so pin the resulting schema: one + ``null``-typed stub column, matching what every other reader produces. + """ + iceberg_ds = IcebergDatasource( + table_identifier=f"{_DB_NAME}.{_TABLE_NAME}", + selected_fields=(), + catalog_kwargs=_CATALOG_KWARGS.copy(), + ) + blocks = [ + block for read_task in iceberg_ds.get_read_tasks(1) for block in read_task() + ] + assert blocks, "fixture table should produce at least one block" + expected_schema = pa.schema( + [pa.field(_BATCH_SIZE_PRESERVING_STUB_COL_NAME, pa.null())] + ) + for block in blocks: + assert block.schema.equals(expected_schema), block.schema + + @pytest.mark.skipif( get_pyarrow_version() < parse_version("14.0.0"), reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", @@ -399,22 +428,20 @@ def test_reported_num_rows_is_unknown_when_scan_stops_early(): get_pyarrow_version() < parse_version("14.0.0"), reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0", ) -def test_empty_projection_respects_pinned_snapshot_schema(): - """The stand-in column must exist in the snapshot being read. - - An empty projection reads one cheap column and slices it away. Choosing it - from the table's *current* schema breaks a read pinned to an older snapshot: - a column added since is absent there, PyIceberg raises ``ValueError: Could - not find column``, and ``count()`` fails instead of counting. +def test_empty_projection_survives_schema_evolution_on_pinned_snapshot(): + """An empty projection must not depend on the table's current schema. + + Reading a pinned snapshot after a column was added is where naming a real + stand-in column breaks: picked from the *current* schema it is absent from + the snapshot, PyIceberg raises ``ValueError: Could not find column``, and + ``count()`` fails instead of counting. The fabricated stub belongs to no + schema version, so no schema change can reach it. """ sql_catalog = pyi_catalog.load_catalog(**_CATALOG_KWARGS.copy()) table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}") old_snapshot_id = table.current_snapshot().snapshot_id expected_rows = 101 # the fixture appends 120 rows, then deletes col_a >= 101 - # Evolve the schema by adding a boolean, which is the *cheapest* column in - # the table and so the one the sentinel picks -- and which does not exist in - # the snapshot above. with table.update_schema() as update: update.add_column("col_d", pyi_types.BooleanType()) table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}") @@ -425,10 +452,6 @@ def test_empty_projection_respects_pinned_snapshot_schema(): catalog_kwargs=_CATALOG_KWARGS.copy(), snapshot_id=old_snapshot_id, ) - assert ( - iceberg_ds._get_cheapest_sentinel_field() != "col_d" - ), "the stand-in column must come from the pinned snapshot's schema" - read_tasks = iceberg_ds.apply_projection({}).get_read_tasks(2) actual = sum(block.num_rows for read_task in read_tasks for block in read_task()) assert actual == expected_rows