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 7dc65b9f61 [python][ray] Add no-shuffle bucket join for co-bucketed
tables (#8397)
7dc65b9f61 is described below
commit 7dc65b9f6175629f6eb505369de41f8d15a277e3
Author: XiaoHongbo <[email protected]>
AuthorDate: Fri Jul 3 21:48:17 2026 +0800
[python][ray] Add no-shuffle bucket join for co-bucketed tables (#8397)
---
docs/docs/pypaimon/ray-data.md | 52 ++++
paimon-python/pypaimon/ray/__init__.py | 2 +
paimon-python/pypaimon/ray/bucket_join.py | 241 +++++++++++++++
.../pypaimon/tests/ray_bucket_join_test.py | 342 +++++++++++++++++++++
4 files changed, 637 insertions(+)
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index 960ab84213..fb56c54ac6 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -339,6 +339,58 @@ table_write = (
table_write.write_ray(ray_dataset)
```
+## Bucket Join
+
+`bucket_join` joins two **co-bucketed** tables (same bucket count and the same
+bucket-key) on the bucket-key, with **no global shuffle**: the same key lands
in
+the same bucket on both sides, so each bucket is read and joined in its own Ray
+task. It returns a `ray.data.Dataset` whose results stay distributed (never
+pulled into the driver).
+
+A common use is looking up a global `_ROW_ID` for a batch of keys without a
+shuffle join against a large table: keep a small co-bucketed `(key, _ROW_ID)`
+side table, `bucket_join` the incoming keys against it, then feed the resulting
+row ids into a row-id update.
+
+```python
+from pypaimon.ray import bucket_join
+
+ds = bucket_join(
+ left="database_name.incoming_keys", # co-bucketed table identifier
+ right="database_name.key_rowid", # co-bucketed table identifier
+ catalog_options={"warehouse": "/path/to/warehouse"},
+ on="url", # must equal the bucket-key
+ left_projection=["url"], # optional; must keep the join key
+ right_projection=["url", "row_id"], # optional; must keep the join key
+)
+# ds: ray.data.Dataset of the joined rows, e.g. {"url": ..., "row_id": ...}
+```
+
+**Parameters:**
+- `left` / `right`: identifiers of the two co-bucketed tables to join.
+- `on`: the join key(s). Must be exactly the bucket-key — equal keys only
+ co-locate by bucket when joining on the bucket-key.
+- `left_projection` / `right_projection`: optional column projections applied
on
+ read. If given, each must include the join key.
+- `join_type`: only `"inner"` is supported (an outer join would need the union
+ of buckets, which per-bucket intersection cannot produce).
+- `ray_remote_args`: Ray remote options applied to each per-bucket join task.
+
+**Returns:** a `ray.data.Dataset` of the joined rows.
+
+**Notes:**
+- Both tables must be fixed-bucket (`bucket > 0`) with the same bucket count
and
+ the same bucket-key (same column names, order, and types); otherwise
+ `bucket_join` raises. For primary-key tables that do not set `bucket-key`
+ explicitly, the bucket-key resolves to the (partition-trimmed) primary key.
+- The two sides must not share columns other than the join key, or the
+ underlying pyarrow join would collide; project them away with
+ `left_projection` / `right_projection` first.
+- Each side is planned at its own latest snapshot, and one bucket is joined by
a
+ single Ray task that reads the whole bucket into memory. Choose a bucket
count
+ that spreads keys evenly to avoid skewed, memory-heavy tasks.
+- Partitioned tables are not supported yet (bucket ids are per-partition).
+
## Merge Into
`merge_into` updates or deletes matched rows and optionally inserts unmatched
diff --git a/paimon-python/pypaimon/ray/__init__.py
b/paimon-python/pypaimon/ray/__init__.py
index 4280187956..ba88fe5c7a 100644
--- a/paimon-python/pypaimon/ray/__init__.py
+++ b/paimon-python/pypaimon/ray/__init__.py
@@ -16,6 +16,7 @@
# under the License.
from pypaimon.ray.ray_paimon import read_paimon, write_paimon
+from pypaimon.ray.bucket_join import bucket_join
from pypaimon.ray.data_evolution_merge_into import (
WhenMatched,
WhenNotMatched,
@@ -30,6 +31,7 @@ from pypaimon.ray.data_evolution_merge_transform import (
__all__ = [
"read_paimon",
"write_paimon",
+ "bucket_join",
"merge_into",
"WhenMatched",
"WhenNotMatched",
diff --git a/paimon-python/pypaimon/ray/bucket_join.py
b/paimon-python/pypaimon/ray/bucket_join.py
new file mode 100644
index 0000000000..818643ec95
--- /dev/null
+++ b/paimon-python/pypaimon/ray/bucket_join.py
@@ -0,0 +1,241 @@
+# 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.
+
+"""Bucket-aligned join on Ray for two co-bucketed Paimon tables.
+
+Same key -> same bucket on both sides, so each bucket is read and joined in
its own
+Ray task with no global shuffle -- the no-shuffle alternative to
``ray.data.join``.
+"""
+
+import threading
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+__all__ = ["bucket_join"]
+
+OnSpec = Union[str, Sequence[str]]
+
+
+def _norm(on: OnSpec) -> List[str]:
+ return [on] if isinstance(on, str) else list(on)
+
+
+def _key_type(table, col):
+ # Logical type without nullability -- a present key hashes the same either
way.
+ return str(table.field_dict[col].type).replace(" NOT NULL", "")
+
+
+def _bucketing(table):
+ # Resolved bucket keys (a PK table without an explicit bucket-key buckets
by its
+ # trimmed primary key) plus the bucket function: same key co-locates only
under both.
+ return (table.options.bucket(),
+ list(table.table_schema.bucket_keys),
+ table.table_schema.options.get("bucket-function.type", "default"))
+
+
+# Per-worker table cache, keyed by schema id (so a schema change invalidates
it) and
+# lock-guarded against concurrent tasks. Planning always loads a fresh table.
+_TABLE_CACHE: Dict = {}
+_TABLE_CACHE_LOCK = threading.Lock()
+
+
+def _get_table(table_id, catalog_options, schema_id=None):
+ from pypaimon.catalog.catalog_factory import CatalogFactory
+ if schema_id is None: # planning: always load the latest schema
+ return CatalogFactory.create(catalog_options).get_table(table_id)
+ key = (table_id, tuple(sorted(catalog_options.items())), schema_id)
+ with _TABLE_CACHE_LOCK:
+ table = _TABLE_CACHE.get(key)
+ if table is None:
+ table = CatalogFactory.create(catalog_options).get_table(table_id)
+ if table.table_schema.id != schema_id:
+ # get_table loads the latest schema; a mismatch means the
schema moved
+ # after the driver planned, so the split plan is stale -- fail
fast.
+ raise ValueError(
+ f"{table_id} schema changed during bucket_join (planned
{schema_id}, "
+ f"now {table.table_schema.id}); retry.")
+ _TABLE_CACHE[key] = table
+ return table
+
+
+def _read_builder(table_id, catalog_options, projection, schema_id=None):
+ rb = _get_table(table_id, catalog_options, schema_id).new_read_builder()
+ return rb.with_projection(projection) if projection is not None else rb
+
+
+def _plan_splits_by_bucket(table_id, catalog_options, projection,
expected_total_buckets):
+ """Plan the manifest and group splits by bucket (driver-side).
+
+ Returns ``(by_bucket, schema_id)``: the schema id of the table instance
that
+ built this plan, so workers validate against the schema the plan was made
with
+ (not a possibly-newer one loaded earlier by the caller).
+ """
+ from pypaimon.common.options.core_options import CoreOptions
+ table = _get_table(table_id, catalog_options) # fresh, latest schema
+ schema_id = table.table_schema.id
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ if snapshot is None:
+ return {}, schema_id
+ # Pin the guard and the split plan to one snapshot, else a commit between
the two
+ # manifest reads could slip stale-bucket files past the guard. Drop any
existing
+ # scan.mode / point-in-time options first so snapshot-id doesn't clash
with them.
+ opts = table.options.options
+ for key in (CoreOptions.SCAN_MODE, CoreOptions.SCAN_SNAPSHOT_ID,
+ CoreOptions.SCAN_TAG_NAME, CoreOptions.SCAN_WATERMARK,
+ CoreOptions.SCAN_TIMESTAMP, CoreOptions.SCAN_TIMESTAMP_MILLIS,
+ CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP,
+ CoreOptions.SCAN_FILE_CREATION_TIME_MILLIS,
+ CoreOptions.SCAN_CREATION_TIME_MILLIS):
+ opts.data.pop(key.key(), None)
+ opts.set(CoreOptions.SCAN_SNAPSHOT_ID, snapshot.id)
+ rb = table.new_read_builder()
+ scan = (rb.with_projection(projection) if projection is not None else
rb).new_scan()
+ # Guard against a rescaled table (old files under a different
total_buckets, which
+ # splits -- bucket only -- can't tell apart); read the entries once for it.
+ fs = scan.file_scanner
+ entries = fs.plan_files()
+ stale = {e.total_buckets for e in entries if e.total_buckets !=
expected_total_buckets}
+ if stale:
+ raise ValueError(
+ f"bucket_join needs {table_id} fully in bucket count
{expected_total_buckets}, "
+ f"but files exist under {sorted(stale)} (rescale in progress);
rewrite first.")
+ # Reuse those entries: scan.plan() re-reads plan_files() (append/pk)
otherwise.
+ fs.plan_files = lambda: entries
+ by_bucket = {}
+ for s in scan.plan().splits():
+ by_bucket.setdefault(s.bucket, []).append(s)
+ return by_bucket, schema_id
+
+
+def _read_splits(table_id, catalog_options, projection, splits, schema_id):
+ # Snapshot-independent but schema-dependent -> cache by schema id (in
_get_table).
+ return _read_builder(
+ table_id, catalog_options, projection,
schema_id).new_read().to_arrow(splits)
+
+
+def bucket_join(
+ left: str,
+ right: str,
+ catalog_options: Dict[str, str],
+ *,
+ on: OnSpec,
+ left_projection: Optional[List[str]] = None,
+ right_projection: Optional[List[str]] = None,
+ join_type: str = "inner",
+ ray_remote_args: Optional[Dict[str, Any]] = None,
+) -> "ray.data.Dataset":
+ """Join two co-bucketed tables (same bucket count + bucket-key, joined on
the
+ bucket-key) with no global shuffle. ``on`` must equal the bucket-key. The
two
+ sides must not share column names other than the join key (pyarrow ``join``
+ would otherwise collide). Returns a ``ray.data.Dataset``."""
+ import ray
+ from pypaimon.catalog.catalog_factory import CatalogFactory
+
+ if not hasattr(ray.data, "from_arrow_refs"):
+ raise RuntimeError(
+ "bucket_join needs a Ray version with ray.data.from_arrow_refs; "
+ f"installed ray is {ray.__version__}.")
+
+ on_cols = _norm(on)
+ cat = CatalogFactory.create(catalog_options)
+ ltable, rtable = cat.get_table(left), cat.get_table(right)
+ lcount, lkey, lfunc = _bucketing(ltable)
+ rcount, rkey, rfunc = _bucketing(rtable)
+
+ if ltable.partition_keys or rtable.partition_keys:
+ # Bucket numbers are per-partition, so the same bucket id lives in
every
+ # partition -- grouping splits by bucket alone would join across
partitions.
+ # Supporting this needs grouping by (partition, bucket); not done yet.
+ raise ValueError(
+ "bucket_join does not support partitioned tables yet; got
partition keys "
+ f"{left}={ltable.partition_keys},
{right}={rtable.partition_keys}.")
+ if not lcount or lcount <= 0 or not rcount or rcount <= 0:
+ raise ValueError(
+ "bucket_join requires both tables to be fixed-bucket (bucket > 0);
"
+ f"got {left}={lcount}, {right}={rcount}.")
+ if lcount != rcount:
+ raise ValueError(
+ f"bucket_join requires the same bucket count; {left}={lcount},
{right}={rcount}.")
+ if lkey != rkey:
+ raise ValueError(
+ f"bucket_join requires the same bucket-key; {left}={lkey},
{right}={rkey}.")
+ if lfunc != rfunc:
+ # Different bucket functions hash the same key to different buckets.
+ raise ValueError(
+ f"bucket_join requires the same bucket-function.type;
{left}={lfunc}, {right}={rfunc}.")
+ if on_cols != lkey:
+ raise ValueError(
+ f"bucket_join requires the join key to be the bucket-key {lkey};
got on={on_cols}. "
+ "Equal keys only co-locate by bucket when joining on the
bucket-key "
+ "(the comparison is order-sensitive for composite keys).")
+ # Same name isn't enough: differing key types (INT vs BIGINT) can hash
apart and
+ # silently drop matches (types compared without nullability).
+ key_type_mismatch = [
+ (c, _key_type(ltable, c), _key_type(rtable, c))
+ for c in on_cols
+ if _key_type(ltable, c) != _key_type(rtable, c)
+ ]
+ if key_type_mismatch:
+ raise ValueError(
+ "bucket_join requires the bucket-key columns to have the same type
on both "
+ f"sides; mismatched (column, left, right): {key_type_mismatch}.")
+ if join_type != "inner":
+ # Outer joins would need the union of buckets (a bucket missing on one
side
+ # still emits rows); only inner is correct with the per-bucket
intersection.
+ raise ValueError(f"bucket_join currently supports only
join_type='inner'; got {join_type!r}.")
+
+ # The join key must survive projection on both sides, or ``Table.join``
has no key.
+ if left_projection is not None and not set(on_cols) <=
set(left_projection):
+ raise ValueError(
+ f"left_projection must include the join key {on_cols}; got
{left_projection}.")
+ if right_projection is not None and not set(on_cols) <=
set(right_projection):
+ raise ValueError(
+ f"right_projection must include the join key {on_cols}; got
{right_projection}.")
+ # The two sides must not share non-key columns, or pyarrow's join collides
on them.
+ # Check up front (against the projected columns) instead of failing inside
a task.
+ lcols = left_projection if left_projection is not None else
ltable.field_names
+ rcols = right_projection if right_projection is not None else
rtable.field_names
+ collisions = sorted((set(lcols) & set(rcols)) - set(on_cols))
+ if collisions:
+ raise ValueError(
+ f"bucket_join sides must not share columns other than the join key
{on_cols}; "
+ f"both have {collisions}. Project or rename them away.")
+
+ # Plan each side's manifest once (driver-side, split metadata only -- the
join
+ # results stay distributed below), then dispatch per-bucket splits to the
tasks.
+ left_by_bucket, l_schema_id = _plan_splits_by_bucket(
+ left, catalog_options, left_projection, lcount)
+ right_by_bucket, r_schema_id = _plan_splits_by_bucket(
+ right, catalog_options, right_projection, rcount)
+
+ def _join_bucket(left_splits, right_splits):
+ left_t = _read_splits(left, catalog_options, left_projection,
left_splits, l_schema_id)
+ right_t = _read_splits(right, catalog_options, right_projection,
right_splits, r_schema_id)
+ return left_t.join(right_t, keys=on_cols, join_type=join_type)
+
+ # ``@ray.remote()`` (empty parens) is rejected by Ray, so wrap
conditionally.
+ remote_fn = ray.remote(**ray_remote_args)(_join_bucket) if ray_remote_args
else ray.remote(_join_bucket)
+ # Inner join: only buckets present on both sides can match.
+ buckets = sorted(set(left_by_bucket) & set(right_by_bucket))
+ if not buckets:
+ # No shared bucket: empty result, but keep the join schema (join two
empties).
+ empty = _read_splits(left, catalog_options, left_projection, [],
l_schema_id).join(
+ _read_splits(right, catalog_options, right_projection, [],
r_schema_id),
+ keys=on_cols, join_type=join_type)
+ return ray.data.from_arrow(empty)
+ # Keep each bucket's result as a distributed object ref -- never pulled
into the driver.
+ refs = [remote_fn.remote(left_by_bucket[b], right_by_bucket[b]) for b in
buckets]
+ return ray.data.from_arrow_refs(refs)
diff --git a/paimon-python/pypaimon/tests/ray_bucket_join_test.py
b/paimon-python/pypaimon/tests/ray_bucket_join_test.py
new file mode 100644
index 0000000000..398ac2a37c
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_bucket_join_test.py
@@ -0,0 +1,342 @@
+# 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.scanner.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()
