| # 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 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 bucket_join | |
| # The package attribute ``pypaimon.ray.bucket_join`` is the function (it shadows the | |
| # submodule); import the module explicitly to reach its internal helpers. | |
| bjmod = importlib.import_module("pypaimon.ray.bucket_join") | |
| class RayBucketJoinTest(unittest.TestCase): | |
| """Bucket-aligned join between two HASH_FIXED tables must equal a global join, | |
| with each bucket joined only against the same bucket (no cross-bucket shuffle).""" | |
| NUM_BUCKETS = 8 | |
| @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 _bucketed_table(self, name, schema, key, data, primary_keys=None, extra_opts=None): | |
| opts = {"bucket": str(self.NUM_BUCKETS)} | |
| if primary_keys is None: # append table: bucket-key must be explicit | |
| opts["bucket-key"] = key | |
| if extra_opts: | |
| opts.update(extra_opts) | |
| # PK table: leave bucket-key unset so it defaults to the primary key. | |
| self.catalog.create_table( | |
| name, | |
| Schema.from_pyarrow_schema(schema, primary_keys=primary_keys, options=opts), | |
| False) | |
| t = self.catalog.get_table(name) | |
| 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 _create_bucketed(self, name, schema, key, num_buckets, | |
| partition_keys=None, extra_opts=None): | |
| opts = {"bucket": str(num_buckets), "bucket-key": key} | |
| if extra_opts: | |
| opts.update(extra_opts) | |
| self.catalog.create_table( | |
| name, | |
| Schema.from_pyarrow_schema(schema, partition_keys=partition_keys, options=opts), | |
| False) | |
| return name | |
| def test_bucket_join_matches_global_join(self): | |
| loc_schema = pa.schema([("url", pa.string()), ("row_id", pa.int64())]) | |
| in_schema = pa.schema([("url", pa.string())]) | |
| self._bucketed_table( | |
| "default.locator", loc_schema, "url", | |
| pa.Table.from_pydict({"url": [f"u{i}" for i in range(1000)], | |
| "row_id": list(range(1000))}, schema=loc_schema)) | |
| self._bucketed_table( | |
| "default.input", in_schema, "url", | |
| pa.Table.from_pydict({"url": [f"u{i}" for i in range(0, 400)]}, schema=in_schema)) | |
| ds = bucket_join( | |
| "default.input", "default.locator", self.catalog_options, | |
| on="url", left_projection=["url"], right_projection=["url", "row_id"]) | |
| got = {r["url"]: r["row_id"] for r in ds.take_all()} | |
| # every input url (u0..u399) is matched to its locator row_id | |
| self.assertEqual(set(got), {f"u{i}" for i in range(400)}) | |
| self.assertEqual(got["u0"], 0) | |
| self.assertEqual(got["u399"], 399) | |
| self.assertTrue(all(got[f"u{i}"] == i for i in range(400))) | |
| def test_fan_out_one_url_many_row_ids(self): | |
| # A url may map to several locator rows; every match must be emitted. | |
| loc_schema = pa.schema([("url", pa.string()), ("row_id", pa.int64())]) | |
| in_schema = pa.schema([("url", pa.string())]) | |
| self._bucketed_table( | |
| "default.loc_fan", loc_schema, "url", | |
| pa.Table.from_pydict({"url": ["u0", "u0", "u1"], "row_id": [0, 1, 2]}, | |
| schema=loc_schema)) | |
| self._bucketed_table( | |
| "default.in_fan", in_schema, "url", | |
| pa.Table.from_pydict({"url": ["u0"]}, schema=in_schema)) | |
| ds = bucket_join( | |
| "default.in_fan", "default.loc_fan", self.catalog_options, | |
| on="url", left_projection=["url"], right_projection=["url", "row_id"]) | |
| self.assertEqual(sorted(r["row_id"] for r in ds.take_all()), [0, 1]) | |
| def test_join_with_explicit_scan_mode(self): | |
| # An explicit scan.mode on the tables must not clash with the internal | |
| # snapshot pin during planning. | |
| loc = pa.schema([("url", pa.string()), ("row_id", pa.int64())]) | |
| ins = pa.schema([("url", pa.string())]) | |
| mode = {"scan.mode": "latest-full"} | |
| self._bucketed_table( | |
| "default.sm_loc", loc, "url", | |
| pa.Table.from_pydict({"url": [f"u{i}" for i in range(50)], | |
| "row_id": list(range(50))}, schema=loc), | |
| extra_opts=mode) | |
| self._bucketed_table( | |
| "default.sm_in", ins, "url", | |
| pa.Table.from_pydict({"url": [f"u{i}" for i in range(20)]}, schema=ins), | |
| extra_opts=mode) | |
| ds = bucket_join( | |
| "default.sm_in", "default.sm_loc", self.catalog_options, | |
| on="url", left_projection=["url"], right_projection=["url", "row_id"]) | |
| got = {r["url"]: r["row_id"] for r in ds.take_all()} | |
| self.assertEqual(got, {f"u{i}": i for i in range(20)}) | |
| def test_dispatches_one_task_per_shared_bucket(self): | |
| # No cross-bucket shuffle: exactly one Ray task (object ref) per shared bucket. | |
| loc = pa.schema([("url", pa.string()), ("row_id", pa.int64())]) | |
| ins = pa.schema([("url", pa.string())]) | |
| self._bucketed_table( | |
| "default.disp_loc", loc, "url", | |
| pa.Table.from_pydict({"url": [f"u{i}" for i in range(200)], | |
| "row_id": list(range(200))}, schema=loc)) | |
| self._bucketed_table( | |
| "default.disp_in", ins, "url", | |
| pa.Table.from_pydict({"url": [f"u{i}" for i in range(100)]}, schema=ins)) | |
| lbb, _ = bjmod._plan_splits_by_bucket( | |
| "default.disp_in", self.catalog_options, ["url"], self.NUM_BUCKETS) | |
| rbb, _ = bjmod._plan_splits_by_bucket( | |
| "default.disp_loc", self.catalog_options, ["url", "row_id"], self.NUM_BUCKETS) | |
| shared = set(lbb) & set(rbb) | |
| self.assertGreater(len(shared), 1) # genuinely spread across buckets | |
| 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 = bucket_join( | |
| "default.disp_in", "default.disp_loc", self.catalog_options, | |
| on="url", left_projection=["url"], right_projection=["url", "row_id"]) | |
| ds.take_all() | |
| self.assertEqual(captured["n"], len(shared)) | |
| def test_composite_bucket_key(self): | |
| # Happy path for a multi-column bucket-key joined on both columns. | |
| loc = pa.schema([("a", pa.string()), ("b", pa.int64()), ("row_id", pa.int64())]) | |
| ins = pa.schema([("a", pa.string()), ("b", pa.int64())]) | |
| self._bucketed_table( | |
| "default.comp_loc", loc, "a,b", | |
| pa.Table.from_pydict({"a": [f"k{i}" for i in range(50)], "b": list(range(50)), | |
| "row_id": list(range(50))}, schema=loc)) | |
| self._bucketed_table( | |
| "default.comp_in", ins, "a,b", | |
| pa.Table.from_pydict({"a": [f"k{i}" for i in range(20)], "b": list(range(20))}, | |
| schema=ins)) | |
| ds = bucket_join( | |
| "default.comp_in", "default.comp_loc", self.catalog_options, | |
| on=["a", "b"], left_projection=["a", "b"], right_projection=["a", "b", "row_id"]) | |
| got = {(r["a"], r["b"]): r["row_id"] for r in ds.take_all()} | |
| self.assertEqual(len(got), 20) | |
| self.assertTrue(all(got[(f"k{i}", i)] == i for i in range(20))) | |
| def test_read_cache_by_schema_id_and_staleness(self): | |
| # Same schema id -> cached (reused). A schema id that is no longer current means | |
| # the plan is stale -> fail fast instead of reading with the wrong schema. | |
| self._create_bucketed("default.cache_t", pa.schema([("url", pa.string())]), "url", 8) | |
| bjmod._TABLE_CACHE.clear() | |
| opts = self.catalog_options | |
| sid = self.catalog.get_table("default.cache_t").table_schema.id | |
| t1 = bjmod._get_table("default.cache_t", opts, sid) | |
| t2 = bjmod._get_table("default.cache_t", opts, sid) | |
| self.assertIs(t1, t2) # same schema id -> cached | |
| with self.assertRaises(ValueError): # stale plan: requested schema not current | |
| bjmod._get_table("default.cache_t", opts, sid + 99) | |
| def test_empty_result_keeps_schema(self): | |
| # No shared bucket -> 0 rows, but the join schema must survive. | |
| loc_schema = pa.schema([("url", pa.string()), ("row_id", pa.int64())]) | |
| in_schema = pa.schema([("url", pa.string())]) | |
| self._bucketed_table( | |
| "default.loc_empty", loc_schema, "url", | |
| pa.Table.from_pydict({"url": ["u0", "u1"], "row_id": [0, 1]}, schema=loc_schema)) | |
| self._bucketed_table( | |
| "default.in_empty", in_schema, "url", | |
| pa.Table.from_pydict({"url": []}, schema=in_schema)) # no rows -> no buckets | |
| ds = bucket_join( | |
| "default.in_empty", "default.loc_empty", self.catalog_options, | |
| on="url", left_projection=["url"], right_projection=["url", "row_id"]) | |
| self.assertEqual(ds.count(), 0) | |
| self.assertIn("row_id", ds.schema().names) | |
| def test_rejects_different_bucket_count(self): | |
| sch = pa.schema([("url", pa.string())]) | |
| self._create_bucketed("default.cnt_8", sch, "url", 8) | |
| self._create_bucketed("default.cnt_16", sch, "url", 16) | |
| with self.assertRaises(ValueError): | |
| bucket_join("default.cnt_8", "default.cnt_16", self.catalog_options, on="url") | |
| def test_rejects_different_bucket_key(self): | |
| sch = pa.schema([("url", pa.string()), ("k", pa.string())]) | |
| self._create_bucketed("default.by_url", sch, "url", 8) | |
| self._create_bucketed("default.by_k", sch, "k", 8) | |
| with self.assertRaises(ValueError): | |
| bucket_join("default.by_url", "default.by_k", self.catalog_options, on="url") | |
| def test_rejects_join_key_not_bucket_key(self): | |
| sch = pa.schema([("url", pa.string()), ("k", pa.string())]) | |
| self._create_bucketed("default.k1", sch, "url", 8) | |
| self._create_bucketed("default.k2", sch, "url", 8) | |
| with self.assertRaises(ValueError): # on=k but bucket-key=url | |
| bucket_join("default.k1", "default.k2", self.catalog_options, on="k") | |
| def test_primary_key_default_bucket_key(self): | |
| # PK tables bucket by their primary key without an explicit bucket-key option; | |
| # bucket_join must resolve that and join on the PK. | |
| loc_schema = pa.schema([("url", pa.string()), ("row_id", pa.int64())]) | |
| in_schema = pa.schema([("url", pa.string())]) | |
| self._bucketed_table( | |
| "default.pk_loc", loc_schema, "url", | |
| pa.Table.from_pydict({"url": [f"u{i}" for i in range(100)], | |
| "row_id": list(range(100))}, schema=loc_schema), | |
| primary_keys=["url"]) | |
| self._bucketed_table( | |
| "default.pk_in", in_schema, "url", | |
| pa.Table.from_pydict({"url": [f"u{i}" for i in range(40)]}, schema=in_schema), | |
| primary_keys=["url"]) | |
| ds = bucket_join( | |
| "default.pk_in", "default.pk_loc", self.catalog_options, | |
| on="url", left_projection=["url"], right_projection=["url", "row_id"]) | |
| got = {r["url"]: r["row_id"] for r in ds.take_all()} | |
| self.assertEqual(set(got), {f"u{i}" for i in range(40)}) | |
| self.assertTrue(all(got[f"u{i}"] == i for i in range(40))) | |
| def test_rejects_projection_missing_join_key(self): | |
| sch = pa.schema([("url", pa.string()), ("v", pa.int64())]) | |
| self._create_bucketed("default.pmj1", sch, "url", 8) | |
| self._create_bucketed("default.pmj2", sch, "url", 8) | |
| with self.assertRaises(ValueError): # left projection drops the join key | |
| bucket_join("default.pmj1", "default.pmj2", self.catalog_options, | |
| on="url", left_projection=["v"], right_projection=["url"]) | |
| def test_rejects_colliding_columns(self): | |
| # Both sides expose a non-key column "v" -> pyarrow join would collide. | |
| sch = pa.schema([("url", pa.string()), ("v", pa.int64())]) | |
| self._create_bucketed("default.col1", sch, "url", 8) | |
| self._create_bucketed("default.col2", sch, "url", 8) | |
| with self.assertRaises(ValueError): | |
| bucket_join("default.col1", "default.col2", self.catalog_options, on="url") | |
| def test_rejects_different_bucket_function(self): | |
| # Same bucket-key/count but different bucket function -> same key may land in | |
| # different buckets, so co-location is not guaranteed. | |
| sch = pa.schema([("k", pa.int64())]) | |
| self._create_bucketed("default.bf_default", sch, "k", 8) | |
| self._create_bucketed("default.bf_mod", sch, "k", 8, | |
| extra_opts={"bucket-function.type": "mod"}) | |
| with self.assertRaises(ValueError): | |
| bucket_join("default.bf_default", "default.bf_mod", self.catalog_options, on="k") | |
| def test_rejects_mismatched_key_type(self): | |
| # Same bucket-key name but different type hashes differently -> reject. | |
| self._create_bucketed("default.ty_str", pa.schema([("k", pa.string())]), "k", 8) | |
| self._create_bucketed("default.ty_int", pa.schema([("k", pa.int64())]), "k", 8) | |
| with self.assertRaises(ValueError): | |
| bucket_join("default.ty_str", "default.ty_int", self.catalog_options, on="k") | |
| def test_rejects_rescaled_mixed_buckets(self): | |
| # Files left under an old bucket count (rescale not yet rewritten) would carry a | |
| # different total_buckets; the same bucket id must not be treated as co-located. | |
| import types | |
| from pypaimon.read.file_scanner import FileScanner | |
| sch = pa.schema([("url", pa.string())]) | |
| self._bucketed_table("default.rs_a", sch, "url", | |
| pa.Table.from_pydict({"url": ["u0", "u1"]}, schema=sch)) | |
| self._bucketed_table("default.rs_b", sch, "url", | |
| pa.Table.from_pydict({"url": ["u0"]}, schema=sch)) | |
| stale = [types.SimpleNamespace(total_buckets=4)] # a file from a 4-bucket era | |
| with mock.patch.object(FileScanner, "plan_files", return_value=stale): | |
| with self.assertRaises(ValueError): | |
| bucket_join("default.rs_a", "default.rs_b", self.catalog_options, on="url") | |
| def test_rejects_partitioned_table(self): | |
| # Bucket ids are per-partition, so bucket-only grouping would join across | |
| # partitions; partitioned tables are rejected until (partition, bucket) grouping. | |
| sch = pa.schema([("url", pa.string()), ("dt", pa.string())]) | |
| self._create_bucketed("default.part_p", sch, "url", 8, partition_keys=["dt"]) | |
| self._create_bucketed("default.part_np", sch, "url", 8) | |
| with self.assertRaises(ValueError): | |
| bucket_join("default.part_p", "default.part_np", self.catalog_options, on="url") | |
| def test_rejects_non_inner_join(self): | |
| sch = pa.schema([("url", pa.string())]) | |
| self._create_bucketed("default.ji1", sch, "url", 8) | |
| self._create_bucketed("default.ji2", sch, "url", 8) | |
| with self.assertRaises(ValueError): | |
| bucket_join("default.ji1", "default.ji2", self.catalog_options, | |
| on="url", join_type="left outer") | |
| if __name__ == "__main__": | |
| unittest.main() |