diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index c45f88ec09c0..a9bd84d38d7a 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -571,9 +571,14 @@ ds = read_by_row_id( target row ids in column `row_id_col`; other columns are ignored. A table-name source is not accepted (a table's system `_ROW_ID` is its own and cannot address the target). - `projection`: top-level columns to read (nested paths are not supported). Blob columns - are resolved to their payloads. Must be non-empty. + are resolved to their payloads, unless overridden via `dynamic_options`. Must be non-empty. - `row_id_col`: the source column holding the row ids (default `_ROW_ID`); set e.g. `row_id_col="row_id"` to consume a `bucket_join` locator directly. +- `dynamic_options`: read options applied via `table.copy`, e.g. + `{"blob-as-descriptor": "true"}` to read blob columns as small `BlobDescriptor` bytes + (resolved later with `map_with_blobs`), or `scan.snapshot-id` / `scan.tag-name` to read a + specific snapshot. Options that flip table invariants (`data-evolution.enabled`, + `row-tracking.enabled`, `deletion-vectors.enabled`) are rejected. - `num_partitions`: parallelism for grouping the row ids by target file; defaults to `max(1, cluster_cpus * 2)`. - `ray_remote_args`: Ray remote options applied to the read tasks. @@ -584,9 +589,10 @@ ds = read_by_row_id( - Lookup/set semantics, like SQL `... WHERE _ROW_ID IN (...)`: one row per **distinct** matched row id (duplicates deduplicated), input order not preserved (rows come out grouped by owning file). An empty source yields an empty but correctly-typed Dataset. -- The row ids must exist in the target's current snapshot; a foreign `_ROW_ID` raises. +- The row ids must exist in the resolved target snapshot (latest, or the one selected via + `dynamic_options`); a foreign `_ROW_ID` raises. - Deletion-vectors-enabled tables are not supported yet, for the same reason as `update_by_row_id`. -- Prefer a materialized `row_ids` source (a `bucket_join` result already is one): the - emptiness check reads one block up front, which would otherwise re-run a lazy source's - first block. +- For a non-empty target, the `row_ids` source is consumed lazily by the downstream + action, not read here. A lazy source missing `row_id_col` raises when the read runs + (a materialized source raises up front). diff --git a/paimon-python/pypaimon/ray/read_by_row_id.py b/paimon-python/pypaimon/ray/read_by_row_id.py index b145823ba035..7c7a6508030b 100644 --- a/paimon-python/pypaimon/ray/read_by_row_id.py +++ b/paimon-python/pypaimon/ray/read_by_row_id.py @@ -48,6 +48,20 @@ def _empty_result(table: "FileStoreTable", read_cols: List[str]): return ray.data.from_arrow(_read_output_schema(table, read_cols).empty_table()) +def _read_snapshot(table): + """The snapshot to route/read on: a time-travel dynamic option if set, else the latest.""" + from pypaimon.snapshot.time_travel_util import SCAN_KEYS, TimeTravelUtil + + opts = table.options.options + if not any(opts.contains_key(k) for k in SCAN_KEYS): + return table.snapshot_manager().get_latest_snapshot() + snap = TimeTravelUtil.try_travel_to_snapshot( + opts, table.tag_manager(), table.snapshot_manager()) + if snap is None: + raise ValueError("could not resolve the time-travel snapshot from dynamic_options.") + return snap + + def read_by_row_id( target: str, row_ids: Any, @@ -55,6 +69,7 @@ def read_by_row_id( *, projection: List[str], row_id_col: Optional[str] = None, + dynamic_options: Optional[Dict[str, str]] = None, num_partitions: Optional[int] = None, ray_remote_args: Optional[Dict[str, Any]] = None, ): @@ -65,7 +80,11 @@ def read_by_row_id( e.g. ``row_id_col="row_id"`` for a ``bucket_join`` locator). Each row id is routed to the data file owning it and only those files -- and only the matched rows -- are read, so the target is never fully scanned and there is no join against it. - ``projection`` lists top-level columns; blob columns are resolved to their payloads. + ``projection`` lists top-level columns; blob columns resolve to payloads by default. + ``dynamic_options`` overrides read options via ``table.copy``: ``{"blob-as-descriptor": + "true"}`` for descriptor bytes (resolve with ``map_with_blobs``), or ``scan.snapshot-id`` / + ``scan.tag-name`` to read that snapshot. Options flipping table invariants + (``data-evolution.enabled`` etc.) are rejected. Requires ``ray >= 2.50`` and a target with ``data-evolution.enabled`` + ``row-tracking.enabled``. @@ -78,6 +97,7 @@ def read_by_row_id( Returns a ``ray.data.Dataset`` of ``(*projection, _ROW_ID)``. """ from pypaimon.catalog.catalog_factory import CatalogFactory + from pypaimon.snapshot.time_travel_util import SCAN_KEYS from pypaimon.table.special_fields import SpecialFields _require_ray_join() @@ -98,6 +118,16 @@ def read_by_row_id( raise ValueError( f"read_by_row_id does not support deletion-vectors-enabled tables yet: " f"'{target}'.") + if dynamic_options: + # Flipping these would bypass the checks above. + bad = sorted({"data-evolution.enabled", "row-tracking.enabled", + "deletion-vectors.enabled"} & set(dynamic_options)) + if bad: + raise ValueError(f"dynamic_options cannot override table invariants {bad}.") + # table.copy's _try_time_travel swallows the multi-key error, so reject it here. + if len([k for k in SCAN_KEYS if k in dynamic_options]) > 1: + raise ValueError(f"dynamic_options may set at most one time-travel key {SCAN_KEYS}.") + table = table.copy(dynamic_options) rid = SpecialFields.ROW_ID.name src_rid_col = row_id_col or rid @@ -111,27 +141,44 @@ def read_by_row_id( "read_by_row_id does not accept a table-name source; pass a ray.data." "Dataset / pyarrow.Table / pandas.DataFrame carrying the target row ids.") source_ds = _normalize_source(row_ids, catalog_options) - if src_rid_col not in set(source_ds.schema().names): + # Only check now if the schema is free; fetching it would execute a lazy source. + known_schema = source_ds.schema(fetch_if_missing=False) + if known_schema is not None and src_rid_col not in set(known_schema.names): raise ValueError(f"row_ids source is missing the {src_rid_col!r} column.") def _project_rid(batch: pa.Table) -> pa.Table: + if src_rid_col not in batch.column_names: + raise ValueError(f"row_ids source is missing the {src_rid_col!r} column.") return pa.table({rid: batch.column(src_rid_col).cast(pa.int64())}) rid_ds = source_ds.map_batches(_project_rid, batch_format="pyarrow") read_cols = list(projection) + ([rid] if rid not in projection else []) - # Empty source -> typed empty Dataset (a zero-row groupby has no schema). - source_empty = rid_ds.limit(1).count() == 0 - - base = table.snapshot_manager().get_latest_snapshot() + base = _read_snapshot(table) + if base is not None and dynamic_options and any(k in dynamic_options for k in SCAN_KEYS): + # A pre-row-tracking snapshot has files without row ids; fail clearly here rather + # than deep in the planner (the persisted-table check above cannot see this). + from pypaimon.common.options.core_options import CoreOptions + from pypaimon.common.options.options import Options + base_schema = table.schema_manager.get_schema(base.schema_id) + if not CoreOptions(Options(base_schema.options)).row_tracking_enabled(): + raise ValueError( + f"the resolved snapshot ({base.id}) predates row-tracking; read_by_row_id needs it.") # No DV (rejected above) -> total_record_count is the live row count; 0 = empty. if base is None or base.total_record_count == 0: - if not source_empty: + # Force an action on the source only in this degenerate branch (like update_by_row_id). + if rid_ds.limit(1).count() > 0: raise ValueError( f"target '{target}' has no rows; every _ROW_ID in the source is foreign.") return _empty_result(table, read_cols) - if source_empty: - return _empty_result(table, read_cols) + # base captures the resolved snapshot; reduce any time-travel key to a plain snapshot-id + # so the planner's own snapshot-id pin does not read as a second, conflicting one. + from pypaimon.common.options.core_options import CoreOptions + present = [k for k in SCAN_KEYS if table.options.options.contains_key(k)] + if present: + overrides = {k: None for k in present} + overrides[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(base.id) + table = table.copy(overrides) try: result = distributed_read_by_row_id( rid_ds, table, projection, @@ -144,4 +191,5 @@ def _project_rid(batch: pa.Table) -> pa.Table: raise # _reraise_inner always raises if result is None: return _empty_result(table, read_cols) - return result + # Lazy result; union a typed-empty block so an empty source still carries the schema. + return result.union(_empty_result(table, read_cols)) diff --git a/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py index 18e67de3e9c6..2fa6cd00c995 100644 --- a/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py @@ -164,6 +164,113 @@ def test_reads_blob_column(self): self.assertEqual(bytes(got[2]["payload"]), payloads[1]) self.assertEqual(bytes(got[4]["payload"]), payloads[3]) + def test_blob_as_descriptor_via_dynamic_options(self): + from pypaimon.ray import map_with_blobs + from pypaimon.table.row.blob import BlobDescriptor + blob_schema = pa.schema([("id", pa.int32()), ("payload", pa.large_binary())]) + target = self._create(schema=blob_schema) + payloads = [bytes([i]) * (i + 3) for i in range(1, 5)] + self._write(target, pa.Table.from_pydict( + {"id": [1, 2, 3, 4], "payload": pa.array(payloads, pa.large_binary())}, + schema=blob_schema)) + rid = self._rowid_by_id(target) + src = pa.table({"_ROW_ID": [rid[2], rid[4]]}, + schema=pa.schema([("_ROW_ID", pa.int64())])) + ds = read_by_row_id(target, ray.data.from_arrow(src), self.catalog_options, + projection=["payload"], + dynamic_options={"blob-as-descriptor": "true"}) + rows = ds.take_all() + self.assertTrue(all(BlobDescriptor.is_blob_descriptor(bytes(r["payload"])) for r in rows)) + + tbl = self.catalog.get_table(target) + + def fn(scalar_batch, blobs): + return pa.table({"_ROW_ID": scalar_batch.column("_ROW_ID").to_pylist(), + "n": [len(b) if b is not None else 0 for b in blobs["payload"]]}) + + res = map_with_blobs(ds, ["payload"], fn, file_io=tbl.file_io, + all_blob_columns=["payload"], batch_size=1) + n = {r["_ROW_ID"]: r["n"] for r in res.take_all()} + self.assertEqual(n[rid[2]], len(payloads[1])) + self.assertEqual(n[rid[4]], len(payloads[3])) + + def test_rejects_invariant_dynamic_options(self): + target = self._create() + self._write(target, pa.Table.from_pydict( + {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) + src = pa.table({"_ROW_ID": [0]}, schema=pa.schema([("_ROW_ID", pa.int64())])) + with self.assertRaisesRegex(ValueError, "invariant"): + read_by_row_id(target, src, self.catalog_options, projection=["age"], + dynamic_options={"deletion-vectors.enabled": "true"}) + + def test_time_travel_via_dynamic_options(self): + target = self._create() + self._write(target, pa.Table.from_pydict( + {"id": [1, 2], "name": ["a", "b"], "age": [1, 2]}, schema=self.pa_schema)) + self._write(target, pa.Table.from_pydict( # snapshot 2 adds ids 3, 4 + {"id": [3, 4], "name": ["c", "d"], "age": [3, 4]}, schema=self.pa_schema)) + rid = self._rowid_by_id(target) + idcol = pa.schema([("_ROW_ID", pa.int64())]) + + ds = read_by_row_id(target, pa.table({"_ROW_ID": [rid[1]]}, schema=idcol), + self.catalog_options, projection=["id", "age"], + dynamic_options={"scan.snapshot-id": "1"}) + got = self._rows_by_id(ds) + self.assertEqual(set(got), {1}) + + # id=3 exists only in snapshot 2; at snapshot 1 its row id is foreign + ds2 = read_by_row_id(target, pa.table({"_ROW_ID": [rid[3]]}, schema=idcol), + self.catalog_options, projection=["id", "age"], + dynamic_options={"scan.snapshot-id": "1"}) + with self.assertRaisesRegex(Exception, "valid range"): + ds2.take_all() + + def test_time_travel_via_tag_dynamic_options(self): + target = self._create() + self._write(target, pa.Table.from_pydict( + {"id": [1, 2], "name": ["a", "b"], "age": [1, 2]}, schema=self.pa_schema)) + self.catalog.get_table(target).create_tag("v1", 1) + self._write(target, pa.Table.from_pydict( + {"id": [3, 4], "name": ["c", "d"], "age": [3, 4]}, schema=self.pa_schema)) + rid = self._rowid_by_id(target) + idcol = pa.schema([("_ROW_ID", pa.int64())]) + + # tag v1 == snapshot 1: id=1 reads, id=3 (snapshot 2 only) is foreign + ds = read_by_row_id(target, pa.table({"_ROW_ID": [rid[1]]}, schema=idcol), + self.catalog_options, projection=["id", "age"], + dynamic_options={"scan.tag-name": "v1"}) + self.assertEqual(set(self._rows_by_id(ds)), {1}) + ds2 = read_by_row_id(target, pa.table({"_ROW_ID": [rid[3]]}, schema=idcol), + self.catalog_options, projection=["id", "age"], + dynamic_options={"scan.tag-name": "v1"}) + with self.assertRaisesRegex(Exception, "valid range"): + ds2.take_all() + + def test_rejects_multiple_time_travel_keys(self): + target = self._create() + self._write(target, pa.Table.from_pydict( + {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) + src = pa.table({"_ROW_ID": [0]}, schema=pa.schema([("_ROW_ID", pa.int64())])) + with self.assertRaisesRegex(ValueError, "at most one time-travel"): + read_by_row_id(target, src, self.catalog_options, projection=["age"], + dynamic_options={"scan.snapshot-id": "1", "scan.tag-name": "x"}) + + def test_time_travel_before_row_tracking_raises(self): + from pypaimon.schema.schema_change import SchemaChange + name = self._create(options={}) # plain table: no data-evolution / row-tracking + self._write(name, pa.Table.from_pydict( + {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) + self.catalog.alter_table(name, [ + SchemaChange.set_option("row-tracking.enabled", "true"), + SchemaChange.set_option("data-evolution.enabled", "true")]) + self._write(name, pa.Table.from_pydict( + {"id": [2], "name": ["b"], "age": [2]}, schema=self.pa_schema)) + src = pa.table({"_ROW_ID": [0]}, schema=pa.schema([("_ROW_ID", pa.int64())])) + # snapshot 1 predates row-tracking -> clear error, not a silent empty read + with self.assertRaisesRegex(ValueError, "row-tracking|data-evolution"): + read_by_row_id(name, src, self.catalog_options, projection=["age"], + dynamic_options={"scan.snapshot-id": "1"}) + def test_pins_base_snapshot(self): import importlib m = importlib.import_module("pypaimon.ray.read_by_row_id") @@ -277,6 +384,26 @@ def test_foreign_row_id_raises(self): with self.assertRaises(Exception): ds.take_all() + def test_returns_lazy_without_executing_source(self): + target = self._create() + self._write(target, pa.Table.from_pydict( + {"id": [1, 2], "name": ["a", "b"], "age": [1, 2]}, schema=self.pa_schema)) + rid = self._rowid_by_id(target) + marker = os.path.join(self.tempdir, f"exec_{uuid.uuid4().hex}") + + def spy(batch): + open(marker, "a").close() + return batch + + src = ray.data.from_arrow( + pa.table({"_ROW_ID": [rid[1]]}, schema=pa.schema([("_ROW_ID", pa.int64())])) + ).map_batches(spy, batch_format="pyarrow") + ds = read_by_row_id(target, src, self.catalog_options, projection=["id", "age"]) + self.assertFalse(os.path.exists(marker), "source was executed at call time") + rows = ds.take_all() + self.assertTrue(os.path.exists(marker)) + self.assertEqual({r["id"] for r in rows}, {1}) + def test_empty_source_non_empty_target_keeps_schema(self): target = self._create() self._write(target, pa.Table.from_pydict(