| # 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_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_rejects_disabling_query_auth_dynamically(self): | |
| target = self._create({**self.de_options, "query-auth.enabled": "true"}) | |
| 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())])) | |
| for value in ("false", False, None, "invalid"): | |
| with self.subTest(value=value): | |
| with self.assertRaisesRegex( | |
| ValueError, "^Option 'query-auth.enabled' is not allowed to be disabled dynamically$"): | |
| read_by_row_id(target, src, self.catalog_options, projection=["age"], | |
| dynamic_options={"query-auth.enabled": value}) | |
| 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_retained_tag_after_snapshot_expiry(self): | |
| from pypaimon.schema.schema_change import SchemaChange | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1, 2], "name": ["a", "b"], "age": [10, 20]}, schema=self.pa_schema)) | |
| table = self.catalog.get_table(target) | |
| table.create_tag("training", 1) | |
| rid = self._rowid_by_id(target) | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [3], "name": ["c"], "age": [30]}, schema=self.pa_schema)) | |
| self.catalog.alter_table(target, [SchemaChange.rename_column("age", "years")]) | |
| table.file_io.delete(table.snapshot_manager().get_snapshot_path(1)) | |
| for selector in ("scan.tag-name", "scan.version"): | |
| with self.subTest(selector=selector): | |
| ds = read_by_row_id( | |
| target, pa.table({"_ROW_ID": [rid[2], rid[1], rid[1]]}), | |
| self.catalog_options, projection=["id", "age"], | |
| dynamic_options={selector: "training", "scan.native-plan.enabled": "true"}) | |
| self.assertEqual(self._rows_by_id(ds), { | |
| 1: {"id": 1, "age": 10, "_ROW_ID": rid[1]}, | |
| 2: {"id": 2, "age": 20, "_ROW_ID": rid[2]}, | |
| }) | |
| def test_rejects_tag_schema_change_during_snapshot_resolution(self): | |
| import importlib | |
| from pypaimon.schema.schema_change import SchemaChange | |
| module = importlib.import_module("pypaimon.ray.read_by_row_id") | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [10]}, schema=self.pa_schema)) | |
| table = self.catalog.get_table(target) | |
| table.create_tag("training", 1) | |
| self.catalog.alter_table(target, [SchemaChange.rename_column("age", "years")]) | |
| schema = pa.schema([("id", pa.int32()), ("name", pa.string()), ("years", pa.int32())]) | |
| self._write(target, pa.table({"id": [2], "name": ["b"], "years": [20]}, schema=schema)) | |
| resolve = module._read_snapshot | |
| def resolve_after_tag_moves(read_table): | |
| table.replace_tag("training") | |
| return resolve(read_table) | |
| with mock.patch.object(module, "_read_snapshot", resolve_after_tag_moves): | |
| with self.assertRaisesRegex(ValueError, "schema changed.*retry"): | |
| read_by_row_id( | |
| target, pa.table({"_ROW_ID": [0]}), self.catalog_options, | |
| projection=["age"], dynamic_options={"scan.tag-name": "training"}) | |
| def test_lazy_tag_read_keeps_resolved_snapshot(self): | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [10]}, schema=self.pa_schema)) | |
| table = self.catalog.get_table(target) | |
| table.create_tag("training", 1) | |
| rid = self._rowid_by_id(target)[1] | |
| ds = read_by_row_id( | |
| target, pa.table({"_ROW_ID": [rid]}), self.catalog_options, | |
| projection=["id", "age"], dynamic_options={"scan.tag-name": "training"}) | |
| from pypaimon.multimodal.table import MultimodalTable | |
| MultimodalTable(self.catalog, target, table).update("id = 1", {"age": 99}) | |
| table.replace_tag("training") | |
| table.file_io.delete(table.snapshot_manager().get_snapshot_path(1)) | |
| self.assertEqual(ds.take_all(), [{"id": 1, "age": 10, "_ROW_ID": rid}]) | |
| latest = read_by_row_id( | |
| target, pa.table({"_ROW_ID": [rid]}), self.catalog_options, | |
| projection=["age"], dynamic_options={"scan.tag-name": "training"}) | |
| self.assertEqual(latest.take_all(), [{"age": 99, "_ROW_ID": rid}]) | |
| def test_row_tracking_cannot_be_enabled_after_data(self): | |
| # row-tracking.enabled / data-evolution.enabled are immutable once the | |
| # table has snapshots (Java parity), so a "snapshot predates | |
| # row-tracking" state can no longer be constructed through ALTER; the | |
| # alter itself is rejected instead. | |
| 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)) | |
| with self.assertRaisesRegex(RuntimeError, "not supported yet"): | |
| self.catalog.alter_table(name, [ | |
| SchemaChange.set_option("row-tracking.enabled", "true"), | |
| SchemaChange.set_option("data-evolution.enabled", "true")]) | |
| 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, | |
| estimated_size_bytes=None, estimated_num_rows=None, | |
| data_context=None): | |
| captured["base_snapshot_id"] = table._read_snapshot.id | |
| captured["num_partitions"] = num_partitions | |
| captured["estimated_size_bytes"] = estimated_size_bytes | |
| captured["estimated_num_rows"] = estimated_num_rows | |
| captured["data_context"] = data_context | |
| 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) | |
| self.assertIsNone(captured["num_partitions"]) | |
| self.assertGreater(captured["estimated_size_bytes"], 0) | |
| self.assertEqual(captured["estimated_num_rows"], 1) | |
| self.assertIsNotNone(captured["data_context"]) | |
| 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_reads_deletion_vectors_and_pins_snapshot(self): | |
| opts = dict(self.de_options, **{"deletion-vectors.enabled": "true"}) | |
| target = self._create(options=opts) | |
| for ids in ([1, 2, 3], [4, 5, 6]): | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": ids, "name": ["a"] * 3, "age": ids}, schema=self.pa_schema)) | |
| table = self.catalog.get_table(target) | |
| src = pa.table({"_ROW_ID": [0, 1, 2, 3, 4, 5, 1]}) | |
| before = read_by_row_id(target, src, self.catalog_options, projection=["id"]) | |
| wb = table.new_batch_write_builder() | |
| commit = wb.new_commit() | |
| try: | |
| commit.commit(wb.new_update().delete_by_row_id([0, 2, 4])) | |
| after = read_by_row_id(target, src, self.catalog_options, projection=["id"]) | |
| self.assertEqual(set(self._rows_by_id(before)), {1, 2, 3, 4, 5, 6}) | |
| self.assertEqual(set(self._rows_by_id(after)), {2, 4, 6}) | |
| deleted = read_by_row_id( | |
| target, pa.table({"_ROW_ID": [0, 2, 4]}), self.catalog_options, projection=["id"]) | |
| self.assertEqual(deleted.take_all(), []) | |
| self.assertEqual(deleted.schema().names, ["id", "_ROW_ID"]) | |
| commit.close() | |
| commit = wb.new_commit() | |
| commit.commit(wb.new_update().delete_by_row_id([1, 3, 5])) | |
| self.assertEqual(read_by_row_id( | |
| target, src, self.catalog_options, projection=["id"]).take_all(), []) | |
| finally: | |
| commit.close() | |
| 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_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( | |
| {"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)) | |
| 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() |