From 7e857ce2965187378acef5a573944c6e69bb1a9a Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 19 Jul 2026 07:05:22 -0700 Subject: [PATCH 01/12] [python] Add range_join: shuffle-free Ray join for range-clustered tables Cut the join-key space into ranges from per-file min/max stats in the manifest, then read and join each range in its own Ray task -- no global shuffle. Complements bucket_join for tables sorted/clustered by the join key rather than co-bucketed. Half-open ranges plus an in-memory range filter keep every matching pair produced exactly once regardless of range count; splits with no usable stats fall back to joining every range. Shared driver/worker helpers (table cache, snapshot pin, split read) are factored into join_common and reused by bucket_join. --- 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 | 291 ++++++++++++++++++ .../pypaimon/tests/ray_range_join_test.py | 242 +++++++++++++++ 5 files changed, 644 insertions(+), 58 deletions(-) create mode 100644 paimon-python/pypaimon/ray/join_common.py create mode 100644 paimon-python/pypaimon/ray/range_join.py create mode 100644 paimon-python/pypaimon/tests/ray_range_join_test.py diff --git a/paimon-python/pypaimon/ray/__init__.py b/paimon-python/pypaimon/ray/__init__.py index bc91c8da4570..03de06824967 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 @@ "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 818643ec9549..2b377126dfd4 100644 --- a/paimon-python/pypaimon/ray/bucket_join.py +++ b/paimon-python/pypaimon/ray/bucket_join.py @@ -21,22 +21,21 @@ 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 000000000000..0892ea9c080c --- /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 000000000000..5f2805fac63b --- /dev/null +++ b/paimon-python/pypaimon/ray/range_join.py @@ -0,0 +1,291 @@ +# 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. + +"""Range-aligned join on Ray for two Paimon tables sorted/clustered by the join key. + +The driver cuts the key space into ranges from per-file min/max stats in the manifest; +each range is read and joined in its own Ray task, with no global shuffle. Works best +when both sides are clustered by the first join key. + +Correctness never depends on stats: a split whose min/max is missing joins every range. +That is also the cost model -- such a split is read once per range, so a mix of a few +no-stats splits with N stats-derived ranges reads those few up to N times. (When *no* +split has stats there are no cut points, so it collapses to a single range and one read.) +""" + +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"] + +_MAX_RANGES = 512 + + +def _stats_field_names(table_field_names, file): + # Names the file's value_stats row is laid out in; None = unknown. + if file.value_stats_cols is not None: + return file.value_stats_cols + if file.write_cols is not None: + return file.write_cols + return table_field_names + + +def _split_key_range(split, col, table_field_names, table_schema_id): + """Min/max of ``col`` over the split's files from manifest stats; (None, None) + when any file lacks usable stats (the split then joins every range).""" + lo, hi = None, None + for f in split.files: + if f.value_stats_cols is None and f.write_cols is None \ + and f.schema_id != table_schema_id: + # Stats laid out in an older schema's field order: name lookup unsafe. + return None, None + names = _stats_field_names(table_field_names, f) + stats = f.value_stats + if col not in names: + return None, None + idx = names.index(col) + if len(stats.min_values) <= idx: + return None, None + fmin = stats.min_values.get_field(idx) + fmax = stats.max_values.get_field(idx) + if fmin is None or fmax is None: # all-null or absent stats + return None, None + lo = fmin if lo is None else min(lo, fmin) + hi = fmax if hi is None else max(hi, fmax) + if lo is None: + return None, None + return lo, hi + + +def _plan_ranged_splits(table_id, catalog_options, projection, range_col): + """Plan the manifest driver-side; returns ``(ranged_splits, schema_id)`` where + ranged_splits is a list of (split, lo, hi).""" + table = get_table(table_id, catalog_options, None, "range_join") + schema_id = table.table_schema.id + 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() + field_names = table.field_names + ranged = [] + for s in scan.plan().splits(): + lo, hi = _split_key_range(s, range_col, field_names, schema_id) + ranged.append((s, lo, hi)) + return ranged, 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 _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_predicate(table_id, catalog_options, projection, schema_id, col, lo, hi): + # Pushdown superset [lo, hi] for row-group pruning; exact filtering happens in-memory. + from pypaimon.common.predicate_builder import PredicateBuilder + from pypaimon.ray.join_common import read_builder + if lo is None and hi is None: + return None + pb = read_builder( + table_id, catalog_options, projection, schema_id, "range_join").new_predicate_builder() + preds = [] + if lo is not None: + preds.append(pb.greater_or_equal(col, lo)) + if hi is not None: + preds.append(pb.less_or_equal(col, hi)) + return PredicateBuilder.and_predicates(preds) + + +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, + 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. 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") + + 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}.") + + # 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}.") + # Non-key columns must not collide, or pyarrow's join collides on them. + 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(lkeys)) & (set(rcols) - set(rkeys))) + if collisions: + raise ValueError( + f"range_join sides must not share columns other than the join keys; " + f"both have {collisions}. Project or rename them 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) + r_ranged, r_schema_id = _plan_ranged_splits( + right, catalog_options, right_projection, r_range_col) + + 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(1, min(_MAX_RANGES, max(len(l_ranged), len(r_ranged)))) + elif num_ranges < 1: + raise ValueError(f"num_ranges must be >= 1; got {num_ranges}.") + ranges = _ranges_from_cuts(_cut_points((l_ranged, r_ranged), num_ranges)) + + def _join_range(left_splits, right_splits, lo, hi): + lt = _restrict_to_range( + read_splits(left, catalog_options, left_projection, left_splits, + l_schema_id, "range_join", + _range_predicate(left, catalog_options, left_projection, + l_schema_id, l_range_col, lo, hi)), + l_range_col, lo, hi) + rt = _restrict_to_range( + read_splits(right, catalog_options, right_projection, right_splits, + r_schema_id, "range_join", + _range_predicate(right, catalog_options, right_projection, + r_schema_id, r_range_col, lo, hi)), + 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 000000000000..61d96fc57e7f --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -0,0 +1,242 @@ +# 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 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) + got = {r["rid"]: 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, "must not share columns"): + 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_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_split_key_range_reads_stats(self): + # The planner reads a file's min/max for the range column from manifest stats. + 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_no_stats_fallback_matches_global_join(self): + # metadata.stats-mode=none -> no per-file min/max -> every split joins every range. + 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()) + # Exactly once, correct, despite the fallback reading every split in every range. + 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) + + +if __name__ == "__main__": + unittest.main() From ee108dd044ec4fc750265e6282d1bd26978c5d7d Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 19 Jul 2026 07:54:04 -0700 Subject: [PATCH 02/12] [python] range_join: read stats via manifest, guard NaN/tz keys, fix join output - Blocker: scan.plan() drops value stats, so every split looked unknown and the join always collapsed to one task. Read per-file min/max straight from the manifest (drop_stats=False), keyed by file name; stats never reach the workers. - Reject FLOAT/DOUBLE range keys (NaN falls out of every range while the hash join still matches NaN==NaN -> dropped matches) and TIMESTAMP_LTZ (naive manifest stats vs tz-aware column can't be compared). - Fix the collision check to catch a left key colliding with a right non-key column; document that the output keeps the left key names. - Cap the range count by an unknown-stats read budget so the every-range fallback re-reads at most one extra full scan. --- paimon-python/pypaimon/ray/range_join.py | 127 +++++++++++++----- .../pypaimon/tests/ray_range_join_test.py | 36 ++++- 2 files changed, 127 insertions(+), 36 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 5f2805fac63b..07be10196e7c 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -22,9 +22,8 @@ when both sides are clustered by the first join key. Correctness never depends on stats: a split whose min/max is missing joins every range. -That is also the cost model -- such a split is read once per range, so a mix of a few -no-stats splits with N stats-derived ranges reads those few up to N times. (When *no* -split has stats there are no cut points, so it collapses to a single range and one read.) +Such a split is read once per range, so the range count is capped to keep those re-reads +under one extra full scan (all-no-stats collapses to a single range). """ from typing import Any, Dict, List, Optional @@ -43,37 +42,61 @@ _MAX_RANGES = 512 -def _stats_field_names(table_field_names, file): - # Names the file's value_stats row is laid out in; None = unknown. +def _file_key_range(file, col, table_field_names, table_schema_id): + """Min/max of ``col`` for one data file from its manifest stats; None when the + file has no usable stats for ``col``.""" + # value_stats row layout: value_stats_cols, else write_cols, else the full schema. if file.value_stats_cols is not None: - return file.value_stats_cols - if file.write_cols is not None: - return file.write_cols - return table_field_names + names = file.value_stats_cols + elif file.write_cols is not None: + names = file.write_cols + elif file.schema_id == table_schema_id: + names = table_field_names + else: + return None # stats laid out in an older schema's field order: unsafe + if col not in names: + return None + idx = names.index(col) + stats = file.value_stats + if len(stats.min_values) <= idx: + return None + fmin = stats.min_values.get_field(idx) + fmax = stats.max_values.get_field(idx) + if fmin is None or fmax is None: # all-null or absent stats + return None + return fmin, fmax + + +def _range_stats_by_file(table, col): + """{file_name: (min, max)} for ``col``, read straight from the manifest with stats. + scan.plan() drops value stats, so read the entries separately (driver-side).""" + from pypaimon.manifest.manifest_file_manager import ManifestFileManager + from pypaimon.manifest.manifest_list_manager import ManifestListManager + snapshot = table.snapshot_manager().get_latest_snapshot() + if snapshot is None: + return {} + manifest_files = ManifestListManager(table).read_all(snapshot) + entries = ManifestFileManager(table).read_entries_parallel(manifest_files, drop_stats=False) + field_names = table.field_names + schema_id = table.table_schema.id + stats = {} + for entry in entries: + rng = _file_key_range(entry.file, col, field_names, schema_id) + if rng is not None: + stats[entry.file.file_name] = rng + return stats -def _split_key_range(split, col, table_field_names, table_schema_id): - """Min/max of ``col`` over the split's files from manifest stats; (None, None) - when any file lacks usable stats (the split then joins every range).""" +def _split_key_range(split, stats_by_file): + """Min/max of the range key over the split's files, or (None, None) when any file + lacks stats (the split then joins every range).""" lo, hi = None, None for f in split.files: - if f.value_stats_cols is None and f.write_cols is None \ - and f.schema_id != table_schema_id: - # Stats laid out in an older schema's field order: name lookup unsafe. - return None, None - names = _stats_field_names(table_field_names, f) - stats = f.value_stats - if col not in names: - return None, None - idx = names.index(col) - if len(stats.min_values) <= idx: - return None, None - fmin = stats.min_values.get_field(idx) - fmax = stats.max_values.get_field(idx) - if fmin is None or fmax is None: # all-null or absent stats + rng = stats_by_file.get(f.file_name) + if rng is None: return None, None - lo = fmin if lo is None else min(lo, fmin) - hi = fmax if hi is None else max(hi, fmax) + 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 @@ -81,17 +104,18 @@ def _split_key_range(split, col, table_field_names, table_schema_id): def _plan_ranged_splits(table_id, catalog_options, projection, range_col): """Plan the manifest driver-side; returns ``(ranged_splits, schema_id)`` where - ranged_splits is a list of (split, lo, hi).""" + ranged_splits is a list of (split, lo, hi). Stats come from a separate manifest + read (scan.plan() strips them) and never reach the workers.""" table = get_table(table_id, catalog_options, None, "range_join") schema_id = table.table_schema.id if pin_latest_snapshot(table) is None: return [], schema_id + stats_by_file = _range_stats_by_file(table, range_col) rb = table.new_read_builder() scan = (rb.with_projection(projection) if projection is not None else rb).new_scan() - field_names = table.field_names ranged = [] for s in scan.plan().splits(): - lo, hi = _split_key_range(s, range_col, field_names, schema_id) + lo, hi = _split_key_range(s, stats_by_file) ranged.append((s, lo, hi)) return ranged, schema_id @@ -122,6 +146,20 @@ def _cut_points(ranged_sides, num_ranges): return cuts +def _range_budget(l_ranged, r_ranged): + # Cap ranges so unknown-stats splits (read in every range) re-read <= one full scan. + total = unknown = 0 + for ranged in (l_ranged, r_ranged): + for split, lo, _ in ranged: + rows = sum(f.row_count for f in split.files) + total += rows + if lo is None: + unknown += rows + if unknown <= 0: + return _MAX_RANGES + return max(1, total // unknown) + + def _ranges_from_cuts(cuts): # Half-open [lo, hi); None = unbounded end. bounds = [None] + cuts + [None] @@ -223,6 +261,21 @@ def range_join( "range_join key columns must have the same type on both sides; " f"mismatched (left, right, left type, right type): {type_mismatch}.") + # The range key is the first pair; reject types that can't be range-partitioned safely. + range_key_type = key_type(ltable, lkeys[0]).upper() + if range_key_type.startswith("FLOAT") or range_key_type.startswith("DOUBLE"): + # NaN compares false to every bound, so it would fall out of every range while + # pyarrow's hash join still matches NaN == NaN -> silently dropped matches. + raise ValueError( + f"range_join range key {lkeys[0]!r} must not be FLOAT/DOUBLE (NaN can't be " + "range-partitioned); use an integer/string/date key.") + if "LOCAL TIME ZONE" in range_key_type: + # Manifest stats decode to naive datetimes; a tz-aware Arrow column can't be + # compared against them. Not supported yet. + raise ValueError( + f"range_join range key {lkeys[0]!r} of type TIMESTAMP WITH LOCAL TIME ZONE " + "is not supported yet.") + # 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( @@ -230,14 +283,15 @@ def range_join( 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}.") - # Non-key columns must not collide, or pyarrow's join collides on them. + # 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(lkeys)) & (set(rcols) - set(rkeys))) + collisions = sorted(set(lcols) & (set(rcols) - set(rkeys))) if collisions: raise ValueError( - f"range_join sides must not share columns other than the join keys; " - f"both have {collisions}. Project or rename them away.") + 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( @@ -259,6 +313,9 @@ def _empty(): num_ranges = max(1, min(_MAX_RANGES, max(len(l_ranged), len(r_ranged)))) elif num_ranges < 1: raise ValueError(f"num_ranges must be >= 1; got {num_ranges}.") + # An unknown-stats split is read in every range. Cap ranges so those re-reads add at + # most one extra full scan, else the fallback can cost more than a shuffle. + num_ranges = min(num_ranges, _range_budget(l_ranged, r_ranged)) ranges = _ranges_from_cuts(_cut_points((l_ranged, r_ranged), num_ranges)) def _join_range(left_splits, right_splits, lo, hi): diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index 61d96fc57e7f..50c0f729bd91 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import collections import os import shutil import tempfile @@ -118,7 +119,8 @@ def test_left_on_right_on_different_names(self): ds = range_join( "default.rj_lr_left", "default.rj_lr_right", self.catalog_options, left_on="lid", right_on="rid", num_ranges=3) - got = {r["rid"]: r["val"] for r in ds.take_all()} + # 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): @@ -187,6 +189,38 @@ 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_range_budget_caps_ranges_when_stats_missing(self): + Split = collections.namedtuple("Split", "files") + File = collections.namedtuple("File", "row_count") + known = [(Split([File(100)]), 0, 99)] # 100 rows with stats + unknown = [(Split([File(100)]), None, None)] # 100 rows without stats + # unknown == total/2 -> budget 2; all-known -> no cap; all-unknown -> 1. + self.assertEqual(rjmod._range_budget(known, unknown), 2) + self.assertEqual(rjmod._range_budget(known, known), rjmod._MAX_RANGES) + self.assertEqual(rjmod._range_budget(unknown, unknown), 1) + def test_split_key_range_reads_stats(self): # The planner reads a file's min/max for the range column from manifest stats. loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) From 7a6cc3776416ad47548b631ec09dcbfe4d344a4b Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 20 Jul 2026 20:51:27 -0700 Subject: [PATCH 03/12] [python] Fix range_join stats extraction to match files_table The value_stats column names come from value_stats_cols when set, else the full schema -- the write_cols/schema_id detour resolved wrong names, so every split looked unknown and the join collapsed to one range. Read min/max via the same path files_table uses. Update the collision-message assertion. --- paimon-python/pypaimon/ray/range_join.py | 24 +++++++------------ .../pypaimon/tests/ray_range_join_test.py | 2 +- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 07be10196e7c..2718e0df2924 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -42,26 +42,21 @@ _MAX_RANGES = 512 -def _file_key_range(file, col, table_field_names, table_schema_id): +def _file_key_range(file, col, table_field_names): """Min/max of ``col`` for one data file from its manifest stats; None when the file has no usable stats for ``col``.""" - # value_stats row layout: value_stats_cols, else write_cols, else the full schema. - if file.value_stats_cols is not None: - names = file.value_stats_cols - elif file.write_cols is not None: - names = file.write_cols - elif file.schema_id == table_schema_id: - names = table_field_names - else: - return None # stats laid out in an older schema's field order: unsafe + # value_stats layout: value_stats_cols when set, else the full schema (see files_table). + cols = file.value_stats_cols + names = list(cols) if cols else table_field_names if col not in names: return None idx = names.index(col) stats = file.value_stats - if len(stats.min_values) <= idx: + min_values = getattr(stats.min_values, "values", None) or [] + max_values = getattr(stats.max_values, "values", None) or [] + if idx >= len(min_values) or idx >= len(max_values): return None - fmin = stats.min_values.get_field(idx) - fmax = stats.max_values.get_field(idx) + fmin, fmax = min_values[idx], max_values[idx] if fmin is None or fmax is None: # all-null or absent stats return None return fmin, fmax @@ -78,10 +73,9 @@ def _range_stats_by_file(table, col): manifest_files = ManifestListManager(table).read_all(snapshot) entries = ManifestFileManager(table).read_entries_parallel(manifest_files, drop_stats=False) field_names = table.field_names - schema_id = table.table_schema.id stats = {} for entry in entries: - rng = _file_key_range(entry.file, col, field_names, schema_id) + rng = _file_key_range(entry.file, col, field_names) if rng is not None: stats[entry.file.file_name] = rng return stats diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index 50c0f729bd91..15e0d7af02d6 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -172,7 +172,7 @@ def test_rejects_shared_non_key_column(self): 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, "must not share columns"): + 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): From b79c16d97ef69001bfc612ee9eb867517661e1a3 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 20 Jul 2026 23:34:47 -0700 Subject: [PATCH 04/12] [python] Read range_join stats from BinaryRow via get_field, not .values value_stats.min/max_values are self-describing BinaryRows: they expose get_field(idx) and their own fields, not a .values list. The previous code read a non-existent .values attribute, so every split looked unknown and the join collapsed to one range. Index by the BinaryRow's own field names. --- paimon-python/pypaimon/ray/range_join.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 2718e0df2924..33881aaf64a3 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -42,21 +42,21 @@ _MAX_RANGES = 512 -def _file_key_range(file, col, table_field_names): +def _file_key_range(file, col): """Min/max of ``col`` for one data file from its manifest stats; None when the file has no usable stats for ``col``.""" - # value_stats layout: value_stats_cols when set, else the full schema (see files_table). - cols = file.value_stats_cols - names = list(cols) if cols else table_field_names + # min_values/max_values are self-describing BinaryRows: index via their own fields + # (get_field), not a re-derived column list. + min_row = file.value_stats.min_values + max_row = file.value_stats.max_values + names = [f.name for f in getattr(min_row, "fields", [])] if col not in names: return None idx = names.index(col) - stats = file.value_stats - min_values = getattr(stats.min_values, "values", None) or [] - max_values = getattr(stats.max_values, "values", None) or [] - if idx >= len(min_values) or idx >= len(max_values): + if idx >= len(min_row) or idx >= len(max_row): return None - fmin, fmax = min_values[idx], max_values[idx] + fmin = min_row.get_field(idx) + fmax = max_row.get_field(idx) if fmin is None or fmax is None: # all-null or absent stats return None return fmin, fmax @@ -72,10 +72,9 @@ def _range_stats_by_file(table, col): return {} manifest_files = ManifestListManager(table).read_all(snapshot) entries = ManifestFileManager(table).read_entries_parallel(manifest_files, drop_stats=False) - field_names = table.field_names stats = {} for entry in entries: - rng = _file_key_range(entry.file, col, field_names) + rng = _file_key_range(entry.file, col) if rng is not None: stats[entry.file.file_name] = rng return stats From a8c7719bd145c88ecba6b0fd2b09e82ea5f13547 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 20 Jul 2026 23:51:39 -0700 Subject: [PATCH 05/12] [python] Read range_join ranges from the parquet footer, not the manifest pypaimon-written tables leave the manifest value stats empty (min/max live only in the parquet footer), so the manifest approach always saw unknown ranges and collapsed to one task. Read each split file's footer min/max via FileIO on the driver; the splits sent to workers still carry no stats. Verified locally on Python 3.11 (pyarrow+ray): ray_range_join_test 14 passed, ray_bucket_join_test 18 passed. --- paimon-python/pypaimon/ray/range_join.py | 93 +++++++++++------------- 1 file changed, 43 insertions(+), 50 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 33881aaf64a3..1facc7fe3c7f 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -17,9 +17,10 @@ """Range-aligned join on Ray for two Paimon tables sorted/clustered by the join key. -The driver cuts the key space into ranges from per-file min/max stats in the manifest; -each range is read and joined in its own Ray task, with no global shuffle. Works best -when both sides are clustered by the first join key. +The driver cuts the key space into ranges from each file's parquet-footer min/max (the +manifest value stats are empty for pypaimon-written tables); each range is read and +joined in its own Ray task, with no global shuffle. Works best when both sides are +clustered by the first join key. Correctness never depends on stats: a split whose min/max is missing joins every range. Such a split is read once per range, so the range count is capped to keep those re-reads @@ -42,50 +43,42 @@ _MAX_RANGES = 512 -def _file_key_range(file, col): - """Min/max of ``col`` for one data file from its manifest stats; None when the - file has no usable stats for ``col``.""" - # min_values/max_values are self-describing BinaryRows: index via their own fields - # (get_field), not a re-derived column list. - min_row = file.value_stats.min_values - max_row = file.value_stats.max_values - names = [f.name for f in getattr(min_row, "fields", [])] - if col not in names: - return None - idx = names.index(col) - if idx >= len(min_row) or idx >= len(max_row): - return None - fmin = min_row.get_field(idx) - fmax = max_row.get_field(idx) - if fmin is None or fmax is None: # all-null or absent stats - return None - return fmin, fmax - - -def _range_stats_by_file(table, col): - """{file_name: (min, max)} for ``col``, read straight from the manifest with stats. - scan.plan() drops value stats, so read the entries separately (driver-side).""" - from pypaimon.manifest.manifest_file_manager import ManifestFileManager - from pypaimon.manifest.manifest_list_manager import ManifestListManager - snapshot = table.snapshot_manager().get_latest_snapshot() - if snapshot is None: - return {} - manifest_files = ManifestListManager(table).read_all(snapshot) - entries = ManifestFileManager(table).read_entries_parallel(manifest_files, drop_stats=False) - stats = {} - for entry in entries: - rng = _file_key_range(entry.file, col) - if rng is not None: - stats[entry.file.file_name] = rng - return stats - - -def _split_key_range(split, stats_by_file): - """Min/max of the range key over the split's files, or (None, None) when any file - lacks stats (the split then joins every range).""" +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 _split_key_range(split, col, file_io): + """Min/max of ``col`` over the split's files, read from their parquet footers; + (None, None) when any file isn't parquet or lacks usable stats (the split then + joins every range). Manifest value stats are empty for pypaimon-written tables, + so the footer is the source of truth.""" + import pyarrow.parquet as pq lo, hi = None, None for f in split.files: - rng = stats_by_file.get(f.file_name) + path = f.external_path if f.external_path else f.file_path + if path is None or not path.endswith(".parquet"): + return None, None + stream = file_io.new_input_stream(path) + try: + rng = _parquet_col_range(pq.read_metadata(stream), col) + except Exception: + return None, None + finally: + stream.close() if rng is None: return None, None lo = rng[0] if lo is None else min(lo, rng[0]) @@ -96,19 +89,19 @@ def _split_key_range(split, stats_by_file): def _plan_ranged_splits(table_id, catalog_options, projection, range_col): - """Plan the manifest driver-side; returns ``(ranged_splits, schema_id)`` where - ranged_splits is a list of (split, lo, hi). Stats come from a separate manifest - read (scan.plan() strips them) and never reach the workers.""" + """Plan driver-side; returns ``(ranged_splits, schema_id)`` where ranged_splits is + a list of (split, lo, hi). Ranges come from each file's parquet footer; the splits + sent to workers carry no stats.""" table = get_table(table_id, catalog_options, None, "range_join") schema_id = table.table_schema.id if pin_latest_snapshot(table) is None: return [], schema_id - stats_by_file = _range_stats_by_file(table, range_col) + file_io = table.file_io rb = table.new_read_builder() scan = (rb.with_projection(projection) if projection is not None else rb).new_scan() ranged = [] for s in scan.plan().splits(): - lo, hi = _split_key_range(s, stats_by_file) + lo, hi = _split_key_range(s, range_col, file_io) ranged.append((s, lo, hi)) return ranged, schema_id From 08901315d4c8a8a1251cedaa5367a9b5647d610a Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 22 Jul 2026 03:52:24 -0700 Subject: [PATCH 06/12] [python] Coerce range_join footer bounds to the join-key type Under DATE->TIMESTAMP schema evolution, old files' footer bounds are dates and new files' are datetimes; comparing/sorting them raw raised TypeError. Coerce every footer min/max to the current join-key arrow type before use; a bound that can't be cast makes the split unknown (joins every range). Verified locally on Python 3.11: real DATE->TIMESTAMP evolution join now returns the correct result; ray_range_join_test 15 passed. --- paimon-python/pypaimon/ray/range_join.py | 38 ++++++++++--- .../pypaimon/tests/ray_range_join_test.py | 53 +++++++++++++++++-- 2 files changed, 80 insertions(+), 11 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 1facc7fe3c7f..d36509627d40 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -61,11 +61,26 @@ def _parquet_col_range(metadata, col): return None if lo is None else (lo, hi) -def _split_key_range(split, col, file_io): - """Min/max of ``col`` over the split's files, read from their parquet footers; - (None, None) when any file isn't parquet or lacks usable stats (the split then - joins every range). Manifest value stats are empty for pypaimon-written tables, - so the footer is the source of truth.""" +_UNCOERCIBLE = object() + + +def _coerce_bound(value, key_type): + """Coerce a footer bound to the current join-key type so bounds from files written + under different schemas (e.g. DATE then TIMESTAMP) stay mutually comparable. Returns + ``_UNCOERCIBLE`` when the value can't be cast (the split then joins every range).""" + import pyarrow as pa + try: + return pa.array([value]).cast(key_type)[0].as_py() + except Exception: + return _UNCOERCIBLE + + +def _split_key_range(split, col, key_type, file_io): + """Min/max of ``col`` over the split's files, read from their parquet footers and + coerced to ``key_type``; (None, None) when any file isn't parquet, lacks usable + stats, or has a bound that can't be coerced (the split then joins every range). + Manifest value stats are empty for pypaimon-written tables, so the footer is the + source of truth.""" import pyarrow.parquet as pq lo, hi = None, None for f in split.files: @@ -81,8 +96,12 @@ def _split_key_range(split, col, file_io): stream.close() 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]) + clo = _coerce_bound(rng[0], key_type) + chi = _coerce_bound(rng[1], key_type) + if clo is _UNCOERCIBLE or chi is _UNCOERCIBLE: + return None, None + lo = clo if lo is None else min(lo, clo) + hi = chi if hi is None else max(hi, chi) if lo is None: return None, None return lo, hi @@ -92,16 +111,19 @@ def _plan_ranged_splits(table_id, catalog_options, projection, range_col): """Plan driver-side; returns ``(ranged_splits, schema_id)`` where ranged_splits is a list of (split, lo, hi). Ranges come from each file's parquet footer; the splits sent to workers carry no stats.""" + from pypaimon.schema.data_types import PyarrowFieldParser table = get_table(table_id, catalog_options, None, "range_join") schema_id = table.table_schema.id if pin_latest_snapshot(table) is None: return [], schema_id file_io = table.file_io + key_type = PyarrowFieldParser.from_paimon_schema( + table.table_schema.fields).field(range_col).type rb = table.new_read_builder() scan = (rb.with_projection(projection) if projection is not None else rb).new_scan() ranged = [] for s in scan.plan().splits(): - lo, hi = _split_key_range(s, range_col, file_io) + lo, hi = _split_key_range(s, range_col, key_type, file_io) ranged.append((s, lo, hi)) return ranged, schema_id diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index 15e0d7af02d6..12f97ce67a56 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -16,6 +16,7 @@ # under the License. import collections +import datetime import os import shutil import tempfile @@ -211,6 +212,51 @@ def test_rejects_left_key_vs_right_column_collision(self): 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_range_budget_caps_ranges_when_stats_missing(self): Split = collections.namedtuple("Split", "files") File = collections.namedtuple("File", "row_count") @@ -234,8 +280,10 @@ def test_split_key_range_reads_stats(self): self.assertEqual(min(los), 10) self.assertEqual(max(his), 20) - def test_no_stats_fallback_matches_global_join(self): - # metadata.stats-mode=none -> no per-file min/max -> every split joins every range. + 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())]) @@ -252,7 +300,6 @@ def test_no_stats_fallback_matches_global_join(self): "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()) - # Exactly once, correct, despite the fallback reading every split in every range. self.assertEqual(got, [(i, i) for i in range(50, 150)]) def test_null_keys_dropped_independent_of_num_ranges(self): From 5f6dd75bed31979aa5c46fcabfbaf9de7e813244 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 22 Jul 2026 04:41:51 -0700 Subject: [PATCH 07/12] [python] range_join: only trust footer stats of the current key type Coercing an old file's footer bounds to the current key type is unsafe when the cast doesn't preserve order: INT->STRING makes 5/100 into '5'/'100', but '42' > '100' as a string, so a coerced range silently drops the '42' row. Use a file's footer stats only when it stored the key in the current type; any other type (schema evolution) marks the split unknown (joins every range). Also drop the range predicate pushdown -- the reader can't compare a new-type bound against an evolved file's physical column; the in-memory clip does the exact, evolution-safe filtering. Verified locally on Python 3.11: INT->STRING and DATE->TIMESTAMP evolutions now return the full correct result; ray_range_join_test 16 passed. --- paimon-python/pypaimon/ray/range_join.py | 62 ++++++------------- .../pypaimon/tests/ray_range_join_test.py | 37 +++++++++++ 2 files changed, 57 insertions(+), 42 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index d36509627d40..70a2724d1ae0 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -61,24 +61,20 @@ def _parquet_col_range(metadata, col): return None if lo is None else (lo, hi) -_UNCOERCIBLE = object() - - -def _coerce_bound(value, key_type): - """Coerce a footer bound to the current join-key type so bounds from files written - under different schemas (e.g. DATE then TIMESTAMP) stay mutually comparable. Returns - ``_UNCOERCIBLE`` when the value can't be cast (the split then joins every range).""" - import pyarrow as pa +def _footer_col_type(metadata, col): + """The arrow type ``col`` is stored as in this parquet file; None if unavailable.""" try: - return pa.array([value]).cast(key_type)[0].as_py() + return metadata.schema.to_arrow_schema().field(col).type except Exception: - return _UNCOERCIBLE + return None def _split_key_range(split, col, key_type, file_io): - """Min/max of ``col`` over the split's files, read from their parquet footers and - coerced to ``key_type``; (None, None) when any file isn't parquet, lacks usable - stats, or has a bound that can't be coerced (the split then joins every range). + """Min/max of ``col`` over the split's files, read from their parquet footers; + (None, None) when any file isn't parquet, lacks usable stats, or stored the key in a + different type than the current one (the split then joins every range). A different + type means schema evolution, and its footer bounds order differently (e.g. INT '10' + < '2' as a string), so they can't be trusted -- only same-type stats are used. Manifest value stats are empty for pypaimon-written tables, so the footer is the source of truth.""" import pyarrow.parquet as pq @@ -89,19 +85,18 @@ def _split_key_range(split, col, key_type, file_io): return None, None stream = file_io.new_input_stream(path) try: - rng = _parquet_col_range(pq.read_metadata(stream), col) + metadata = pq.read_metadata(stream) except Exception: return None, None finally: stream.close() - if rng is None: + if _footer_col_type(metadata, col) != key_type: return None, None - clo = _coerce_bound(rng[0], key_type) - chi = _coerce_bound(rng[1], key_type) - if clo is _UNCOERCIBLE or chi is _UNCOERCIBLE: + rng = _parquet_col_range(metadata, col) + if rng is None: return None, None - lo = clo if lo is None else min(lo, clo) - hi = chi if hi is None else max(hi, chi) + 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 @@ -192,22 +187,6 @@ def _restrict_to_range(arrow_table, col, lo, hi): return arrow_table.filter(mask) -def _range_predicate(table_id, catalog_options, projection, schema_id, col, lo, hi): - # Pushdown superset [lo, hi] for row-group pruning; exact filtering happens in-memory. - from pypaimon.common.predicate_builder import PredicateBuilder - from pypaimon.ray.join_common import read_builder - if lo is None and hi is None: - return None - pb = read_builder( - table_id, catalog_options, projection, schema_id, "range_join").new_predicate_builder() - preds = [] - if lo is not None: - preds.append(pb.greater_or_equal(col, lo)) - if hi is not None: - preds.append(pb.less_or_equal(col, hi)) - return PredicateBuilder.and_predicates(preds) - - def range_join( left: str, right: str, @@ -327,17 +306,16 @@ def _empty(): ranges = _ranges_from_cuts(_cut_points((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", - _range_predicate(left, catalog_options, left_projection, - l_schema_id, l_range_col, lo, hi)), + 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", - _range_predicate(right, catalog_options, right_projection, - r_schema_id, r_range_col, lo, hi)), + r_schema_id, "range_join"), r_range_col, lo, hi) return lt.join(rt, keys=lkeys, right_keys=rkeys, join_type=join_type) diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index 12f97ce67a56..de41eb3cf900 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -257,6 +257,43 @@ def test_date_to_timestamp_schema_evolution(self): 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_range_budget_caps_ranges_when_stats_missing(self): Split = collections.namedtuple("Split", "files") File = collections.namedtuple("File", "row_count") From 2b3ccb4469719556f5b50b55b01f4cd35cb18d62 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 22 Jul 2026 05:37:45 -0700 Subject: [PATCH 08/12] [python] range_join: partition pruning + parallel footer reads - left_partitions/right_partitions ({column: value} on partition columns) prune each side to those partitions before planning, so a join can target specific partitions instead of whole tables. - Read each split's parquet footer in a thread pool instead of serially, cutting driver-side planning latency on many-file tables. Verified locally on Python 3.11: partition-filtered and partitioned-table joins, many-to-many correctness; ray_range_join_test 18 passed. --- paimon-python/pypaimon/ray/range_join.py | 41 +++++++++++------ .../pypaimon/tests/ray_range_join_test.py | 44 +++++++++++++++++++ 2 files changed, 72 insertions(+), 13 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 70a2724d1ae0..24908b7126ab 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -102,10 +102,14 @@ def _split_key_range(split, col, key_type, file_io): return lo, hi -def _plan_ranged_splits(table_id, catalog_options, projection, range_col): +def _plan_ranged_splits(table_id, catalog_options, projection, range_col, partitions=None): """Plan driver-side; returns ``(ranged_splits, schema_id)`` where ranged_splits is - a list of (split, lo, hi). Ranges come from each file's parquet footer; the splits - sent to workers carry no stats.""" + a list of (split, lo, hi). Ranges come from each file's parquet footer (read in + parallel); the splits sent to workers carry no stats. ``partitions`` (a {column: + value} dict on partition columns) prunes to those partitions before planning.""" + 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 @@ -115,12 +119,20 @@ def _plan_ranged_splits(table_id, catalog_options, projection, range_col): key_type = PyarrowFieldParser.from_paimon_schema( table.table_schema.fields).field(range_col).type rb = table.new_read_builder() - scan = (rb.with_projection(projection) if projection is not None else rb).new_scan() - ranged = [] - for s in scan.plan().splits(): - lo, hi = _split_key_range(s, range_col, key_type, file_io) - ranged.append((s, lo, hi)) - return ranged, schema_id + if partitions: + # Build the partition predicate before projection, so its field list still has + # the partition columns. + pb = rb.new_predicate_builder() + rb = rb.with_partition_filter( + PredicateBuilder.and_predicates([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()) + 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, range_col, key_type, file_io), splits) + return [(s, lo, hi) for s, (lo, hi) in zip(splits, bounds)], schema_id def _cut_points(ranged_sides, num_ranges): @@ -198,6 +210,8 @@ def range_join( 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": @@ -205,8 +219,9 @@ def range_join( ``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. Sides must not share column names other than ``on`` keys. - Returns a ``ray.data.Dataset``. + 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 @@ -282,9 +297,9 @@ def range_join( 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, 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, catalog_options, right_projection, r_range_col, right_partitions) def _empty(): empty = read_splits( diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index de41eb3cf900..09ea66c2df3e 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -186,6 +186,50 @@ def test_rejects_key_type_mismatch(self): 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 From be96fd455724b99f1dc9e8d93e783d7626c03064 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 22 Jul 2026 20:45:15 -0700 Subject: [PATCH 09/12] [python] range_join: bound re-read of wide known splits, not just unknown ones The old budget only counted unknown-stats splits, but a known split whose [lo,hi] spans many ranges is also read once per range (no pushdown; read the whole split, clip in memory). On poorly clustered input that re-read could exceed a shuffle. Cap ranges by the actual total re-read -- rows times ranges each split overlaps -- halving num_ranges until it stays within a couple of full scans. Document that the feature only benefits clustered inputs. Verified locally: ray_range_join_test 18 passed (incl. the budget contract for wide/unknown splits), ray_bucket_join_test 18 passed. --- paimon-python/pypaimon/ray/range_join.py | 56 ++++++++++++------- .../pypaimon/tests/ray_range_join_test.py | 27 ++++++--- 2 files changed, 56 insertions(+), 27 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 24908b7126ab..4fbfc5b3134d 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -19,12 +19,13 @@ The driver cuts the key space into ranges from each file's parquet-footer min/max (the manifest value stats are empty for pypaimon-written tables); each range is read and -joined in its own Ray task, with no global shuffle. Works best when both sides are -clustered by the first join key. +joined in its own Ray task, with no global shuffle. Only beneficial when both sides are +clustered by the first join key: a poorly clustered (wide) split overlaps many ranges and +is read in each, so on unclustered input this can cost more than a shuffle. -Correctness never depends on stats: a split whose min/max is missing joins every range. -Such a split is read once per range, so the range count is capped to keep those re-reads -under one extra full scan (all-no-stats collapses to a single range). +Correctness never depends on stats: a split whose min/max is missing (or is wide) joins +every range it overlaps and is clipped in memory. The range count is reduced until that +total re-read stays within a couple of full scans (all-unknown collapses to one range). """ from typing import Any, Dict, List, Optional @@ -41,6 +42,8 @@ __all__ = ["range_join"] _MAX_RANGES = 512 +# Cap total re-read to this many full scans (see _bounded_ranges). +_REREAD_BUDGET = 2 def _parquet_col_range(metadata, col): @@ -161,18 +164,32 @@ def _cut_points(ranged_sides, num_ranges): return cuts -def _range_budget(l_ranged, r_ranged): - # Cap ranges so unknown-stats splits (read in every range) re-read <= one full scan. - total = unknown = 0 +def _split_rows(split): + return sum(f.row_count for f in split.files) + + +def _total_reads(l_ranged, r_ranged, ranges): + """Rows physically read = each split's rows 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, _ in ranged: - rows = sum(f.row_count for f in split.files) - total += rows - if lo is None: - unknown += rows - if unknown <= 0: - return _MAX_RANGES - return max(1, total // unknown) + 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_rows(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, so poorly clustered input can't cost far more than one scan.""" + total_rows = sum(_split_rows(s) + for ranged in (l_ranged, r_ranged) for s, _, _ in ranged) + budget = _REREAD_BUDGET * max(1, total_rows) + 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): @@ -315,10 +332,9 @@ def _empty(): num_ranges = max(1, min(_MAX_RANGES, max(len(l_ranged), len(r_ranged)))) elif num_ranges < 1: raise ValueError(f"num_ranges must be >= 1; got {num_ranges}.") - # An unknown-stats split is read in every range. Cap ranges so those re-reads add at - # most one extra full scan, else the fallback can cost more than a shuffle. - num_ranges = min(num_ranges, _range_budget(l_ranged, r_ranged)) - ranges = _ranges_from_cuts(_cut_points((l_ranged, r_ranged), num_ranges)) + # Reduce ranges until the total re-read (wide known splits + unknown-stats splits, + # both read in every range they overlap) 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 diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index 09ea66c2df3e..affd6699ad49 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -338,15 +338,28 @@ def test_int_to_string_schema_evolution_no_dropped_rows(self): got = sorted((r["k"], r["val"]) for r in ds.take_all()) self.assertEqual(got, [("100", "v100"), ("42", "v42"), ("5", "v5")]) - def test_range_budget_caps_ranges_when_stats_missing(self): + def test_reread_budget_bounds_wide_and_unknown_splits(self): Split = collections.namedtuple("Split", "files") File = collections.namedtuple("File", "row_count") - known = [(Split([File(100)]), 0, 99)] # 100 rows with stats - unknown = [(Split([File(100)]), None, None)] # 100 rows without stats - # unknown == total/2 -> budget 2; all-known -> no cap; all-unknown -> 1. - self.assertEqual(rjmod._range_budget(known, unknown), 2) - self.assertEqual(rjmod._range_budget(known, known), rjmod._MAX_RANGES) - self.assertEqual(rjmod._range_budget(unknown, unknown), 1) + + def rng(lo, hi, rows=100): + return (Split([File(rows)]), lo, hi) + + # _total_reads = rows x ranges a split overlaps. + self.assertEqual(rjmod._total_reads([rng(0, 10)], [], [(None, 5), (5, None)]), 200) + # All-unknown collapses to a single range. + unknown = [(Split([File(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): # The planner reads a file's min/max for the range column from manifest stats. From 8bb988e55ebd696bf24519a94b4ffb2e7d32dd37 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 22 Jul 2026 21:27:58 -0700 Subject: [PATCH 10/12] [python] range_join: trust footer bounds only when safe, budget by bytes, validate inputs Review fixes: - Resolve the range key's footer column by field id, so a renamed/swapped (same-type) column reads its own stats instead of another column's. - Treat a PK table's non-PK range key as unknown: merge (aggregation/partial update) can move it past the file min/max, which would silently drop matches. - Bound re-read by file bytes, not row count, so a wide-row split can't be re-read hundreds of times within budget. - Validate partition keys (reject non-partition columns; None -> is_null) and num_ranges (int >= 1, capped at _MAX_RANGES even when passed explicitly). - Warn instead of silently degrading when a parquet footer read fails. --- paimon-python/pypaimon/ray/range_join.py | 109 +++++++++++++----- .../pypaimon/tests/ray_range_join_test.py | 51 +++++++- 2 files changed, 126 insertions(+), 34 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 4fbfc5b3134d..7dfc85919786 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -28,6 +28,8 @@ total re-read stays within a couple of full scans (all-unknown collapses to one range). """ +import logging +import threading from typing import Any, Dict, List, Optional from pypaimon.ray.join_common import ( @@ -41,8 +43,10 @@ __all__ = ["range_join"] +_LOG = logging.getLogger(__name__) + _MAX_RANGES = 512 -# Cap total re-read to this many full scans (see _bounded_ranges). +# Cap total re-read to this many full scans of the read bytes (see _bounded_ranges). _REREAD_BUDGET = 2 @@ -72,27 +76,34 @@ def _footer_col_type(metadata, col): return None -def _split_key_range(split, col, key_type, file_io): - """Min/max of ``col`` over the split's files, read from their parquet footers; - (None, None) when any file isn't parquet, lacks usable stats, or stored the key in a - different type than the current one (the split then joins every range). A different - type means schema evolution, and its footer bounds order differently (e.g. INT '10' - < '2' as a string), so they can't be trusted -- only same-type stats are used. - Manifest value stats are empty for pypaimon-written tables, so the footer is the - source of truth.""" +def _split_key_range(split, name_for_schema, key_type, file_io): + """Min/max of the range key over the split's parquet footers; (None, None) -> the + split joins every range. ``name_for_schema`` maps a file's schema id to the key's + physical name by field id, so a renamed/swapped column reads its own stats. Unknown + when the key is absent in that schema, its footer type differs from the current key + type (order-incompatible evolution), or stats are missing. Footer, not manifest: + pypaimon writes empty manifest stats.""" import pyarrow.parquet as pq lo, hi = None, None for f in split.files: + 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 - stream = file_io.new_input_stream(path) try: - metadata = pq.read_metadata(stream) - except Exception: + 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 - finally: - stream.close() if _footer_col_type(metadata, col) != key_type: return None, None rng = _parquet_col_range(metadata, col) @@ -121,20 +132,45 @@ def _plan_ranged_splits(table_id, catalog_options, projection, range_col, partit 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. + # 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.equal(c, v) for c, v in partitions.items()])) + 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()) + # Footer min/max bounds the key only if merge can't move it out of range: safe when + # append-only, or the key is a primary key (merge never rewrites PKs). Otherwise (an + # aggregated/updated value column) every split joins every range, clipped in memory. + if table.primary_keys and range_col not in table.primary_keys: + return [(s, None, None) for s in splits], schema_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, range_col, key_type, file_io), splits) + lambda s: _split_key_range(s, name_for_schema, key_type, file_io), splits) return [(s, lo, hi) for s, (lo, hi) in zip(splits, bounds)], schema_id @@ -168,23 +204,28 @@ 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): - """Rows physically read = each split's rows times the ranges it overlaps (each range + """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_rows(split) * spans + 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, so poorly clustered input can't cost far more than one scan.""" - total_rows = sum(_split_rows(s) - for ranged in (l_ranged, r_ranged) for s, _, _ in ranged) - budget = _REREAD_BUDGET * max(1, total_rows) + 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: @@ -266,6 +307,16 @@ def range_join( 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 + # (a non-partition column) and read the whole table. + 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: @@ -329,11 +380,11 @@ def _empty(): return _empty() if num_ranges is None: - num_ranges = max(1, min(_MAX_RANGES, max(len(l_ranged), len(r_ranged)))) - elif num_ranges < 1: - raise ValueError(f"num_ranges must be >= 1; got {num_ranges}.") - # Reduce ranges until the total re-read (wide known splits + unknown-stats splits, - # both read in every range they overlap) stays bounded -- see _bounded_ranges. + 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): diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index affd6699ad49..e22a8007d8e2 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -340,15 +340,18 @@ def test_int_to_string_schema_evolution_no_dropped_rows(self): def test_reread_budget_bounds_wide_and_unknown_splits(self): Split = collections.namedtuple("Split", "files") - File = collections.namedtuple("File", "row_count") + File = collections.namedtuple("File", "row_count file_size") - def rng(lo, hi, rows=100): - return (Split([File(rows)]), lo, hi) + def rng(lo, hi, rows=100, size=100): + return (Split([File(rows, size)]), lo, hi) - # _total_reads = rows x ranges a split overlaps. + # _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)]), None, None)] + 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)] @@ -413,5 +416,43 @@ def test_null_keys_dropped_independent_of_num_ranges(self): 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_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) + + if __name__ == "__main__": unittest.main() From 0b396a40d95b725e44d177b2dc5323b9ead40fd8 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 22 Jul 2026 23:08:42 -0700 Subject: [PATCH 11/12] [python] range_join: reject nested and TIMESTAMP_LTZ range keys at the driver Fail fast for range keys that can't be range-partitioned instead of raising deep inside a Ray worker: nested/complex types (ARRAY/MAP/ROW/...) and both string forms of TIMESTAMP WITH LOCAL TIME ZONE (TIMESTAMP_LTZ(p) as well as ... WITH LOCAL TIME ZONE, only the latter was caught before). --- paimon-python/pypaimon/ray/range_join.py | 29 ++++++++++--------- .../pypaimon/tests/ray_range_join_test.py | 13 +++++++++ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 7dfc85919786..92705a47cb6c 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -307,8 +307,7 @@ def range_join( 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 - # (a non-partition column) and read the whole table. + # 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 [] @@ -331,20 +330,22 @@ def range_join( "range_join key columns must have the same type on both sides; " f"mismatched (left, right, left type, right type): {type_mismatch}.") - # The range key is the first pair; reject types that can't be range-partitioned safely. + # The range key is the first pair; reject up front (not inside a worker) the types + # that can't be range-partitioned safely. range_key_type = key_type(ltable, lkeys[0]).upper() - if range_key_type.startswith("FLOAT") or range_key_type.startswith("DOUBLE"): - # NaN compares false to every bound, so it would fall out of every range while - # pyarrow's hash join still matches NaN == NaN -> silently dropped matches. + 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" + elif "<" in range_key_type: # ARRAY<>/MAP<>/ROW<>/... -- not orderable/hashable here + reason = "a nested/complex type" + if reason: raise ValueError( - f"range_join range key {lkeys[0]!r} must not be FLOAT/DOUBLE (NaN can't be " - "range-partitioned); use an integer/string/date key.") - if "LOCAL TIME ZONE" in range_key_type: - # Manifest stats decode to naive datetimes; a tz-aware Arrow column can't be - # compared against them. Not supported yet. - raise ValueError( - f"range_join range key {lkeys[0]!r} of type TIMESTAMP WITH LOCAL TIME ZONE " - "is not supported yet.") + 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): diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index e22a8007d8e2..fe9d9ab050a0 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -453,6 +453,19 @@ def test_num_ranges_validation(self): 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") + if __name__ == "__main__": unittest.main() From 0ffa79858b55e9bc9fa09eb9d92e82140d873fbe Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 22 Jul 2026 23:35:16 -0700 Subject: [PATCH 12/12] [python] range_join: untrust footer under auth column-masking; validate all keys - P0: a range key under auth column-masking is rewritten at read time, so raw parquet footer min/max no longer bounds it -> treat as unknown (join every range, clip in memory), same fallback as a merge-rewritten non-PK key. - Validate every join key (not just the range key) up front and reject VARIANT as well as nested types, instead of failing as ArrowInvalid inside a worker. - Fix a stale test comment (footer, not manifest, stats). --- paimon-python/pypaimon/ray/range_join.py | 24 ++++++++++++------- .../pypaimon/tests/ray_range_join_test.py | 9 +++++-- 2 files changed, 23 insertions(+), 10 deletions(-) diff --git a/paimon-python/pypaimon/ray/range_join.py b/paimon-python/pypaimon/ray/range_join.py index 92705a47cb6c..3451b3b2e2b0 100644 --- a/paimon-python/pypaimon/ray/range_join.py +++ b/paimon-python/pypaimon/ray/range_join.py @@ -162,10 +162,14 @@ def name_for_schema(sid): if projection is not None: rb = rb.with_projection(projection) splits = list(rb.new_scan().plan().splits()) - # Footer min/max bounds the key only if merge can't move it out of range: safe when - # append-only, or the key is a primary key (merge never rewrites PKs). Otherwise (an - # aggregated/updated value column) every split joins every range, clipped in memory. - if table.primary_keys and range_col not in table.primary_keys: + # Footer min/max bounds the key only if the read can't move it out of range. Untrust + # (join every range, clip in memory) when it can: (1) a PK table's non-PK key, which + # merge (aggregation/partial-update) may rewrite; (2) a key under auth column-masking, + # rewritten at read time so raw footer stats no longer bound it. + 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) + if masked or (table.primary_keys and range_col not in table.primary_keys): return [(s, None, None) for s in splits], schema_id workers = min(16, (os.cpu_count() or 4) * 4, len(splits) or 1) with ThreadPoolExecutor(max_workers=workers) as pool: @@ -330,8 +334,14 @@ def range_join( "range_join key columns must have the same type on both sides; " f"mismatched (left, right, left type, right type): {type_mismatch}.") - # The range key is the first pair; reject up front (not inside a worker) the types - # that can't be range-partitioned safely. + # 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")): @@ -340,8 +350,6 @@ def range_join( 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" - elif "<" in range_key_type: # ARRAY<>/MAP<>/ROW<>/... -- not orderable/hashable here - reason = "a nested/complex type" if reason: raise ValueError( f"range_join range key {lkeys[0]!r} must not be {reason}; " diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py index fe9d9ab050a0..7ca14ebbb39b 100644 --- a/paimon-python/pypaimon/tests/ray_range_join_test.py +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -365,7 +365,7 @@ def rng(lo, hi, rows=100, size=100): len(rjmod._bounded_ranges(wide, wide, 4))) def test_split_key_range_reads_stats(self): - # The planner reads a file's min/max for the range column from manifest stats. + # The planner reads a file's min/max for the range column from the parquet 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)]) @@ -415,7 +415,6 @@ def test_null_keys_dropped_independent_of_num_ranges(self): 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. @@ -465,6 +464,12 @@ def test_rejects_unrangeable_key_types(self): 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__":