| # 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 collections |
| import datetime |
| import os |
| import shutil |
| import tempfile |
| import unittest |
| |
| import pyarrow as pa |
| import pytest |
| |
| pypaimon = pytest.importorskip("pypaimon") |
| ray = pytest.importorskip("ray") |
| |
| import importlib |
| from unittest import mock |
| |
| from pypaimon import CatalogFactory, Schema |
| from pypaimon.ray import range_join |
| |
| rjmod = importlib.import_module("pypaimon.ray.range_join") |
| |
| |
| class RayRangeJoinTest(unittest.TestCase): |
| """Range-aligned join must equal a global inner join, cutting the key space from |
| per-file min/max stats so each range is read/joined in its own task (no shuffle).""" |
| |
| @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=4) |
| |
| @classmethod |
| def tearDownClass(cls): |
| try: |
| if ray.is_initialized(): |
| ray.shutdown() |
| except Exception: |
| pass |
| shutil.rmtree(cls.tempdir, ignore_errors=True) |
| |
| def _table(self, name, schema, commits, primary_keys=None, options=None): |
| """Create a table and write each arrow table in ``commits`` as its own commit, |
| so the manifest holds several data files with distinct key ranges.""" |
| self.catalog.create_table( |
| name, |
| Schema.from_pyarrow_schema(schema, primary_keys=primary_keys, options=options), |
| False) |
| t = self.catalog.get_table(name) |
| for data in commits: |
| wb = t.new_batch_write_builder() |
| w = wb.new_write() |
| w.write_arrow(data) |
| wb.new_commit().commit(w.prepare_commit()) |
| w.close() |
| return name |
| |
| def test_range_join_matches_global_join(self): |
| loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) |
| ins = pa.schema([("k", pa.int64())]) |
| # locator: k in 0..599 spread across three files with disjoint key ranges. |
| self._table("default.rj_loc", loc, [ |
| pa.Table.from_pydict({"k": list(range(0, 200)), |
| "row_id": list(range(0, 200))}, schema=loc), |
| pa.Table.from_pydict({"k": list(range(200, 400)), |
| "row_id": list(range(200, 400))}, schema=loc), |
| pa.Table.from_pydict({"k": list(range(400, 600)), |
| "row_id": list(range(400, 600))}, schema=loc), |
| ]) |
| self._table("default.rj_in", ins, [ |
| pa.Table.from_pydict({"k": list(range(0, 250))}, schema=ins), |
| ]) |
| ds = range_join( |
| "default.rj_in", "default.rj_loc", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k", "row_id"], num_ranges=4) |
| got = {r["k"]: r["row_id"] for r in ds.take_all()} |
| self.assertEqual(set(got), set(range(250))) |
| self.assertTrue(all(got[i] == i for i in range(250))) |
| |
| def test_fan_out_one_key_many_rows(self): |
| loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_fan_loc", loc, [ |
| pa.Table.from_pydict({"k": [5, 5, 7], "row_id": [0, 1, 2]}, schema=loc)]) |
| self._table("default.rj_fan_in", ins, [ |
| pa.Table.from_pydict({"k": [5]}, schema=ins)]) |
| ds = range_join( |
| "default.rj_fan_in", "default.rj_fan_loc", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k", "row_id"]) |
| self.assertEqual(sorted(r["row_id"] for r in ds.take_all()), [0, 1]) |
| |
| def test_left_on_right_on_different_names(self): |
| right = pa.schema([("rid", pa.int64()), ("val", pa.string())]) |
| left = pa.schema([("lid", pa.int64())]) |
| self._table("default.rj_lr_right", right, [ |
| pa.Table.from_pydict({"rid": list(range(100)), |
| "val": [f"v{i}" for i in range(100)]}, schema=right)]) |
| self._table("default.rj_lr_left", left, [ |
| pa.Table.from_pydict({"lid": list(range(30))}, schema=left)]) |
| ds = range_join( |
| "default.rj_lr_left", "default.rj_lr_right", self.catalog_options, |
| left_on="lid", right_on="rid", num_ranges=3) |
| # Output keeps the left key name (pyarrow coalesces the right key into it). |
| got = {r["lid"]: r["val"] for r in ds.take_all()} |
| self.assertEqual(got, {i: f"v{i}" for i in range(30)}) |
| |
| def test_num_ranges_one_is_correct(self): |
| # A single range degenerates to one local join and must still be exact. |
| loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_one_loc", loc, [ |
| pa.Table.from_pydict({"k": list(range(50)), |
| "row_id": list(range(50))}, schema=loc)]) |
| self._table("default.rj_one_in", ins, [ |
| pa.Table.from_pydict({"k": list(range(20))}, schema=ins)]) |
| ds = range_join( |
| "default.rj_one_in", "default.rj_one_loc", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k", "row_id"], num_ranges=1) |
| got = {r["k"]: r["row_id"] for r in ds.take_all()} |
| self.assertEqual(got, {i: i for i in range(20)}) |
| |
| def test_dispatches_multiple_range_tasks(self): |
| # No global shuffle: several disjoint-range files produce more than one task. |
| loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_disp_loc", loc, [ |
| pa.Table.from_pydict({"k": list(range(0, 300)), |
| "row_id": list(range(0, 300))}, schema=loc), |
| pa.Table.from_pydict({"k": list(range(300, 600)), |
| "row_id": list(range(300, 600))}, schema=loc), |
| ]) |
| self._table("default.rj_disp_in", ins, [ |
| pa.Table.from_pydict({"k": list(range(0, 600))}, schema=ins)]) |
| |
| captured = {} |
| real = ray.data.from_arrow_refs |
| |
| def spy(refs): |
| captured["n"] = len(refs) |
| return real(refs) |
| |
| with mock.patch.object(ray.data, "from_arrow_refs", spy): |
| ds = range_join( |
| "default.rj_disp_in", "default.rj_disp_loc", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k", "row_id"], num_ranges=4) |
| ds.take_all() |
| self.assertGreater(captured["n"], 1) |
| |
| def test_rejects_shared_non_key_column(self): |
| loc = pa.schema([("k", pa.int64()), ("v", pa.int64())]) |
| ins = pa.schema([("k", pa.int64()), ("v", pa.int64())]) |
| self._table("default.rj_col_loc", loc, [ |
| pa.Table.from_pydict({"k": [1], "v": [1]}, schema=loc)]) |
| self._table("default.rj_col_in", ins, [ |
| pa.Table.from_pydict({"k": [1], "v": [2]}, schema=ins)]) |
| with self.assertRaisesRegex(ValueError, "collide"): |
| range_join("default.rj_col_in", "default.rj_col_loc", self.catalog_options, on="k") |
| |
| def test_rejects_key_type_mismatch(self): |
| loc = pa.schema([("k", pa.int32()), ("row_id", pa.int64())]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_ty_loc", loc, [ |
| pa.Table.from_pydict({"k": pa.array([1], pa.int32()), "row_id": [1]}, schema=loc)]) |
| self._table("default.rj_ty_in", ins, [ |
| pa.Table.from_pydict({"k": [1]}, schema=ins)]) |
| with self.assertRaisesRegex(ValueError, "same type"): |
| range_join("default.rj_ty_in", "default.rj_ty_loc", self.catalog_options, on="k") |
| |
| def test_partition_filter_and_partitioned_table(self): |
| # range_join works on partitioned tables (unlike bucket_join); left_partitions |
| # prunes to the requested partition before planning. |
| loc = pa.schema([("p", pa.string()), ("k", pa.int64())]) |
| self.catalog.create_table( |
| "default.rj_pf_l", |
| Schema.from_pyarrow_schema(loc, partition_keys=["p"]), False) |
| t = self.catalog.get_table("default.rj_pf_l") |
| wb = t.new_batch_write_builder() |
| w = wb.new_write() |
| w.write_arrow(pa.Table.from_pydict( |
| {"p": ["a", "a", "b", "b"], "k": [1, 2, 3, 4]}, schema=loc)) |
| wb.new_commit().commit(w.prepare_commit()) |
| w.close() |
| self._table("default.rj_pf_r", pa.schema([("k2", pa.int64()), ("val", pa.string())]), [ |
| pa.Table.from_pydict({"k2": [1, 2, 3, 4], "val": ["v1", "v2", "v3", "v4"]}, |
| schema=pa.schema([("k2", pa.int64()), ("val", pa.string())]))]) |
| |
| ds = range_join("default.rj_pf_l", "default.rj_pf_r", self.catalog_options, |
| left_on="k", right_on="k2", left_projection=["k"], |
| right_projection=["k2", "val"], left_partitions={"p": "a"}, num_ranges=2) |
| self.assertEqual(sorted((r["k"], r["val"]) for r in ds.take_all()), |
| [(1, "v1"), (2, "v2")]) |
| # Whole partitioned table joins fine when unfiltered. |
| ds = range_join("default.rj_pf_l", "default.rj_pf_r", self.catalog_options, |
| left_on="k", right_on="k2", left_projection=["k"], |
| right_projection=["k2", "val"], num_ranges=2) |
| self.assertEqual(sorted(r["k"] for r in ds.take_all()), [1, 2, 3, 4]) |
| |
| def test_many_to_many_matches_global_join(self): |
| loc = pa.schema([("k", pa.int64()), ("rid", pa.int64())]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_mm_loc", loc, [ |
| pa.Table.from_pydict({"k": [1, 1, 2], "rid": [10, 11, 20]}, schema=loc)]) |
| self._table("default.rj_mm_in", ins, [ |
| pa.Table.from_pydict({"k": [1, 1, 2]}, schema=ins)]) |
| # key 1: 2 left x 2 right = 4 rows; key 2: 1 x 1 = 1 row. |
| for num_ranges in (1, 3): |
| ds = range_join("default.rj_mm_in", "default.rj_mm_loc", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k", "rid"], |
| num_ranges=num_ranges) |
| got = sorted(r["rid"] for r in ds.take_all()) |
| self.assertEqual(got, [10, 10, 11, 11, 20]) |
| |
| def test_rejects_bad_on_spec(self): |
| with self.assertRaisesRegex(ValueError, "exactly one of"): |
| range_join("a", "b", self.catalog_options) # neither on nor left_on/right_on |
| |
| def test_rejects_float_range_key(self): |
| schema = pa.schema([("k", pa.float64()), ("v", pa.int64())]) |
| self._table("default.rj_float_a", schema, [ |
| pa.Table.from_pydict({"k": [1.0], "v": [1]}, schema=schema)]) |
| self._table("default.rj_float_b", schema, [ |
| pa.Table.from_pydict({"k": [1.0], "v": [2]}, schema=schema)]) |
| with self.assertRaisesRegex(ValueError, "FLOAT/DOUBLE"): |
| range_join("default.rj_float_a", "default.rj_float_b", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k"]) |
| |
| def test_rejects_left_key_vs_right_column_collision(self): |
| left = pa.schema([("lid", pa.int64()), ("x", pa.int64())]) |
| right = pa.schema([("rid", pa.int64()), ("lid", pa.int64())]) |
| self._table("default.rj_xn_left", left, [ |
| pa.Table.from_pydict({"lid": [1], "x": [1]}, schema=left)]) |
| self._table("default.rj_xn_right", right, [ |
| pa.Table.from_pydict({"rid": [1], "lid": [9]}, schema=right)]) |
| # Left key 'lid' collides with the right non-key column 'lid' in the output. |
| with self.assertRaisesRegex(ValueError, "collide"): |
| range_join("default.rj_xn_left", "default.rj_xn_right", self.catalog_options, |
| left_on="lid", right_on="rid") |
| |
| def test_date_to_timestamp_schema_evolution(self): |
| # A DATE->TIMESTAMP evolved key yields date footers in old files and datetime in |
| # new ones; the planner must coerce both to the key type, not compare them raw. |
| from pypaimon.schema.data_types import AtomicType |
| from pypaimon.schema.schema_change import SchemaChange |
| |
| a_date = pa.schema([("k", pa.date32())]) |
| self.catalog.create_table( |
| "default.rj_ev_a", Schema.from_pyarrow_schema(a_date), False) |
| t = self.catalog.get_table("default.rj_ev_a") |
| wb = t.new_batch_write_builder() |
| w = wb.new_write() |
| w.write_arrow(pa.Table.from_pydict( |
| {"k": [datetime.date(2020, 1, 1), datetime.date(2020, 1, 2)]}, schema=a_date)) |
| wb.new_commit().commit(w.prepare_commit()) |
| w.close() |
| self.catalog.alter_table( |
| "default.rj_ev_a", |
| [SchemaChange.update_column_type("k", AtomicType("TIMESTAMP(6)"))], False) |
| t = self.catalog.get_table("default.rj_ev_a") |
| a_ts = pa.schema([("k", pa.timestamp("us"))]) |
| wb = t.new_batch_write_builder() |
| w = wb.new_write() |
| w.write_arrow(pa.Table.from_pydict( |
| {"k": [datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 2)]}, schema=a_ts)) |
| wb.new_commit().commit(w.prepare_commit()) |
| w.close() |
| |
| b = pa.schema([("bk", pa.timestamp("us")), ("val", pa.string())]) |
| self.catalog.create_table("default.rj_ev_b", Schema.from_pyarrow_schema(b), False) |
| t = self.catalog.get_table("default.rj_ev_b") |
| wb = t.new_batch_write_builder() |
| w = wb.new_write() |
| w.write_arrow(pa.Table.from_pydict( |
| {"bk": [datetime.datetime(2020, 1, 1), datetime.datetime(2020, 6, 1)], |
| "val": ["jan1", "jun1"]}, schema=b)) |
| wb.new_commit().commit(w.prepare_commit()) |
| w.close() |
| |
| ds = range_join("default.rj_ev_a", "default.rj_ev_b", self.catalog_options, |
| left_on="k", right_on="bk", num_ranges=3) |
| got = sorted((str(r["k"]), r["val"]) for r in ds.take_all()) |
| self.assertEqual(got, [("2020-01-01 00:00:00", "jan1"), |
| ("2020-06-01 00:00:00", "jun1")]) |
| |
| def test_int_to_string_schema_evolution_no_dropped_rows(self): |
| # INT->STRING isn't order-preserving ('10' < '2'), so an old INT file's footer |
| # bounds are invalid under the new string order. Such files must be treated as |
| # unknown (join every range), not pruned, or rows are silently dropped. |
| from pypaimon.schema.data_types import AtomicType |
| from pypaimon.schema.schema_change import SchemaChange |
| |
| a_int = pa.schema([("k", pa.int32())]) |
| self.catalog.create_table( |
| "default.rj_is_a", Schema.from_pyarrow_schema(a_int), False) |
| t = self.catalog.get_table("default.rj_is_a") |
| wb = t.new_batch_write_builder() |
| w = wb.new_write() |
| # int order 5<42<100, but as strings '100'<'42'<'5'. |
| w.write_arrow(pa.Table.from_pydict({"k": [5, 42, 100]}, schema=a_int)) |
| wb.new_commit().commit(w.prepare_commit()) |
| w.close() |
| self.catalog.alter_table( |
| "default.rj_is_a", |
| [SchemaChange.update_column_type("k", AtomicType("STRING"))], False) |
| |
| b = pa.schema([("bk", pa.string()), ("val", pa.string())]) |
| self.catalog.create_table("default.rj_is_b", Schema.from_pyarrow_schema(b), False) |
| t = self.catalog.get_table("default.rj_is_b") |
| wb = t.new_batch_write_builder() |
| w = wb.new_write() |
| w.write_arrow(pa.Table.from_pydict( |
| {"bk": ["5", "42", "100"], "val": ["v5", "v42", "v100"]}, schema=b)) |
| wb.new_commit().commit(w.prepare_commit()) |
| w.close() |
| |
| for num_ranges in (1, 3): |
| ds = range_join("default.rj_is_a", "default.rj_is_b", self.catalog_options, |
| left_on="k", right_on="bk", num_ranges=num_ranges) |
| got = sorted((r["k"], r["val"]) for r in ds.take_all()) |
| self.assertEqual(got, [("100", "v100"), ("42", "v42"), ("5", "v5")]) |
| |
| def test_reread_budget_bounds_wide_and_unknown_splits(self): |
| Split = collections.namedtuple("Split", "files") |
| File = collections.namedtuple("File", "row_count file_size") |
| |
| def rng(lo, hi, rows=100, size=100): |
| return (Split([File(rows, size)]), lo, hi) |
| |
| # _total_reads = bytes x ranges a split overlaps. |
| self.assertEqual(rjmod._total_reads([rng(0, 10)], [], [(None, 5), (5, None)]), 200) |
| # Budget is bytes, not rows: a wide split (few rows, large files) counts its bytes. |
| wide_row = [rng(0, 10, rows=1, size=1000)] |
| self.assertEqual(rjmod._total_reads(wide_row, [], [(None, 5), (5, None)]), 2000) |
| # All-unknown collapses to a single range. |
| unknown = [(Split([File(100, 100)]), None, None)] |
| self.assertEqual(len(rjmod._bounded_ranges(unknown, unknown, 8)), 1) |
| # Wide known splits (each overlaps many ranges) are bounded by the budget. |
| wide = [rng(0, 100), rng(0, 100), rng(0, 100), rng(0, 100)] |
| ranges = rjmod._bounded_ranges(wide, wide, 16) |
| budget = rjmod._REREAD_BUDGET * (8 * 100) |
| self.assertTrue(len(ranges) == 1 |
| or rjmod._total_reads(wide, wide, ranges) <= budget) |
| # Clustered (disjoint) splits keep at least as much parallelism as wide ones. |
| clustered = [rng(0, 9), rng(10, 19), rng(20, 29), rng(30, 39)] |
| self.assertGreaterEqual(len(rjmod._bounded_ranges(clustered, clustered, 4)), |
| len(rjmod._bounded_ranges(wide, wide, 4))) |
| |
| def test_split_key_range_reads_stats(self): |
| # With default metadata.stats-mode=none, the planner falls back to the footer. |
| loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) |
| self._table("default.rj_stats", loc, [ |
| pa.Table.from_pydict({"k": [10, 20, 15], "row_id": [1, 2, 3]}, schema=loc)]) |
| ranged, _ = rjmod._plan_ranged_splits( |
| "default.rj_stats", self.catalog_options, None, "k") |
| self.assertTrue(ranged) |
| los = [lo for _, lo, _ in ranged if lo is not None] |
| his = [hi for _, _, hi in ranged if hi is not None] |
| self.assertEqual(min(los), 10) |
| self.assertEqual(max(his), 20) |
| |
| def test_manifest_stats_avoid_footer_reads(self): |
| schema = pa.schema([("k", pa.int64()), ("v", pa.string())]) |
| self._table("default.rj_manifest", schema, [ |
| pa.Table.from_pydict({"k": [10, 20, 15], "v": ["a", "b", "c"]}, schema=schema) |
| ], options={"metadata.stats-mode": "full"}) |
| |
| with mock.patch("pyarrow.parquet.read_metadata", |
| side_effect=AssertionError("footer should not be read")): |
| ranged, _ = rjmod._plan_ranged_splits( |
| "default.rj_manifest", self.catalog_options, None, "k") |
| self.assertEqual([(lo, hi) for _, lo, hi in ranged], [(10, 20)]) |
| |
| def test_key_stats_avoid_footer_reads(self): |
| schema = pa.schema([("k", pa.int64()), ("v", pa.string())]) |
| self._table("default.rj_key_stats", schema, [ |
| pa.Table.from_pydict({"k": [10, 20, 15], "v": ["a", "b", "c"]}, schema=schema) |
| ], primary_keys=["k"], options={"bucket": "1"}) |
| |
| with mock.patch("pyarrow.parquet.read_metadata", |
| side_effect=AssertionError("footer should not be read")): |
| ranged, _ = rjmod._plan_ranged_splits( |
| "default.rj_key_stats", self.catalog_options, None, "k") |
| self.assertEqual([(lo, hi) for _, lo, hi in ranged], [(10, 20)]) |
| |
| def test_manifest_stats_follow_field_id_after_rename(self): |
| from pypaimon.schema.schema_change import SchemaChange |
| |
| schema = pa.schema([("k", pa.int64()), ("v", pa.string())]) |
| name = "default.rj_manifest_rename" |
| self._table(name, schema, [ |
| pa.Table.from_pydict({"k": [3, 9, 6], "v": ["a", "b", "c"]}, schema=schema) |
| ], options={"metadata.stats-mode": "full"}) |
| self.catalog.alter_table( |
| name, [SchemaChange.rename_column("k", "renamed")], False) |
| |
| with mock.patch("pyarrow.parquet.read_metadata", |
| side_effect=AssertionError("footer should not be read")): |
| ranged, _ = rjmod._plan_ranged_splits( |
| name, self.catalog_options, None, "renamed") |
| self.assertEqual([(lo, hi) for _, lo, hi in ranged], [(3, 9)]) |
| |
| def test_footer_failure_degrades_to_unknown(self): |
| schema = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_footer_failure", schema, [ |
| pa.Table.from_pydict({"k": [1, 2, 3]}, schema=schema)]) |
| |
| with self.assertLogs(rjmod._LOG, level="WARNING"): |
| with mock.patch("pyarrow.parquet.read_metadata", |
| side_effect=OSError("not seekable")): |
| ranged, _ = rjmod._plan_ranged_splits( |
| "default.rj_footer_failure", self.catalog_options, None, "k") |
| self.assertTrue(ranged) |
| self.assertTrue(all(lo is None and hi is None for _, lo, hi in ranged)) |
| |
| def test_stats_mode_none_still_correct(self): |
| # metadata.stats-mode=none only drops manifest stats; the parquet footer still |
| # carries min/max (range_join's actual source), so ranges still work. The |
| # unknown-split fallback itself is covered by the planning-logic tests. |
| no_stats = {"metadata.stats-mode": "none"} |
| loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_ns_loc", loc, [ |
| pa.Table.from_pydict({"k": list(range(0, 100)), |
| "row_id": list(range(0, 100))}, schema=loc), |
| pa.Table.from_pydict({"k": list(range(100, 200)), |
| "row_id": list(range(100, 200))}, schema=loc), |
| ], options=no_stats) |
| self._table("default.rj_ns_in", ins, [ |
| pa.Table.from_pydict({"k": list(range(50, 150))}, schema=ins)], |
| options=no_stats) |
| ds = range_join( |
| "default.rj_ns_in", "default.rj_ns_loc", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k", "row_id"], num_ranges=4) |
| got = sorted((r["k"], r["row_id"]) for r in ds.take_all()) |
| self.assertEqual(got, [(i, i) for i in range(50, 150)]) |
| |
| def test_null_keys_dropped_independent_of_num_ranges(self): |
| loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_null_loc", loc, [ |
| pa.Table.from_pydict({"k": [1, 2, None, 3], "row_id": [1, 2, 99, 3]}, schema=loc)]) |
| self._table("default.rj_null_in", ins, [ |
| pa.Table.from_pydict({"k": [1, None, 3, None]}, schema=ins)]) |
| expected = [(1, 1), (3, 3)] # null never matches; no duplicates |
| for num_ranges in (1, 5): |
| ds = range_join( |
| "default.rj_null_in", "default.rj_null_loc", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k", "row_id"], |
| num_ranges=num_ranges) |
| got = sorted((r["k"], r["row_id"]) for r in ds.take_all()) |
| self.assertEqual(got, expected) |
| |
| def test_pk_nonkey_range_col_untrusted(self): |
| # A PK table's non-PK column may be rewritten by merge (aggregation/partial-update) |
| # beyond the footer min/max, so its bounds are untrusted -> every split unknown. |
| schema = pa.schema([("id", pa.int64()), ("g", pa.int64())]) |
| self._table("default.rj_agg", schema, [ |
| pa.Table.from_pydict({"id": [1, 2], "g": [10, 20]}, schema=schema)], |
| primary_keys=["id"], options={"bucket": "1"}) |
| ranged, _ = rjmod._plan_ranged_splits( |
| "default.rj_agg", self.catalog_options, None, "g") |
| self.assertTrue(ranged) |
| self.assertTrue(all(lo is None and hi is None for _, lo, hi in ranged)) |
| |
| def test_masked_range_col_untrusted(self): |
| from pypaimon.catalog.table_query_auth import TableQueryAuthResult |
| from pypaimon.read.query_auth_split import QueryAuthSplit |
| |
| schema = pa.schema([("k", pa.int64())]) |
| name = "default.rj_masked" |
| self._table(name, schema, [ |
| pa.Table.from_pydict({"k": [1, 2, 3]}, schema=schema)]) |
| table = self.catalog.get_table(name) |
| splits = list(table.new_read_builder().new_scan().plan().splits()) |
| auth = TableQueryAuthResult( |
| None, {"k": "CAST(0 AS BIGINT)"}) |
| masked = [QueryAuthSplit(split, auth) for split in splits] |
| |
| self.assertFalse(rjmod._range_stats_trusted(table, masked, "k")) |
| |
| def test_partition_key_validation(self): |
| loc = pa.schema([("k", pa.int64()), ("v", pa.string())]) |
| self._table("default.rj_pv", loc, [ |
| pa.Table.from_pydict({"k": [1], "v": ["a"]}, schema=loc)]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_pv_in", ins, [ |
| pa.Table.from_pydict({"k": [1]}, schema=ins)]) |
| with self.assertRaises(ValueError): # not a partition column |
| range_join("default.rj_pv_in", "default.rj_pv", self.catalog_options, |
| on="k", left_projection=["k"], right_projection=["k", "v"], |
| right_partitions={"nope": "a"}) |
| |
| def test_num_ranges_validation(self): |
| loc = pa.schema([("k", pa.int64()), ("v", pa.string())]) |
| self._table("default.rj_nr", loc, [ |
| pa.Table.from_pydict({"k": [1, 2], "v": ["a", "b"]}, schema=loc)]) |
| ins = pa.schema([("k", pa.int64())]) |
| self._table("default.rj_nr_in", ins, [ |
| pa.Table.from_pydict({"k": [1]}, schema=ins)]) |
| for bad in (0, -1, "5", 2.0): |
| with self.assertRaises(ValueError): |
| range_join("default.rj_nr_in", "default.rj_nr", self.catalog_options, |
| on="k", left_projection=["k"], |
| right_projection=["k", "v"], num_ranges=bad) |
| |
| def test_rejects_unrangeable_key_types(self): |
| # Rejected at the driver (not inside a worker): nested and tz-aware keys. |
| arr = pa.schema([("k", pa.list_(pa.int64())), ("v", pa.string())]) |
| self._table("default.rj_arr_a", arr, []) |
| self._table("default.rj_arr_b", arr, []) |
| with self.assertRaisesRegex(ValueError, "must not be"): |
| range_join("default.rj_arr_a", "default.rj_arr_b", self.catalog_options, on="k") |
| ltz = pa.schema([("k", pa.timestamp("us", tz="UTC")), ("v", pa.string())]) |
| self._table("default.rj_ltz_a", ltz, []) |
| self._table("default.rj_ltz_b", ltz, []) |
| with self.assertRaisesRegex(ValueError, "must not be"): |
| range_join("default.rj_ltz_a", "default.rj_ltz_b", self.catalog_options, on="k") |
| # A nested SECOND key is rejected too (every key is validated, not just the range key). |
| multi = pa.schema([("k", pa.int64()), ("k2", pa.list_(pa.int64())), ("v", pa.string())]) |
| self._table("default.rj_mk_a", multi, []) |
| self._table("default.rj_mk_b", multi, []) |
| with self.assertRaisesRegex(ValueError, "nested/complex"): |
| range_join("default.rj_mk_a", "default.rj_mk_b", self.catalog_options, on=["k", "k2"]) |
| |
| |
| if __name__ == "__main__": |
| unittest.main() |
