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 0dd0bcc31c [python] Add range_join: shuffle-free Ray join for
range-clustered tables (#8738)
0dd0bcc31c is described below
commit 0dd0bcc31c92fbe796a304692d5a1bf2cb106045
Author: XiaoHongbo <[email protected]>
AuthorDate: Fri Jul 31 10:17:08 2026 +0800
[python] Add range_join: shuffle-free Ray join for range-clustered tables
(#8738)
---
.github/workflows/paimon-python-checks.yml | 3 +-
docs/docs/pypaimon/ray-data.md | 29 ++
paimon-python/pypaimon/ray/__init__.py | 2 +
paimon-python/pypaimon/ray/bucket_join.py | 75 +--
paimon-python/pypaimon/ray/join_common.py | 92 ++++
paimon-python/pypaimon/ray/range_join.py | 492 +++++++++++++++++++
.../pypaimon/tests/ray_range_join_test.py | 546 +++++++++++++++++++++
7 files changed, 1180 insertions(+), 59 deletions(-)
diff --git a/.github/workflows/paimon-python-checks.yml
b/.github/workflows/paimon-python-checks.yml
index b11ae24e3d..d8ff36e5f9 100755
--- a/.github/workflows/paimon-python-checks.yml
+++ b/.github/workflows/paimon-python-checks.yml
@@ -292,7 +292,8 @@ jobs:
python -m pip install --no-cache-dir -q ray==$ray_version
python -c "import ray; print(f'Ray version: {ray.__version__}')"
python -c "from packaging.version import parse; import ray; assert
parse(ray.__version__) == parse('$ray_version'), f'Expected Ray $ray_version,
got {ray.__version__}'"
- python -m pytest pypaimon/tests/ray_data_test.py::RayDataTest -v
--tb=short || {
+ python -m pytest pypaimon/tests/ray_data_test.py::RayDataTest \
+ pypaimon/tests/ray_range_join_test.py::RayRangeJoinTest -v
--tb=short || {
echo "Tests failed for Ray $ray_version"; python -m pip
uninstall -y ray; exit 1;
}
python -m pip uninstall -y ray
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index 88c0a7143a..7f56169eb4 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -397,6 +397,35 @@ ds = bucket_join(
that spreads keys evenly to avoid skewed, memory-heavy tasks.
- Partitioned tables are not supported yet (bucket ids are per-partition).
+## Range Join
+
+`range_join` joins tables clustered by the first join key without a global
+shuffle. Each key range runs in one Ray task.
+
+```python
+from pypaimon.ray import range_join
+
+ds = range_join(
+ left="database_name.incoming_keys",
+ right="database_name.key_rowid",
+ catalog_options={"warehouse": "/path/to/warehouse"},
+ left_on="url",
+ right_on="lookup_url",
+ left_projection=["url"],
+ right_projection=["lookup_url", "row_id"],
+ left_partitions={"dt": "2026-07-30"}, # optional
+ num_ranges=64, # optional
+)
+```
+
+Use `on="url"` when key names match. Multiple keys are supported; the first
+defines ranges. Only inner join is supported.
+
+Manifest/key stats are preferred; Parquet footers are the fallback. Missing
+stats safely reduce parallelism, possibly to one task. Unclustered files may be
+read repeatedly. Float/double and local-time-zone timestamp range keys are not
+supported.
+
## 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 bc91c8da45..03de068249 100644
--- a/paimon-python/pypaimon/ray/__init__.py
+++ b/paimon-python/pypaimon/ray/__init__.py
@@ -17,6 +17,7 @@
from pypaimon.ray.ray_paimon import map_with_blobs, read_paimon, write_paimon
from pypaimon.ray.bucket_join import bucket_join
+from pypaimon.ray.range_join import range_join
from pypaimon.ray.data_evolution_merge_into import (
WhenMatched,
WhenNotMatched,
@@ -35,6 +36,7 @@ __all__ = [
"map_with_blobs",
"write_paimon",
"bucket_join",
+ "range_join",
"merge_into",
"update_by_row_id",
"read_by_row_id",
diff --git a/paimon-python/pypaimon/ray/bucket_join.py
b/paimon-python/pypaimon/ray/bucket_join.py
index 818643ec95..2b377126df 100644
--- a/paimon-python/pypaimon/ray/bucket_join.py
+++ b/paimon-python/pypaimon/ray/bucket_join.py
@@ -21,22 +21,21 @@ Same key -> same bucket on both sides, so each bucket is
read and joined in its
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
+from typing import Any, Dict, List, Optional
+
+from pypaimon.ray.join_common import (
+ OnSpec,
+ get_table as _shared_get_table,
+ key_type as _key_type,
+ norm_on as _norm,
+ pin_latest_snapshot,
+ read_splits,
+)
+# The table cache now lives in join_common; keep the old name for
callers/tests.
+from pypaimon.ray.join_common import _TABLE_CACHE # noqa: F401
__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
@@ -46,34 +45,8 @@ def _bucketing(table):
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
+ return _shared_get_table(table_id, catalog_options, schema_id,
"bucket_join")
def _plan_splits_by_bucket(table_id, catalog_options, projection,
expected_total_buckets):
@@ -83,24 +56,12 @@ def _plan_splits_by_bucket(table_id, catalog_options,
projection, expected_total
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)
+ # manifest reads could slip stale-bucket files past the guard.
+ if pin_latest_snapshot(table) is None:
+ return {}, schema_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
@@ -121,9 +82,7 @@ def _plan_splits_by_bucket(table_id, catalog_options,
projection, expected_total
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)
+ return read_splits(table_id, catalog_options, projection, splits,
schema_id, "bucket_join")
def bucket_join(
diff --git a/paimon-python/pypaimon/ray/join_common.py
b/paimon-python/pypaimon/ray/join_common.py
new file mode 100644
index 0000000000..0892ea9c08
--- /dev/null
+++ b/paimon-python/pypaimon/ray/join_common.py
@@ -0,0 +1,92 @@
+# 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.
+
+"""Shared driver/worker helpers for the co-located Ray joins (bucket_join,
range_join)."""
+
+import threading
+from typing import Dict, List, Optional, Sequence, Union
+
+OnSpec = Union[str, Sequence[str]]
+
+
+def norm_on(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 compares the same
either way.
+ return str(table.field_dict[col].type).replace(" NOT NULL", "")
+
+
+# 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, join_name="join"):
+ 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 {join_name} (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,
join_name="join"):
+ rb = get_table(table_id, catalog_options, schema_id,
join_name).new_read_builder()
+ return rb.with_projection(projection) if projection is not None else rb
+
+
+def read_splits(table_id, catalog_options, projection, splits, schema_id,
+ join_name="join", predicate=None):
+ # Snapshot-independent but schema-dependent -> cache by schema id (in
get_table).
+ rb = read_builder(table_id, catalog_options, projection, schema_id,
join_name)
+ if predicate is not None:
+ rb = rb.with_filter(predicate)
+ return rb.new_read().to_arrow(splits)
+
+
+def pin_latest_snapshot(table) -> Optional[int]:
+ """Pin the table instance to its latest snapshot; returns the snapshot id
or None
+ when the table is empty. Pinning keeps every manifest read of a plan
consistent."""
+ from pypaimon.common.options.core_options import CoreOptions
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ if snapshot is None:
+ return None
+ # 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)
+ return snapshot.id
diff --git a/paimon-python/pypaimon/ray/range_join.py
b/paimon-python/pypaimon/ray/range_join.py
new file mode 100644
index 0000000000..6886d4bfce
--- /dev/null
+++ b/paimon-python/pypaimon/ray/range_join.py
@@ -0,0 +1,492 @@
+# 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.
+
+"""Shuffle-free Ray join for tables clustered by the first join key.
+
+Ranges use manifest/key stats with a Parquet-footer fallback. Missing stats are
+safe: affected splits join every overlapping range and are filtered in memory.
+"""
+
+import logging
+import threading
+from typing import Any, Dict, List, Optional
+
+from pypaimon.ray.join_common import (
+ OnSpec,
+ get_table,
+ key_type,
+ norm_on,
+ pin_latest_snapshot,
+ read_splits,
+)
+
+__all__ = ["range_join"]
+
+_LOG = logging.getLogger(__name__)
+
+_MAX_RANGES = 512
+# Cap total re-read to this many full scans of the read bytes (see
_bounded_ranges).
+_REREAD_BUDGET = 2
+
+
+def _stats_range(stats, field_id, key_type):
+ """Read matching field-id/type bounds from manifest stats."""
+ from pypaimon.schema.data_types import PyarrowFieldParser
+
+ min_row, max_row = stats.min_values, stats.max_values
+ min_fields = getattr(min_row, "fields", []) or []
+ max_fields = getattr(max_row, "fields", []) or []
+ idx = next((i for i, field in enumerate(min_fields) if field.id ==
field_id), None)
+ if idx is None or idx >= len(max_fields) or max_fields[idx].id != field_id:
+ return None
+ try:
+ stored_type = PyarrowFieldParser.from_paimon_type(min_fields[idx].type)
+ max_type = PyarrowFieldParser.from_paimon_type(max_fields[idx].type)
+ if stored_type != key_type or max_type != key_type:
+ return None
+ lo, hi = min_row.get_field(idx), max_row.get_field(idx)
+ except Exception:
+ return None
+ return None if lo is None or hi is None else (lo, hi)
+
+
+def _file_stats_range(file, field_id, key_type):
+ """Prefer key stats, then value stats."""
+ return (_stats_range(file.key_stats, field_id, key_type)
+ or _stats_range(file.value_stats, field_id, key_type))
+
+
+def _file_identity(file):
+ # External paths disambiguate imported files.
+ return file.external_path or file.file_name
+
+
+def _manifest_stats_by_file(table, field_id, key_type, snapshot_id):
+ """Read active-file stats for the pinned snapshot."""
+ from pypaimon.manifest.manifest_file_manager import ManifestFileManager
+ from pypaimon.manifest.manifest_list_manager import ManifestListManager
+
+ try:
+ snapshot = table.snapshot_manager().get_snapshot_by_id(snapshot_id)
+ manifests = ManifestListManager(table).read_all(snapshot)
+ entries = ManifestFileManager(table).read_entries_parallel(
+ manifests, drop_stats=False)
+ except Exception as e:
+ # Stats are optional.
+ _LOG.warning(
+ "range_join: manifest stats read failed (%s); falling back to file
footers", e)
+ return {}
+
+ by_file, ambiguous = {}, set()
+ for entry in entries:
+ file = entry.file
+ rng = _file_stats_range(file, field_id, key_type)
+ if rng is None:
+ continue
+ identity = _file_identity(file)
+ previous = by_file.get(identity)
+ if previous is not None and previous != rng:
+ ambiguous.add(identity)
+ else:
+ by_file[identity] = rng
+ for identity in ambiguous:
+ by_file.pop(identity, None)
+ return by_file
+
+
+def _parquet_col_range(metadata, col):
+ """Min/max of ``col`` across a parquet file's row groups; None when a row
group
+ lacks usable stats for ``col``."""
+ lo, hi = None, None
+ for i in range(metadata.num_row_groups):
+ rg = metadata.row_group(i)
+ stats = None
+ for j in range(rg.num_columns):
+ if rg.column(j).path_in_schema == col:
+ stats = rg.column(j).statistics
+ break
+ if stats is None or not stats.has_min_max:
+ return None
+ lo = stats.min if lo is None else min(lo, stats.min)
+ hi = stats.max if hi is None else max(hi, stats.max)
+ return None if lo is None else (lo, hi)
+
+
+def _footer_col_type(metadata, col):
+ """The arrow type ``col`` is stored as in this parquet file; None if
unavailable."""
+ try:
+ return metadata.schema.to_arrow_schema().field(col).type
+ except Exception:
+ return None
+
+
+def _split_key_range(
+ split, name_for_schema, field_id, key_type, file_io, manifest_stats):
+ """Split bounds from key/manifest stats, then Parquet footers."""
+ import pyarrow.parquet as pq
+ lo, hi = None, None
+ for f in split.files:
+ rng = (_file_stats_range(f, field_id, key_type)
+ or manifest_stats.get(_file_identity(f)))
+ if rng is not None:
+ lo = rng[0] if lo is None else min(lo, rng[0])
+ hi = rng[1] if hi is None else max(hi, rng[1])
+ continue
+ col = name_for_schema(f.schema_id)
+ if col is None:
+ return None, None
+ path = f.external_path if f.external_path else f.file_path
+ if path is None or not path.endswith(".parquet"):
+ return None, None
+ try:
+ stream = file_io.new_input_stream(path)
+ try:
+ metadata = pq.read_metadata(stream)
+ finally:
+ stream.close()
+ except Exception as e:
+ # Footer read can fail (e.g. local-cache streams aren't seekable):
degrade to
+ # unknown, but warn -- the fallback would otherwise be silent.
+ _LOG.warning("range_join: parquet footer read failed for %s (%s);
treating "
+ "its range as unknown", path, e)
+ return None, None
+ if _footer_col_type(metadata, col) != key_type:
+ return None, None
+ rng = _parquet_col_range(metadata, col)
+ if rng is None:
+ return None, None
+ lo = rng[0] if lo is None else min(lo, rng[0])
+ hi = rng[1] if hi is None else max(hi, rng[1])
+ if lo is None:
+ return None, None
+ return lo, hi
+
+
+def _range_stats_trusted(table, splits, range_col):
+ """Whether stored bounds still enclose the read value."""
+ from pypaimon.read.query_auth_split import QueryAuthSplit
+
+ masked = any(isinstance(s, QueryAuthSplit) and s.auth_result.column_masking
+ and range_col in s.auth_result.column_masking for s in splits)
+ # Merge engines may rewrite non-PK values.
+ return not masked and not (
+ table.primary_keys and range_col not in table.primary_keys)
+
+
+def _plan_ranged_splits(table_id, catalog_options, projection, range_col,
partitions=None):
+ """Plan ``(split, min, max)`` entries on the driver."""
+ import os
+ from concurrent.futures import ThreadPoolExecutor
+ from pypaimon.common.predicate_builder import PredicateBuilder
+ from pypaimon.schema.data_types import PyarrowFieldParser
+ table = get_table(table_id, catalog_options, None, "range_join")
+ schema_id = table.table_schema.id
+ snapshot_id = pin_latest_snapshot(table)
+ if snapshot_id is None:
+ return [], schema_id
+ file_io = table.file_io
+ key_type = PyarrowFieldParser.from_paimon_schema(
+ table.table_schema.fields).field(range_col).type
+
+ # Range key's physical name in a file's schema, by field id (rename/swap
safe).
+ # Cached; current-schema files skip the schema load.
+ key_field_id = next(f.id for f in table.table_schema.fields if f.name ==
range_col)
+ name_cache, cache_lock = {schema_id: range_col}, threading.Lock()
+
+ def name_for_schema(sid):
+ with cache_lock:
+ if sid in name_cache:
+ return name_cache[sid]
+ try:
+ fields = table.schema_manager.get_schema(sid).fields
+ name = next((f.name for f in fields if f.id == key_field_id), None)
+ except Exception:
+ name = None
+ with cache_lock:
+ name_cache[sid] = name
+ return name
+
+ rb = table.new_read_builder()
+ if partitions:
+ # Build the partition predicate before projection, so its field list
still has
+ # the partition columns. None means the null partition (is_null, not =
None).
+ pb = rb.new_predicate_builder()
+ rb = rb.with_partition_filter(PredicateBuilder.and_predicates(
+ [pb.is_null(c) if v is None else pb.equal(c, v)
+ for c, v in partitions.items()]))
+ if projection is not None:
+ rb = rb.with_projection(projection)
+ splits = list(rb.new_scan().plan().splits())
+ # Merges or masking can move values beyond stored bounds.
+ if not _range_stats_trusted(table, splits, range_col):
+ return [(s, None, None) for s in splits], schema_id
+ manifest_stats = _manifest_stats_by_file(
+ table, key_field_id, key_type, snapshot_id)
+ workers = min(16, (os.cpu_count() or 4) * 4, len(splits) or 1)
+ with ThreadPoolExecutor(max_workers=workers) as pool:
+ bounds = pool.map(
+ lambda s: _split_key_range(
+ s, name_for_schema, key_field_id, key_type, file_io,
manifest_stats),
+ splits)
+ return [(s, lo, hi) for s, (lo, hi) in zip(splits, bounds)], schema_id
+
+
+def _cut_points(ranged_sides, num_ranges):
+ """Pick ``num_ranges - 1`` cut values from row-count-weighted file
boundaries."""
+ points = []
+ for ranged in ranged_sides:
+ for split, lo, hi in ranged:
+ if lo is None:
+ continue
+ rows = sum(f.row_count for f in split.files)
+ points.append((lo, rows / 2.0))
+ points.append((hi, rows / 2.0))
+ if not points:
+ return []
+ points.sort(key=lambda p: p[0])
+ total = sum(w for _, w in points)
+ cuts, acc, k = [], 0.0, 1
+ for value, weight in points:
+ acc += weight
+ if k >= num_ranges:
+ break
+ if acc >= total * k / num_ranges:
+ if not cuts or value > cuts[-1]: # strictly increasing
+ cuts.append(value)
+ k += 1
+ return cuts
+
+
+def _split_rows(split):
+ return sum(f.row_count for f in split.files)
+
+
+def _split_bytes(split):
+ # Re-read cost is bytes, not rows: a wide-row split is cheap by rows,
costly by I/O.
+ return sum(f.file_size for f in split.files)
+
+
+def _total_reads(l_ranged, r_ranged, ranges):
+ """Bytes physically read = each split's bytes times the ranges it overlaps
(each range
+ reads the whole split and clips). Counts unknown-stats and wide known
splits alike."""
+ reads = 0
+ for ranged in (l_ranged, r_ranged):
+ for split, lo, hi in ranged:
+ spans = sum(1 for r_lo, r_hi in ranges if _overlaps(lo, hi, r_lo,
r_hi))
+ reads += _split_bytes(split) * spans
+ return reads
+
+
+def _bounded_ranges(l_ranged, r_ranged, num_ranges):
+ """Cut into ``num_ranges`` ranges, halving until total re-read <=
_REREAD_BUDGET full
+ scans of bytes, so poorly clustered input can't cost far more than one
scan."""
+ total_bytes = sum(_split_bytes(s)
+ for ranged in (l_ranged, r_ranged) for s, _, _ in ranged)
+ budget = _REREAD_BUDGET * max(1, total_bytes)
+ while True:
+ ranges = _ranges_from_cuts(_cut_points((l_ranged, r_ranged),
num_ranges))
+ if num_ranges <= 1 or _total_reads(l_ranged, r_ranged, ranges) <=
budget:
+ return ranges
+ num_ranges = max(1, num_ranges // 2)
+
+
+def _ranges_from_cuts(cuts):
+ # Half-open [lo, hi); None = unbounded end.
+ bounds = [None] + cuts + [None]
+ return [(bounds[i], bounds[i + 1]) for i in range(len(bounds) - 1)]
+
+
+def _overlaps(lo, hi, r_lo, r_hi):
+ if lo is None: # unknown split range: belongs to every range
+ return True
+ return (r_lo is None or hi >= r_lo) and (r_hi is None or lo < r_hi)
+
+
+def _restrict_to_range(arrow_table, col, lo, hi):
+ """Keep rows with ``lo <= col < hi``. Null keys are always dropped (an
inner
+ join never matches them), which also keeps the result independent of
num_ranges."""
+ import pyarrow.compute as pc
+ mask = pc.is_valid(arrow_table[col])
+ if lo is not None:
+ mask = pc.and_(mask, pc.greater_equal(arrow_table[col], lo))
+ if hi is not None:
+ mask = pc.and_(mask, pc.less(arrow_table[col], hi))
+ return arrow_table.filter(mask)
+
+
+def range_join(
+ left: str,
+ right: str,
+ catalog_options: Dict[str, str],
+ *,
+ on: Optional[OnSpec] = None,
+ left_on: Optional[OnSpec] = None,
+ right_on: Optional[OnSpec] = None,
+ num_ranges: Optional[int] = None,
+ left_projection: Optional[List[str]] = None,
+ right_projection: Optional[List[str]] = None,
+ left_partitions: Optional[Dict[str, Any]] = None,
+ right_partitions: Optional[Dict[str, Any]] = None,
+ join_type: str = "inner",
+ ray_remote_args: Optional[Dict[str, Any]] = None,
+) -> "ray.data.Dataset":
+ """Join two tables clustered by the first join key with no global shuffle.
+
+ ``on`` when both sides use the same column names, or
``left_on``/``right_on``
+ when they differ (positionally paired). The first pair is the range key
used to
+ cut the key space. ``left_partitions``/``right_partitions`` ({column:
value} dicts
+ on partition columns) prune each side to those partitions first. Sides
must not
+ share column names other than ``on`` keys. Returns a ``ray.data.Dataset``.
+ """
+ import ray
+
+ if not hasattr(ray.data, "from_arrow_refs"):
+ raise RuntimeError(
+ "range_join needs a Ray version with ray.data.from_arrow_refs; "
+ f"installed ray is {ray.__version__}.")
+
+ if (on is None) == (left_on is None and right_on is None):
+ raise ValueError("range_join requires exactly one of on= or
left_on=/right_on=.")
+ if on is not None:
+ lkeys = rkeys = norm_on(on)
+ else:
+ if left_on is None or right_on is None:
+ raise ValueError("range_join requires both left_on= and
right_on=.")
+ lkeys, rkeys = norm_on(left_on), norm_on(right_on)
+ if len(lkeys) != len(rkeys) or not lkeys:
+ raise ValueError(
+ f"range_join join keys must pair up non-empty; got
left_on={lkeys}, right_on={rkeys}.")
+ if join_type != "inner":
+ # Outer joins would need every unmatched row emitted exactly once
across
+ # ranges plus null-key handling; only inner is supported for now.
+ raise ValueError(f"range_join currently supports only
join_type='inner'; got {join_type!r}.")
+
+ ltable = get_table(left, catalog_options, None, "range_join")
+ rtable = get_table(right, catalog_options, None, "range_join")
+
+ # Partition filters must name partition columns, else they'd be silently
ignored.
+ for name, tbl, parts in (("left_partitions", ltable, left_partitions),
+ ("right_partitions", rtable, right_partitions)):
+ bad = sorted(set(parts) - set(tbl.partition_keys)) if parts else []
+ if bad:
+ raise ValueError(
+ f"range_join {name} keys {bad} are not partition columns; "
+ f"partition columns are {list(tbl.partition_keys)}.")
+
+ missing = [c for c in lkeys if c not in ltable.field_dict] \
+ + [c for c in rkeys if c not in rtable.field_dict]
+ if missing:
+ raise ValueError(f"range_join keys not found in table schema:
{missing}.")
+ type_mismatch = [
+ (lc, rc, key_type(ltable, lc), key_type(rtable, rc))
+ for lc, rc in zip(lkeys, rkeys)
+ if key_type(ltable, lc) != key_type(rtable, rc)
+ ]
+ if type_mismatch:
+ raise ValueError(
+ "range_join key columns must have the same type on both sides; "
+ f"mismatched (left, right, left type, right type):
{type_mismatch}.")
+
+ # Reject unsupported key types up front (not inside a worker as
ArrowInvalid). Every
+ # join key must be hashable; nested (ARRAY<>/MAP<>/ROW<>/...) and VARIANT
are not.
+ for c in lkeys:
+ t = key_type(ltable, c).upper()
+ if "<" in t or t.startswith("VARIANT"):
+ raise ValueError(
+ f"range_join join key {c!r} must not be a nested/complex type;
got {t}.")
+ # The range key (first pair) additionally must be range-partitionable.
+ range_key_type = key_type(ltable, lkeys[0]).upper()
+ reason = None
+ if range_key_type.startswith(("FLOAT", "DOUBLE")):
+ # NaN falls out of every range while the hash join still matches it ->
drops rows.
+ reason = "FLOAT/DOUBLE"
+ elif "LOCAL TIME ZONE" in range_key_type or "TIMESTAMP_LTZ" in
range_key_type:
+ # Footer stats decode to naive datetimes; a tz-aware column can't
compare to them.
+ reason = "TIMESTAMP WITH LOCAL TIME ZONE"
+ if reason:
+ raise ValueError(
+ f"range_join range key {lkeys[0]!r} must not be {reason}; "
+ "use an integer/string/date/timestamp key.")
+
+ # The join keys must survive projection, or the local join has no key.
+ if left_projection is not None and not set(lkeys) <= set(left_projection):
+ raise ValueError(
+ f"left_projection must include the join keys {lkeys}; got
{left_projection}.")
+ if right_projection is not None and not set(rkeys) <=
set(right_projection):
+ raise ValueError(
+ f"right_projection must include the join keys {rkeys}; got
{right_projection}.")
+ # pyarrow drops the right keys (coalesced into the left), so the output
keeps the LEFT
+ # key names. A right non-key column sharing a left column name collides ->
reject it.
+ 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(rkeys)))
+ if collisions:
+ raise ValueError(
+ f"range_join output columns collide: {collisions}. The output
keeps the left "
+ "key names and the right non-key columns; project or rename the
overlap away.")
+
+ l_range_col, r_range_col = lkeys[0], rkeys[0]
+ l_ranged, l_schema_id = _plan_ranged_splits(
+ left, catalog_options, left_projection, l_range_col, left_partitions)
+ r_ranged, r_schema_id = _plan_ranged_splits(
+ right, catalog_options, right_projection, r_range_col,
right_partitions)
+
+ def _empty():
+ empty = read_splits(
+ left, catalog_options, left_projection, [], l_schema_id,
"range_join").join(
+ read_splits(right, catalog_options, right_projection, [],
r_schema_id, "range_join"),
+ keys=lkeys, right_keys=rkeys, join_type=join_type)
+ return ray.data.from_arrow(empty)
+
+ if not l_ranged or not r_ranged: # inner join: one empty side, empty
result
+ return _empty()
+
+ if num_ranges is None:
+ num_ranges = max(len(l_ranged), len(r_ranged))
+ elif not isinstance(num_ranges, int) or num_ranges < 1:
+ raise ValueError(f"num_ranges must be an int >= 1; got
{num_ranges!r}.")
+ num_ranges = max(1, min(_MAX_RANGES, num_ranges)) # cap tasks even when
explicit
+ # Reduce ranges until total re-read stays bounded -- see _bounded_ranges.
+ ranges = _bounded_ranges(l_ranged, r_ranged, num_ranges)
+
+ def _join_range(left_splits, right_splits, lo, hi):
+ # No predicate pushdown: the range key may be schema-evolved (e.g. a
file stored
+ # as INT read as STRING), which the reader can't compare against a
new-type bound.
+ # The in-memory clip below does the exact, evolution-safe filtering.
+ lt = _restrict_to_range(
+ read_splits(left, catalog_options, left_projection, left_splits,
+ l_schema_id, "range_join"),
+ l_range_col, lo, hi)
+ rt = _restrict_to_range(
+ read_splits(right, catalog_options, right_projection, right_splits,
+ r_schema_id, "range_join"),
+ r_range_col, lo, hi)
+ return lt.join(rt, keys=lkeys, right_keys=rkeys, join_type=join_type)
+
+ # ``@ray.remote()`` (empty parens) is rejected by Ray, so wrap
conditionally.
+ remote_fn = ray.remote(**ray_remote_args)(_join_range) if ray_remote_args
else ray.remote(_join_range)
+ refs = []
+ for r_lo, r_hi in ranges:
+ ls = [s for s, lo, hi in l_ranged if _overlaps(lo, hi, r_lo, r_hi)]
+ rs = [s for s, lo, hi in r_ranged if _overlaps(lo, hi, r_lo, r_hi)]
+ if not ls or not rs: # inner join: a one-sided range can't match
+ continue
+ refs.append(remote_fn.remote(ls, rs, r_lo, r_hi))
+ if not refs:
+ return _empty()
+ # Keep each range's result as a distributed object ref -- never pulled
into the driver.
+ return ray.data.from_arrow_refs(refs)
diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py
b/paimon-python/pypaimon/tests/ray_range_join_test.py
new file mode 100644
index 0000000000..3696bf851f
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_range_join_test.py
@@ -0,0 +1,546 @@
+# 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()
