This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new a9a3d9277a [python][ray] Add distributed read_by_row_id for
data-evolution tables (#8465)
a9a3d9277a is described below
commit a9a3d9277a2574a3faa6d8cfb5d9d2d189f06d9d
Author: XiaoHongbo <[email protected]>
AuthorDate: Sun Jul 5 16:08:05 2026 +0800
[python][ray] Add distributed read_by_row_id for data-evolution tables
(#8465)
Add Ray support for distributed `read_by_row_id` on data-evolution
tables.
This is the read-side counterpart of `update_by_row_id`. It allows a Ray
Dataset carrying target `_ROW_ID`s, for example produced by
`bucket_join`, to read projected columns from the target table without
scanning or joining the whole target table.
This is useful for workflows such as:
1. match input keys to target `_ROW_ID`s with `bucket_join`
2. read only matched rows / blob columns by row id
3. transform or run inference in Ray
4. write results back with `update_by_row_id`
---
docs/docs/pypaimon/ray-data.md | 48 +++
paimon-python/pypaimon/ray/__init__.py | 2 +
.../pypaimon/ray/data_evolution_merge_join.py | 134 +++++++++
paimon-python/pypaimon/ray/read_by_row_id.py | 147 ++++++++++
.../pypaimon/tests/ray_read_by_row_id_test.py | 325 +++++++++++++++++++++
5 files changed, 656 insertions(+)
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index fc1e6cc9a9..c45f88ec09 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -542,3 +542,51 @@ print(metrics) # {"num_updated": 50}
- Partition columns cannot be updated (in-place rewrite can't move a row
across partitions).
- Deletion-vectors-enabled tables are not supported yet: a DV-deleted row
still lives
in its data file, so it can't be told apart from a live row without reading
the target.
+
+## Read By Row Id
+
+`read_by_row_id` is the read-side mirror of `update_by_row_id`: it reads
columns
+(including blob) of a **data-evolution** table for a set of `_ROW_ID`s, without
+scanning or joining the whole target. Each row id is routed to the data file
that
+owns it and only those files — and only the matched rows — are read. It pairs
with
+`bucket_join` (which produces the row ids) and feeds `update_by_row_id`: match
by
+key → read the matched rows → transform → write back by row id. Requires
+`ray >= 2.50` and a target with `data-evolution.enabled` and
`row-tracking.enabled`.
+
+```python
+from pypaimon.ray import read_by_row_id
+
+ds = read_by_row_id(
+ target="database_name.table_name",
+ row_ids=locator_ds, # ray.data.Dataset / pa.Table / pandas,
carrying the row ids
+ catalog_options={"warehouse": "/path/to/warehouse"},
+ projection=["image", "feature"], # columns to read; may include blob
columns
+ row_id_col="row_id", # source column holding the row ids (default
"_ROW_ID")
+)
+# ds: ray.data.Dataset of (image, feature, _ROW_ID) for the matched rows
+```
+
+**Parameters:**
+- `row_ids`: a `ray.data.Dataset`, `pyarrow.Table`, or `pandas.DataFrame`
carrying the
+ 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.
+- `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.
+- `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.
+
+**Returns:** a `ray.data.Dataset` of `(*projection, _ROW_ID)`.
+
+**Notes:**
+- 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.
+- 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.
diff --git a/paimon-python/pypaimon/ray/__init__.py
b/paimon-python/pypaimon/ray/__init__.py
index 52b4575b5b..bc91c8da45 100644
--- a/paimon-python/pypaimon/ray/__init__.py
+++ b/paimon-python/pypaimon/ray/__init__.py
@@ -28,6 +28,7 @@ from pypaimon.ray.data_evolution_merge_transform import (
lit,
)
from pypaimon.ray.update_by_row_id import update_by_row_id
+from pypaimon.ray.read_by_row_id import read_by_row_id
__all__ = [
"read_paimon",
@@ -36,6 +37,7 @@ __all__ = [
"bucket_join",
"merge_into",
"update_by_row_id",
+ "read_by_row_id",
"WhenMatched",
"WhenNotMatched",
"source_col",
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
index 03fc068ea2..5535b7ec1b 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
@@ -584,6 +584,140 @@ def distributed_update_apply(
return all_msgs, num_updated, action_row_ids
+def _read_output_schema(table, read_cols: Sequence[str]) -> "pa.Schema":
+ """Result schema: each projected column's type plus int64 ``_ROW_ID``, in
+ ``read_cols`` order. Shared by the empty-result paths so they can't
drift."""
+ from pypaimon.schema.data_types import PyarrowFieldParser
+ from pypaimon.table.special_fields import SpecialFields
+
+ rid = SpecialFields.ROW_ID.name
+ full = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
+ # Keep each field's nullability so an empty result matches a non-empty
read.
+ return pa.schema([
+ pa.field(rid, pa.int64(), nullable=False) if col == rid else
full.field(col)
+ for col in read_cols
+ ])
+
+
+def distributed_read_by_row_id(
+ row_ids_ds,
+ table,
+ projection: Sequence[str],
+ *,
+ num_partitions: int,
+ ray_remote_args: Optional[Dict[str, Any]] = None,
+ base_snapshot_id: Optional[int] = None,
+):
+ """Read ``projection`` for the ``_ROW_ID``s in ``row_ids_ds``, routing
each to its
+ owning file and reading only the matched rows via ``IndexedSplit`` slicing
(blob
+ resolved). Returns a ``ray.data.Dataset`` of ``(*projection, _ROW_ID)``,
or ``None``
+ if the target is empty. Read-side mirror of ``distributed_update_apply``.
+ """
+ import numpy as np
+ import uuid
+
+ import ray
+
+ from pypaimon.common.options.core_options import CoreOptions
+ from pypaimon.globalindex.indexed_split import IndexedSplit
+ from pypaimon.read.split import DataSplit
+ from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
+ from pypaimon.table.special_fields import SpecialFields
+ from pypaimon.utils.range import Range
+ from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+
+ row_id_name = SpecialFields.ROW_ID.name
+ read_cols = list(projection)
+ if row_id_name not in read_cols:
+ read_cols.append(row_id_name)
+
+ # Typed empty block so all output blocks share one schema.
+ empty_out = _read_output_schema(table, read_cols).empty_table()
+
+ # Read-only planner (only scans the manifest); pinned to the base snapshot
for stable routing.
+ scan_table = (
+ table.copy({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)})
+ if base_snapshot_id is not None else table
+ )
+ planner = TableUpdateByRowId(
+ scan_table,
+ "_read_by_row_id_planner_" + uuid.uuid4().hex[:8],
+ BATCH_COMMIT_IDENTIFIER,
+ )
+ sorted_first_row_ids = list(planner.first_row_ids)
+ if not sorted_first_row_ids:
+ return None
+
+ precomputed_info_ref = ray.put(planner._snapshot_files_info())
+ frid_col = "_FIRST_ROW_ID"
+ sorted_arr = np.asarray(sorted_first_row_ids, dtype=np.int64)
+ valid_ranges = planner.valid_row_id_ranges
+ range_starts = np.asarray([r.from_ for r in valid_ranges], dtype=np.int64)
+ range_ends = np.asarray([r.to for r in valid_ranges], dtype=np.int64)
+
+ def _assign_frid(batch: pa.Table) -> pa.Table:
+ if batch.num_rows == 0:
+ return batch.append_column(frid_col, pa.array([], type=pa.int64()))
+ rid_col = batch.column(row_id_name)
+ if rid_col.null_count:
+ raise ValueError(
+ "_ROW_ID is null; the planner snapshot is stale or the row ids
"
+ "come from a different table."
+ )
+ rids = rid_col.to_numpy(zero_copy_only=False)
+ # Foreign-id check: valid_ranges are sorted+merged, so one
searchsorted finds
+ # the candidate range (O(rows log ranges), like
distributed_delete_apply).
+ ridx = np.searchsorted(range_starts, rids, side="right") - 1
+ safe = np.clip(ridx, 0, len(range_starts) - 1)
+ in_range = (
+ (ridx >= 0)
+ & (rids >= range_starts[safe])
+ & (rids <= range_ends[safe])
+ )
+ if not in_range.all():
+ bad = rids[~in_range][0]
+ raise ValueError(
+ f"_ROW_ID {bad} does not belong to any valid range "
+ f"{[f'[{r.from_}, {r.to}]' for r in valid_ranges]}; the
planner "
+ f"snapshot is stale or the row ids come from a different
table."
+ )
+ idx = np.searchsorted(sorted_arr, rids, side="right") - 1
+ return batch.append_column(
+ frid_col, pa.array(sorted_arr[idx], type=pa.int64())
+ )
+
+ captured_table = scan_table # read at the same pinned snapshot the
planner routed on
+ captured_read_cols = read_cols
+ captured_empty = empty_out
+
+ def _read_group(group: pa.Table) -> pa.Table:
+ if group.num_rows == 0:
+ return captured_empty
+ frid = int(group.column(frid_col)[0].as_py())
+ info = ray.get(precomputed_info_ref)
+ owning_split, target_files = info.first_row_id_index[frid]
+ origin_split = DataSplit(
+ files=target_files,
+ partition=owning_split.partition,
+ bucket=owning_split.bucket,
+ raw_convertible=True,
+ )
+ # Only matched rows (deduped, contiguous ids -> ranges); blob gets
row-index pushdown.
+ wanted = set(group.column(row_id_name).to_pylist())
+ indexed = IndexedSplit(origin_split, Range.to_ranges(list(wanted)))
+ read = captured_table.new_read_builder().with_projection(
+ captured_read_cols
+ ).new_read()
+ return read.to_arrow([indexed])
+
+ map_kwargs = _map_kwargs(ray_remote_args)
+ with_frid = row_ids_ds.map_batches(_assign_frid, **map_kwargs)
+ group_partitions = max(1, min(len(sorted_first_row_ids), num_partitions))
+ return with_frid.groupby(frid_col,
num_partitions=group_partitions).map_groups(
+ _read_group, **map_kwargs
+ )
+
+
def distributed_delete_apply(
delete_ds,
table,
diff --git a/paimon-python/pypaimon/ray/read_by_row_id.py
b/paimon-python/pypaimon/ray/read_by_row_id.py
new file mode 100644
index 0000000000..b145823ba0
--- /dev/null
+++ b/paimon-python/pypaimon/ray/read_by_row_id.py
@@ -0,0 +1,147 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Distributed row-id read on Ray for data-evolution tables.
+
+The read-side mirror of ``update_by_row_id``: read columns (including blob)
for a set
+of ``_ROW_ID``s by routing each to its owning data file -- no full-target
read, no
+shuffle join. Pairs with ``bucket_join``, which produces the row ids.
+"""
+
+from typing import Any, Dict, List, Optional
+
+import pyarrow as pa
+
+from pypaimon.ray.data_evolution_merge_into import (
+ _normalize_source,
+ _reraise_inner,
+ _require_ray_join,
+ _resolve_num_partitions,
+)
+from pypaimon.ray.data_evolution_merge_join import (
+ _read_output_schema,
+ distributed_read_by_row_id,
+)
+
+__all__ = ["read_by_row_id"]
+
+
+def _empty_result(table: "FileStoreTable", read_cols: List[str]):
+ """An empty ``ray.data.Dataset`` with the projected read schema (empty
source
+ or target). Uses the same schema builder as the read path so they can't
drift."""
+ import ray
+
+ return ray.data.from_arrow(_read_output_schema(table,
read_cols).empty_table())
+
+
+def read_by_row_id(
+ target: str,
+ row_ids: Any,
+ catalog_options: Dict[str, str],
+ *,
+ projection: List[str],
+ row_id_col: Optional[str] = None,
+ num_partitions: Optional[int] = None,
+ ray_remote_args: Optional[Dict[str, Any]] = None,
+):
+ """Read ``projection`` columns of a data-evolution table by ``_ROW_ID``.
+
+ ``row_ids`` (a ``ray.data.Dataset`` / ``pyarrow.Table`` /
``pandas.DataFrame``)
+ must carry the target row ids in column ``row_id_col`` (default
``_ROW_ID``; set
+ 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.
+ Requires ``ray >= 2.50`` and a target with ``data-evolution.enabled`` +
+ ``row-tracking.enabled``.
+
+ Lookup/set semantics, like SQL ``... WHERE _ROW_ID IN (...)``: the result
has one
+ row per *distinct* matched row id -- duplicate row ids are deduplicated,
source
+ columns other than ``row_id_col`` are dropped, and the input row order is
not
+ preserved (rows come out grouped by owning file). An empty source yields
an empty
+ but correctly-typed Dataset.
+
+ Returns a ``ray.data.Dataset`` of ``(*projection, _ROW_ID)``.
+ """
+ from pypaimon.catalog.catalog_factory import CatalogFactory
+ from pypaimon.table.special_fields import SpecialFields
+
+ _require_ray_join()
+ if not projection:
+ raise ValueError("projection must be non-empty.")
+ projection = list(dict.fromkeys(projection))
+ num_partitions = _resolve_num_partitions(num_partitions)
+
+ table = CatalogFactory.create(catalog_options).get_table(target)
+ if not table.options.data_evolution_enabled():
+ raise ValueError(
+ f"read_by_row_id requires 'data-evolution.enabled'='true' on
'{target}'.")
+ if not table.options.row_tracking_enabled():
+ raise ValueError(
+ f"read_by_row_id requires 'row-tracking.enabled'='true' on
'{target}'.")
+ if table.options.deletion_vectors_enabled():
+ # A DV-deleted row still lives in its file, so slicing would surface
it.
+ raise ValueError(
+ f"read_by_row_id does not support deletion-vectors-enabled tables
yet: "
+ f"'{target}'.")
+
+ rid = SpecialFields.ROW_ID.name
+ src_rid_col = row_id_col or rid
+ for col in projection:
+ if col != rid and col not in table.field_names:
+ raise ValueError(f"projection column {col!r} is not in target
'{target}'.")
+
+ if isinstance(row_ids, str):
+ # A source table's _ROW_ID is its own, not the target's; require
in-memory ids.
+ raise ValueError(
+ "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):
+ raise ValueError(f"row_ids source is missing the {src_rid_col!r}
column.")
+
+ def _project_rid(batch: pa.Table) -> pa.Table:
+ 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()
+ # 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:
+ 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)
+ try:
+ result = distributed_read_by_row_id(
+ rid_ds, table, projection,
+ num_partitions=num_partitions,
+ ray_remote_args=ray_remote_args,
+ base_snapshot_id=base.id,
+ )
+ except Exception as e:
+ _reraise_inner(e)
+ raise # _reraise_inner always raises
+ if result is None:
+ return _empty_result(table, read_cols)
+ return result
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
new file mode 100644
index 0000000000..18e67de3e9
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_read_by_row_id_test.py
@@ -0,0 +1,325 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+import shutil
+import tempfile
+import unittest
+import uuid
+from unittest import mock
+
+import pyarrow as pa
+import pytest
+
+pypaimon = pytest.importorskip("pypaimon")
+ray = pytest.importorskip("ray")
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.ray import read_by_row_id
+
+
+class RayReadByRowIdTest(unittest.TestCase):
+ """Distributed row-id read: read only the files owning the given row ids
(and
+ only the matched rows), without reading or joining the whole target. The
+ read-side mirror of update_by_row_id."""
+
+ pa_schema = pa.schema([
+ ("id", pa.int32()),
+ ("name", pa.string()),
+ ("age", pa.int32()),
+ ])
+ de_options = {"row-tracking.enabled": "true", "data-evolution.enabled":
"true"}
+
+ @classmethod
+ def setUpClass(cls):
+ cls.tempdir = tempfile.mkdtemp()
+ cls.catalog_options = {"warehouse": os.path.join(cls.tempdir, "wh")}
+ cls.catalog = CatalogFactory.create(cls.catalog_options)
+ cls.catalog.create_database("default", True)
+ if not ray.is_initialized():
+ ray.init(ignore_reinit_error=True, num_cpus=2)
+
+ @classmethod
+ def tearDownClass(cls):
+ try:
+ if ray.is_initialized():
+ ray.shutdown()
+ except Exception:
+ pass
+ shutil.rmtree(cls.tempdir, ignore_errors=True)
+
+ def _create(self, options=None, schema=None):
+ name = f"default.r_{uuid.uuid4().hex[:8]}"
+ opts = self.de_options if options is None else options
+ self.catalog.create_table(
+ name, Schema.from_pyarrow_schema(schema or self.pa_schema,
options=opts), False)
+ return name
+
+ def _write(self, target, data):
+ t = self.catalog.get_table(target)
+ wb = t.new_batch_write_builder()
+ w = wb.new_write()
+ w.write_arrow(data)
+ wb.new_commit().commit(w.prepare_commit())
+ w.close()
+
+ def _read(self, target, projection=None):
+ t = self.catalog.get_table(target)
+ rb = t.new_read_builder()
+ if projection is not None:
+ rb = rb.with_projection(projection)
+ return rb.new_read().to_arrow(rb.new_scan().plan().splits())
+
+ def _rowid_by_id(self, target):
+ tab = self._read(target, ["_ROW_ID", "id"])
+ return dict(zip(tab.column("id").to_pylist(),
tab.column("_ROW_ID").to_pylist()))
+
+ def _rows_by_id(self, ds):
+ return {r["id"]: r for r in ds.take_all()}
+
+ def test_read_by_row_id_basic(self):
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": list(range(1, 7)), "name": [f"n{i}" for i in range(1, 7)],
+ "age": [i * 10 for i in range(1, 7)]}, schema=self.pa_schema))
+ rid = self._rowid_by_id(target)
+
+ want = [2, 5]
+ src = pa.table({"_ROW_ID": [rid[i] for i in want]},
+ schema=pa.schema([("_ROW_ID", pa.int64())]))
+ ds = read_by_row_id(target, ray.data.from_arrow(src),
self.catalog_options,
+ projection=["id", "name", "age"])
+ got = self._rows_by_id(ds)
+ self.assertEqual(set(got), set(want))
+ self.assertEqual(got[2]["age"], 20)
+ self.assertEqual(got[5]["name"], "n5")
+ self.assertEqual(got[2]["_ROW_ID"], rid[2])
+
+ def test_reads_correct_row_across_files(self):
+ target = self._create()
+ for chunk in ([10, 11, 12], [20, 21], [30, 31, 32, 33]):
+ self._write(target, pa.Table.from_pydict(
+ {"id": chunk, "name": ["x"] * len(chunk), "age": [c for c in
chunk]},
+ schema=self.pa_schema))
+ rid = self._rowid_by_id(target)
+ src = pa.table({"_ROW_ID": [rid[21]]}, schema=pa.schema([("_ROW_ID",
pa.int64())]))
+ ds = read_by_row_id(target, ray.data.from_arrow(src),
self.catalog_options,
+ projection=["id", "age"])
+ rows = ds.take_all()
+ self.assertEqual(len(rows), 1)
+ self.assertEqual(rows[0]["id"], 21)
+ self.assertEqual(rows[0]["age"], 21)
+
+ def test_reads_across_evolution_split_files(self):
+ # update_by_row_id writes a column delta (splits the row's range); the
read must merge original + delta.
+ from pypaimon.ray import update_by_row_id
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": [1, 2, 3], "name": ["a", "b", "c"], "age": [10, 20, 30]},
+ schema=self.pa_schema))
+ rid = self._rowid_by_id(target)
+ update_by_row_id(
+ target,
+ pa.table({"_ROW_ID": [rid[2]], "age": [999]},
+ schema=pa.schema([("_ROW_ID", pa.int64()), ("age",
pa.int32())])),
+ self.catalog_options, update_cols=["age"])
+ ds = read_by_row_id(
+ target,
+ pa.table({"_ROW_ID": [rid[2]]}, schema=pa.schema([("_ROW_ID",
pa.int64())])),
+ self.catalog_options, projection=["id", "name", "age"])
+ rows = ds.take_all()
+ self.assertEqual(len(rows), 1)
+ self.assertEqual(rows[0]["id"], 2)
+ self.assertEqual(rows[0]["name"], "b")
+ self.assertEqual(rows[0]["age"], 999)
+
+ def test_reads_blob_column(self):
+ blob_schema = pa.schema([("id", pa.int32()), ("payload",
pa.large_binary())])
+ target = self._create(schema=blob_schema)
+ payloads = [bytes([i]) * (i + 1) 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=["id", "payload"])
+ got = self._rows_by_id(ds)
+ self.assertEqual(set(got), {2, 4})
+ self.assertEqual(bytes(got[2]["payload"]), payloads[1])
+ self.assertEqual(bytes(got[4]["payload"]), payloads[3])
+
+ def test_pins_base_snapshot(self):
+ import importlib
+ m = importlib.import_module("pypaimon.ray.read_by_row_id")
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": [1, 2], "name": ["a", "b"], "age": [1, 2]},
schema=self.pa_schema))
+ expected_sid = self.catalog.get_table(
+ target).snapshot_manager().get_latest_snapshot().id
+ rid = self._rowid_by_id(target)
+ src = pa.table({"_ROW_ID": [rid[1]]}, schema=pa.schema([("_ROW_ID",
pa.int64())]))
+
+ captured = {}
+
+ def fake_read(rid_ds, table, projection, *, num_partitions,
+ ray_remote_args=None, base_snapshot_id=None):
+ captured["base_snapshot_id"] = base_snapshot_id
+ return ray.data.from_arrow(pa.table({"_ROW_ID": pa.array([],
pa.int64())}))
+
+ with mock.patch.object(m, "distributed_read_by_row_id", fake_read):
+ read_by_row_id(target, src, self.catalog_options,
projection=["age"])
+ self.assertEqual(captured["base_snapshot_id"], expected_sid)
+
+ def test_accepts_pyarrow_and_pandas_source(self):
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": [1, 2, 3], "name": ["a", "b", "c"], "age": [1, 2, 3]},
+ schema=self.pa_schema))
+ rid = self._rowid_by_id(target)
+ ds = read_by_row_id(
+ target, pa.table({"_ROW_ID": [rid[2]]},
schema=pa.schema([("_ROW_ID", pa.int64())])),
+ self.catalog_options, projection=["name"])
+ self.assertEqual([r["name"] for r in ds.take_all()], ["b"])
+ import pandas as pd
+ ds = read_by_row_id(
+ target, pd.DataFrame({"_ROW_ID": pd.array([rid[3]],
dtype="int64")}),
+ self.catalog_options, projection=["name"])
+ self.assertEqual([r["name"] for r in ds.take_all()], ["c"])
+
+ def test_ignores_extra_source_columns_and_dedups(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)
+ src = pa.table({"_ROW_ID": [rid[2], rid[2]], "junk": ["x", "y"]},
+ schema=pa.schema([("_ROW_ID", pa.int64()), ("junk",
pa.string())]))
+ ds = read_by_row_id(target, ray.data.from_arrow(src),
self.catalog_options,
+ projection=["id", "name"])
+ rows = ds.take_all()
+ self.assertEqual(len(rows), 1)
+ self.assertEqual(rows[0]["name"], "b")
+
+ def test_rejects_table_name_source(self):
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema))
+ with self.assertRaises(ValueError):
+ read_by_row_id(target, "default.some_source", self.catalog_options,
+ projection=["age"])
+
+ def test_rejects_non_data_evolution_table(self):
+ target = self._create(options={})
+ 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.assertRaises(ValueError):
+ read_by_row_id(target, src, self.catalog_options,
projection=["age"])
+
+ def test_rejects_deletion_vectors_table(self):
+ opts = dict(self.de_options, **{"deletion-vectors.enabled": "true"})
+ target = self._create(options=opts)
+ 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.assertRaises(ValueError):
+ read_by_row_id(target, src, self.catalog_options,
projection=["age"])
+
+ def test_rejects_missing_row_id_column(self):
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema))
+ src = pa.table({"id": [1]}, schema=pa.schema([("id", pa.int32())]))
+ with self.assertRaises(ValueError):
+ read_by_row_id(target, src, self.catalog_options,
projection=["age"])
+
+ def test_rejects_unknown_and_empty_projection(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.assertRaises(ValueError):
+ read_by_row_id(target, src, self.catalog_options,
projection=["nope"])
+ with self.assertRaises(ValueError):
+ read_by_row_id(target, src, self.catalog_options, projection=[])
+
+ def test_empty_target(self):
+ src = pa.table({"_ROW_ID": [0]}, schema=pa.schema([("_ROW_ID",
pa.int64())]))
+ empty_src = pa.table({"_ROW_ID": pa.array([], pa.int64())})
+
+ target = self._create()
+ with self.assertRaises(ValueError):
+ read_by_row_id(target, src, self.catalog_options,
projection=["age"])
+ ds = read_by_row_id(target, empty_src, self.catalog_options,
projection=["age"])
+ self.assertEqual(ds.count(), 0)
+
+ def test_foreign_row_id_raises(self):
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": [1, 2], "name": ["a", "b"], "age": [1, 2]},
schema=self.pa_schema))
+ src = pa.table({"_ROW_ID": [10_000]}, schema=pa.schema([("_ROW_ID",
pa.int64())]))
+ ds = read_by_row_id(target, src, self.catalog_options,
projection=["age"])
+ with self.assertRaises(Exception):
+ ds.take_all()
+
+ def test_empty_source_non_empty_target_keeps_schema(self):
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": [1, 2], "name": ["a", "b"], "age": [1, 2]},
schema=self.pa_schema))
+ empty_src = pa.table({"_ROW_ID": pa.array([], pa.int64())})
+ ds = read_by_row_id(target, empty_src, self.catalog_options,
projection=["id", "age"])
+ self.assertEqual(ds.count(), 0)
+ self.assertIsNotNone(ds.schema())
+ self.assertEqual(set(ds.schema().names), {"id", "age", "_ROW_ID"})
+
+ def test_empty_result_schema_matches_nonempty_read(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)
+ proj = ["id", "age"]
+ nonempty = read_by_row_id(
+ target, pa.table({"_ROW_ID": [rid[1]]},
schema=pa.schema([("_ROW_ID", pa.int64())])),
+ self.catalog_options, projection=proj)
+ real_schema = None
+ for b in nonempty.iter_batches(batch_format="pyarrow"):
+ real_schema = b.schema
+ break
+ empty = read_by_row_id(
+ target, pa.table({"_ROW_ID": pa.array([], pa.int64())}),
+ self.catalog_options, projection=proj)
+ self.assertTrue(real_schema.equals(empty.schema().base_schema))
+ self.assertFalse(empty.schema().base_schema.field("_ROW_ID").nullable)
+
+ def test_custom_row_id_col(self):
+ target = self._create()
+ self._write(target, pa.Table.from_pydict(
+ {"id": [1, 2, 3], "name": ["a", "b", "c"], "age": [1, 2, 3]},
+ schema=self.pa_schema))
+ rid = self._rowid_by_id(target)
+ src = pa.table({"row_id": [rid[2]], "url": ["u2"]},
+ schema=pa.schema([("row_id", pa.int64()), ("url",
pa.string())]))
+ ds = read_by_row_id(target, src, self.catalog_options,
+ projection=["name"], row_id_col="row_id")
+ rows = ds.take_all()
+ self.assertEqual([r["name"] for r in rows], ["b"])
+ self.assertEqual(rows[0]["_ROW_ID"], rid[2])
+
+
+if __name__ == "__main__":
+ unittest.main()
