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..3451b3b2e2b0 --- /dev/null +++ b/paimon-python/pypaimon/ray/range_join.py @@ -0,0 +1,425 @@ +# 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 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. 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 (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). +""" + +import logging +import threading +from typing import Any, Dict, List, Optional + +from pypaimon.ray.join_common import ( + OnSpec, + get_table, + key_type, + norm_on, + pin_latest_snapshot, + read_splits, +) + +__all__ = ["range_join"] + +_LOG = logging.getLogger(__name__) + +_MAX_RANGES = 512 +# Cap total re-read to this many full scans of the read bytes (see _bounded_ranges). +_REREAD_BUDGET = 2 + + +def _parquet_col_range(metadata, col): + """Min/max of ``col`` across a parquet file's row groups; None when a row group + lacks usable stats for ``col``.""" + lo, hi = None, None + for i in range(metadata.num_row_groups): + rg = metadata.row_group(i) + stats = None + for j in range(rg.num_columns): + if rg.column(j).path_in_schema == col: + stats = rg.column(j).statistics + break + if stats is None or not stats.has_min_max: + return None + lo = stats.min if lo is None else min(lo, stats.min) + hi = stats.max if hi is None else max(hi, stats.max) + return None if lo is None else (lo, hi) + + +def _footer_col_type(metadata, col): + """The arrow type ``col`` is stored as in this parquet file; None if unavailable.""" + try: + return metadata.schema.to_arrow_schema().field(col).type + except Exception: + return None + + +def _split_key_range(split, name_for_schema, 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 + try: + stream = file_io.new_input_stream(path) + try: + metadata = pq.read_metadata(stream) + finally: + stream.close() + except Exception as e: + # Footer read can fail (e.g. local-cache streams aren't seekable): degrade to + # unknown, but warn -- the fallback would otherwise be silent. + _LOG.warning("range_join: parquet footer read failed for %s (%s); treating " + "its range as unknown", path, e) + return None, None + if _footer_col_type(metadata, col) != key_type: + return None, None + rng = _parquet_col_range(metadata, col) + if rng is None: + return None, None + lo = rng[0] if lo is None else min(lo, rng[0]) + hi = rng[1] if hi is None else max(hi, rng[1]) + if lo is None: + return None, None + return lo, hi + + +def _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 (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 + 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 + + # Range key's physical name in a file's schema, by field id (rename/swap safe). + # Cached; current-schema files skip the schema load. + key_field_id = next(f.id for f in table.table_schema.fields if f.name == range_col) + name_cache, cache_lock = {schema_id: range_col}, threading.Lock() + + def name_for_schema(sid): + with cache_lock: + if sid in name_cache: + return name_cache[sid] + try: + fields = table.schema_manager.get_schema(sid).fields + name = next((f.name for f in fields if f.id == key_field_id), None) + except Exception: + name = None + with cache_lock: + name_cache[sid] = name + return name + + rb = table.new_read_builder() + if partitions: + # Build the partition predicate before projection, so its field list still has + # the partition columns. None means the null partition (is_null, not = None). + pb = rb.new_predicate_builder() + rb = rb.with_partition_filter(PredicateBuilder.and_predicates( + [pb.is_null(c) if v is None else pb.equal(c, v) + for c, v in partitions.items()])) + if projection is not None: + rb = rb.with_projection(projection) + splits = list(rb.new_scan().plan().splits()) + # 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: + bounds = pool.map( + 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 + + +def _cut_points(ranged_sides, num_ranges): + """Pick ``num_ranges - 1`` cut values from row-count-weighted file boundaries.""" + points = [] + for ranged in ranged_sides: + for split, lo, hi in ranged: + if lo is None: + continue + rows = sum(f.row_count for f in split.files) + points.append((lo, rows / 2.0)) + points.append((hi, rows / 2.0)) + if not points: + return [] + points.sort(key=lambda p: p[0]) + total = sum(w for _, w in points) + cuts, acc, k = [], 0.0, 1 + for value, weight in points: + acc += weight + if k >= num_ranges: + break + if acc >= total * k / num_ranges: + if not cuts or value > cuts[-1]: # strictly increasing + cuts.append(value) + k += 1 + return cuts + + +def _split_rows(split): + return sum(f.row_count for f in split.files) + + +def _split_bytes(split): + # Re-read cost is bytes, not rows: a wide-row split is cheap by rows, costly by I/O. + return sum(f.file_size for f in split.files) + + +def _total_reads(l_ranged, r_ranged, ranges): + """Bytes physically read = each split's bytes times the ranges it overlaps (each range + reads the whole split and clips). Counts unknown-stats and wide known splits alike.""" + reads = 0 + for ranged in (l_ranged, r_ranged): + for split, lo, hi in ranged: + spans = sum(1 for r_lo, r_hi in ranges if _overlaps(lo, hi, r_lo, r_hi)) + reads += _split_bytes(split) * spans + return reads + + +def _bounded_ranges(l_ranged, r_ranged, num_ranges): + """Cut into ``num_ranges`` ranges, halving until total re-read <= _REREAD_BUDGET full + scans of bytes, so poorly clustered input can't cost far more than one scan.""" + total_bytes = sum(_split_bytes(s) + for ranged in (l_ranged, r_ranged) for s, _, _ in ranged) + budget = _REREAD_BUDGET * max(1, total_bytes) + while True: + ranges = _ranges_from_cuts(_cut_points((l_ranged, r_ranged), num_ranges)) + if num_ranges <= 1 or _total_reads(l_ranged, r_ranged, ranges) <= budget: + return ranges + num_ranges = max(1, num_ranges // 2) + + +def _ranges_from_cuts(cuts): + # Half-open [lo, hi); None = unbounded end. + bounds = [None] + cuts + [None] + return [(bounds[i], bounds[i + 1]) for i in range(len(bounds) - 1)] + + +def _overlaps(lo, hi, r_lo, r_hi): + if lo is None: # unknown split range: belongs to every range + return True + return (r_lo is None or hi >= r_lo) and (r_hi is None or lo < r_hi) + + +def _restrict_to_range(arrow_table, col, lo, hi): + """Keep rows with ``lo <= col < hi``. Null keys are always dropped (an inner + join never matches them), which also keeps the result independent of num_ranges.""" + import pyarrow.compute as pc + mask = pc.is_valid(arrow_table[col]) + if lo is not None: + mask = pc.and_(mask, pc.greater_equal(arrow_table[col], lo)) + if hi is not None: + mask = pc.and_(mask, pc.less(arrow_table[col], hi)) + return arrow_table.filter(mask) + + +def range_join( + left: str, + right: str, + catalog_options: Dict[str, str], + *, + on: Optional[OnSpec] = None, + left_on: Optional[OnSpec] = None, + right_on: Optional[OnSpec] = None, + num_ranges: Optional[int] = None, + left_projection: Optional[List[str]] = None, + right_projection: Optional[List[str]] = None, + left_partitions: Optional[Dict[str, Any]] = None, + right_partitions: Optional[Dict[str, Any]] = None, + join_type: str = "inner", + ray_remote_args: Optional[Dict[str, Any]] = None, +) -> "ray.data.Dataset": + """Join two tables clustered by the first join key with no global shuffle. + + ``on`` when both sides use the same column names, or ``left_on``/``right_on`` + when they differ (positionally paired). The first pair is the range key used to + cut the key space. ``left_partitions``/``right_partitions`` ({column: value} dicts + on partition columns) prune each side to those partitions first. Sides must not + share column names other than ``on`` keys. Returns a ``ray.data.Dataset``. + """ + import ray + + if not hasattr(ray.data, "from_arrow_refs"): + raise RuntimeError( + "range_join needs a Ray version with ray.data.from_arrow_refs; " + f"installed ray is {ray.__version__}.") + + if (on is None) == (left_on is None and right_on is None): + raise ValueError("range_join requires exactly one of on= or left_on=/right_on=.") + if on is not None: + lkeys = rkeys = norm_on(on) + else: + if left_on is None or right_on is None: + raise ValueError("range_join requires both left_on= and right_on=.") + lkeys, rkeys = norm_on(left_on), norm_on(right_on) + if len(lkeys) != len(rkeys) or not lkeys: + raise ValueError( + f"range_join join keys must pair up non-empty; got left_on={lkeys}, right_on={rkeys}.") + if join_type != "inner": + # Outer joins would need every unmatched row emitted exactly once across + # ranges plus null-key handling; only inner is supported for now. + raise ValueError(f"range_join currently supports only join_type='inner'; got {join_type!r}.") + + ltable = get_table(left, catalog_options, None, "range_join") + rtable = get_table(right, catalog_options, None, "range_join") + + # Partition filters must name partition columns, else they'd be silently ignored. + for name, tbl, parts in (("left_partitions", ltable, left_partitions), + ("right_partitions", rtable, right_partitions)): + bad = sorted(set(parts) - set(tbl.partition_keys)) if parts else [] + if bad: + raise ValueError( + f"range_join {name} keys {bad} are not partition columns; " + f"partition columns are {list(tbl.partition_keys)}.") + + missing = [c for c in lkeys if c not in ltable.field_dict] \ + + [c for c in rkeys if c not in rtable.field_dict] + if missing: + raise ValueError(f"range_join keys not found in table schema: {missing}.") + type_mismatch = [ + (lc, rc, key_type(ltable, lc), key_type(rtable, rc)) + for lc, rc in zip(lkeys, rkeys) + if key_type(ltable, lc) != key_type(rtable, rc) + ] + if type_mismatch: + raise ValueError( + "range_join key columns must have the same type on both sides; " + f"mismatched (left, right, left type, right type): {type_mismatch}.") + + # Reject unsupported key types up front (not inside a worker as ArrowInvalid). Every + # join key must be hashable; nested (ARRAY<>/MAP<>/ROW<>/...) and VARIANT are not. + for c in lkeys: + t = key_type(ltable, c).upper() + if "<" in t or t.startswith("VARIANT"): + raise ValueError( + f"range_join join key {c!r} must not be a nested/complex type; got {t}.") + # The range key (first pair) additionally must be range-partitionable. + range_key_type = key_type(ltable, lkeys[0]).upper() + reason = None + if range_key_type.startswith(("FLOAT", "DOUBLE")): + # NaN falls out of every range while the hash join still matches it -> drops rows. + reason = "FLOAT/DOUBLE" + elif "LOCAL TIME ZONE" in range_key_type or "TIMESTAMP_LTZ" in range_key_type: + # Footer stats decode to naive datetimes; a tz-aware column can't compare to them. + reason = "TIMESTAMP WITH LOCAL TIME ZONE" + if reason: + raise ValueError( + f"range_join range key {lkeys[0]!r} must not be {reason}; " + "use an integer/string/date/timestamp key.") + + # The join keys must survive projection, or the local join has no key. + if left_projection is not None and not set(lkeys) <= set(left_projection): + raise ValueError( + f"left_projection must include the join keys {lkeys}; got {left_projection}.") + if right_projection is not None and not set(rkeys) <= set(right_projection): + raise ValueError( + f"right_projection must include the join keys {rkeys}; got {right_projection}.") + # pyarrow drops the right keys (coalesced into the left), so the output keeps the LEFT + # key names. A right non-key column sharing a left column name collides -> reject it. + lcols = left_projection if left_projection is not None else ltable.field_names + rcols = right_projection if right_projection is not None else rtable.field_names + collisions = sorted(set(lcols) & (set(rcols) - set(rkeys))) + if collisions: + raise ValueError( + f"range_join output columns collide: {collisions}. The output keeps the left " + "key names and the right non-key columns; project or rename the overlap away.") + + l_range_col, r_range_col = lkeys[0], rkeys[0] + l_ranged, l_schema_id = _plan_ranged_splits( + left, catalog_options, left_projection, l_range_col, left_partitions) + r_ranged, r_schema_id = _plan_ranged_splits( + right, catalog_options, right_projection, r_range_col, right_partitions) + + def _empty(): + empty = read_splits( + left, catalog_options, left_projection, [], l_schema_id, "range_join").join( + read_splits(right, catalog_options, right_projection, [], r_schema_id, "range_join"), + keys=lkeys, right_keys=rkeys, join_type=join_type) + return ray.data.from_arrow(empty) + + if not l_ranged or not r_ranged: # inner join: one empty side, empty result + return _empty() + + if num_ranges is None: + num_ranges = max(len(l_ranged), len(r_ranged)) + elif not isinstance(num_ranges, int) or num_ranges < 1: + raise ValueError(f"num_ranges must be an int >= 1; got {num_ranges!r}.") + num_ranges = max(1, min(_MAX_RANGES, num_ranges)) # cap tasks even when explicit + # Reduce ranges until total re-read stays bounded -- see _bounded_ranges. + ranges = _bounded_ranges(l_ranged, r_ranged, num_ranges) + + def _join_range(left_splits, right_splits, lo, hi): + # No predicate pushdown: the range key may be schema-evolved (e.g. a file stored + # as INT read as STRING), which the reader can't compare against a new-type bound. + # The in-memory clip below does the exact, evolution-safe filtering. + lt = _restrict_to_range( + read_splits(left, catalog_options, left_projection, left_splits, + l_schema_id, "range_join"), + l_range_col, lo, hi) + rt = _restrict_to_range( + read_splits(right, catalog_options, right_projection, right_splits, + r_schema_id, "range_join"), + r_range_col, lo, hi) + return lt.join(rt, keys=lkeys, right_keys=rkeys, join_type=join_type) + + # ``@ray.remote()`` (empty parens) is rejected by Ray, so wrap conditionally. + remote_fn = ray.remote(**ray_remote_args)(_join_range) if ray_remote_args else ray.remote(_join_range) + refs = [] + for r_lo, r_hi in ranges: + ls = [s for s, lo, hi in l_ranged if _overlaps(lo, hi, r_lo, r_hi)] + rs = [s for s, lo, hi in r_ranged if _overlaps(lo, hi, r_lo, r_hi)] + if not ls or not rs: # inner join: a one-sided range can't match + continue + refs.append(remote_fn.remote(ls, rs, r_lo, r_hi)) + if not refs: + return _empty() + # Keep each range's result as a distributed object ref -- never pulled into the driver. + return ray.data.from_arrow_refs(refs) diff --git a/paimon-python/pypaimon/tests/ray_range_join_test.py b/paimon-python/pypaimon/tests/ray_range_join_test.py new file mode 100644 index 000000000000..7ca14ebbb39b --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_range_join_test.py @@ -0,0 +1,476 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import collections +import datetime +import os +import shutil +import tempfile +import unittest + +import pyarrow as pa +import pytest + +pypaimon = pytest.importorskip("pypaimon") +ray = pytest.importorskip("ray") + +import importlib +from unittest import mock + +from pypaimon import CatalogFactory, Schema +from pypaimon.ray import range_join + +rjmod = importlib.import_module("pypaimon.ray.range_join") + + +class RayRangeJoinTest(unittest.TestCase): + """Range-aligned join must equal a global inner join, cutting the key space from + per-file min/max stats so each range is read/joined in its own task (no shuffle).""" + + @classmethod + def setUpClass(cls): + cls.tempdir = tempfile.mkdtemp() + cls.catalog_options = {"warehouse": os.path.join(cls.tempdir, "wh")} + cls.catalog = CatalogFactory.create(cls.catalog_options) + cls.catalog.create_database("default", True) + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True, num_cpus=4) + + @classmethod + def tearDownClass(cls): + try: + if ray.is_initialized(): + ray.shutdown() + except Exception: + pass + shutil.rmtree(cls.tempdir, ignore_errors=True) + + def _table(self, name, schema, commits, primary_keys=None, options=None): + """Create a table and write each arrow table in ``commits`` as its own commit, + so the manifest holds several data files with distinct key ranges.""" + self.catalog.create_table( + name, + Schema.from_pyarrow_schema(schema, primary_keys=primary_keys, options=options), + False) + t = self.catalog.get_table(name) + for data in commits: + wb = t.new_batch_write_builder() + w = wb.new_write() + w.write_arrow(data) + wb.new_commit().commit(w.prepare_commit()) + w.close() + return name + + def test_range_join_matches_global_join(self): + loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) + ins = pa.schema([("k", pa.int64())]) + # locator: k in 0..599 spread across three files with disjoint key ranges. + self._table("default.rj_loc", loc, [ + pa.Table.from_pydict({"k": list(range(0, 200)), + "row_id": list(range(0, 200))}, schema=loc), + pa.Table.from_pydict({"k": list(range(200, 400)), + "row_id": list(range(200, 400))}, schema=loc), + pa.Table.from_pydict({"k": list(range(400, 600)), + "row_id": list(range(400, 600))}, schema=loc), + ]) + self._table("default.rj_in", ins, [ + pa.Table.from_pydict({"k": list(range(0, 250))}, schema=ins), + ]) + ds = range_join( + "default.rj_in", "default.rj_loc", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k", "row_id"], num_ranges=4) + got = {r["k"]: r["row_id"] for r in ds.take_all()} + self.assertEqual(set(got), set(range(250))) + self.assertTrue(all(got[i] == i for i in range(250))) + + def test_fan_out_one_key_many_rows(self): + loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_fan_loc", loc, [ + pa.Table.from_pydict({"k": [5, 5, 7], "row_id": [0, 1, 2]}, schema=loc)]) + self._table("default.rj_fan_in", ins, [ + pa.Table.from_pydict({"k": [5]}, schema=ins)]) + ds = range_join( + "default.rj_fan_in", "default.rj_fan_loc", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k", "row_id"]) + self.assertEqual(sorted(r["row_id"] for r in ds.take_all()), [0, 1]) + + def test_left_on_right_on_different_names(self): + right = pa.schema([("rid", pa.int64()), ("val", pa.string())]) + left = pa.schema([("lid", pa.int64())]) + self._table("default.rj_lr_right", right, [ + pa.Table.from_pydict({"rid": list(range(100)), + "val": [f"v{i}" for i in range(100)]}, schema=right)]) + self._table("default.rj_lr_left", left, [ + pa.Table.from_pydict({"lid": list(range(30))}, schema=left)]) + ds = range_join( + "default.rj_lr_left", "default.rj_lr_right", self.catalog_options, + left_on="lid", right_on="rid", num_ranges=3) + # Output keeps the left key name (pyarrow coalesces the right key into it). + got = {r["lid"]: r["val"] for r in ds.take_all()} + self.assertEqual(got, {i: f"v{i}" for i in range(30)}) + + def test_num_ranges_one_is_correct(self): + # A single range degenerates to one local join and must still be exact. + loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_one_loc", loc, [ + pa.Table.from_pydict({"k": list(range(50)), + "row_id": list(range(50))}, schema=loc)]) + self._table("default.rj_one_in", ins, [ + pa.Table.from_pydict({"k": list(range(20))}, schema=ins)]) + ds = range_join( + "default.rj_one_in", "default.rj_one_loc", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k", "row_id"], num_ranges=1) + got = {r["k"]: r["row_id"] for r in ds.take_all()} + self.assertEqual(got, {i: i for i in range(20)}) + + def test_dispatches_multiple_range_tasks(self): + # No global shuffle: several disjoint-range files produce more than one task. + loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_disp_loc", loc, [ + pa.Table.from_pydict({"k": list(range(0, 300)), + "row_id": list(range(0, 300))}, schema=loc), + pa.Table.from_pydict({"k": list(range(300, 600)), + "row_id": list(range(300, 600))}, schema=loc), + ]) + self._table("default.rj_disp_in", ins, [ + pa.Table.from_pydict({"k": list(range(0, 600))}, schema=ins)]) + + captured = {} + real = ray.data.from_arrow_refs + + def spy(refs): + captured["n"] = len(refs) + return real(refs) + + with mock.patch.object(ray.data, "from_arrow_refs", spy): + ds = range_join( + "default.rj_disp_in", "default.rj_disp_loc", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k", "row_id"], num_ranges=4) + ds.take_all() + self.assertGreater(captured["n"], 1) + + def test_rejects_shared_non_key_column(self): + loc = pa.schema([("k", pa.int64()), ("v", pa.int64())]) + ins = pa.schema([("k", pa.int64()), ("v", pa.int64())]) + self._table("default.rj_col_loc", loc, [ + pa.Table.from_pydict({"k": [1], "v": [1]}, schema=loc)]) + self._table("default.rj_col_in", ins, [ + pa.Table.from_pydict({"k": [1], "v": [2]}, schema=ins)]) + with self.assertRaisesRegex(ValueError, "collide"): + range_join("default.rj_col_in", "default.rj_col_loc", self.catalog_options, on="k") + + def test_rejects_key_type_mismatch(self): + loc = pa.schema([("k", pa.int32()), ("row_id", pa.int64())]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_ty_loc", loc, [ + pa.Table.from_pydict({"k": pa.array([1], pa.int32()), "row_id": [1]}, schema=loc)]) + self._table("default.rj_ty_in", ins, [ + pa.Table.from_pydict({"k": [1]}, schema=ins)]) + with self.assertRaisesRegex(ValueError, "same type"): + range_join("default.rj_ty_in", "default.rj_ty_loc", self.catalog_options, on="k") + + def test_partition_filter_and_partitioned_table(self): + # range_join works on partitioned tables (unlike bucket_join); left_partitions + # prunes to the requested partition before planning. + loc = pa.schema([("p", pa.string()), ("k", pa.int64())]) + self.catalog.create_table( + "default.rj_pf_l", + Schema.from_pyarrow_schema(loc, partition_keys=["p"]), False) + t = self.catalog.get_table("default.rj_pf_l") + wb = t.new_batch_write_builder() + w = wb.new_write() + w.write_arrow(pa.Table.from_pydict( + {"p": ["a", "a", "b", "b"], "k": [1, 2, 3, 4]}, schema=loc)) + wb.new_commit().commit(w.prepare_commit()) + w.close() + self._table("default.rj_pf_r", pa.schema([("k2", pa.int64()), ("val", pa.string())]), [ + pa.Table.from_pydict({"k2": [1, 2, 3, 4], "val": ["v1", "v2", "v3", "v4"]}, + schema=pa.schema([("k2", pa.int64()), ("val", pa.string())]))]) + + ds = range_join("default.rj_pf_l", "default.rj_pf_r", self.catalog_options, + left_on="k", right_on="k2", left_projection=["k"], + right_projection=["k2", "val"], left_partitions={"p": "a"}, num_ranges=2) + self.assertEqual(sorted((r["k"], r["val"]) for r in ds.take_all()), + [(1, "v1"), (2, "v2")]) + # Whole partitioned table joins fine when unfiltered. + ds = range_join("default.rj_pf_l", "default.rj_pf_r", self.catalog_options, + left_on="k", right_on="k2", left_projection=["k"], + right_projection=["k2", "val"], num_ranges=2) + self.assertEqual(sorted(r["k"] for r in ds.take_all()), [1, 2, 3, 4]) + + def test_many_to_many_matches_global_join(self): + loc = pa.schema([("k", pa.int64()), ("rid", pa.int64())]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_mm_loc", loc, [ + pa.Table.from_pydict({"k": [1, 1, 2], "rid": [10, 11, 20]}, schema=loc)]) + self._table("default.rj_mm_in", ins, [ + pa.Table.from_pydict({"k": [1, 1, 2]}, schema=ins)]) + # key 1: 2 left x 2 right = 4 rows; key 2: 1 x 1 = 1 row. + for num_ranges in (1, 3): + ds = range_join("default.rj_mm_in", "default.rj_mm_loc", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k", "rid"], + num_ranges=num_ranges) + got = sorted(r["rid"] for r in ds.take_all()) + self.assertEqual(got, [10, 10, 11, 11, 20]) + + def test_rejects_bad_on_spec(self): + with self.assertRaisesRegex(ValueError, "exactly one of"): + range_join("a", "b", self.catalog_options) # neither on nor left_on/right_on + + def test_rejects_float_range_key(self): + schema = pa.schema([("k", pa.float64()), ("v", pa.int64())]) + self._table("default.rj_float_a", schema, [ + pa.Table.from_pydict({"k": [1.0], "v": [1]}, schema=schema)]) + self._table("default.rj_float_b", schema, [ + pa.Table.from_pydict({"k": [1.0], "v": [2]}, schema=schema)]) + with self.assertRaisesRegex(ValueError, "FLOAT/DOUBLE"): + range_join("default.rj_float_a", "default.rj_float_b", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k"]) + + def test_rejects_left_key_vs_right_column_collision(self): + left = pa.schema([("lid", pa.int64()), ("x", pa.int64())]) + right = pa.schema([("rid", pa.int64()), ("lid", pa.int64())]) + self._table("default.rj_xn_left", left, [ + pa.Table.from_pydict({"lid": [1], "x": [1]}, schema=left)]) + self._table("default.rj_xn_right", right, [ + pa.Table.from_pydict({"rid": [1], "lid": [9]}, schema=right)]) + # Left key 'lid' collides with the right non-key column 'lid' in the output. + with self.assertRaisesRegex(ValueError, "collide"): + range_join("default.rj_xn_left", "default.rj_xn_right", self.catalog_options, + left_on="lid", right_on="rid") + + def test_date_to_timestamp_schema_evolution(self): + # A DATE->TIMESTAMP evolved key yields date footers in old files and datetime in + # new ones; the planner must coerce both to the key type, not compare them raw. + from pypaimon.schema.data_types import AtomicType + from pypaimon.schema.schema_change import SchemaChange + + a_date = pa.schema([("k", pa.date32())]) + self.catalog.create_table( + "default.rj_ev_a", Schema.from_pyarrow_schema(a_date), False) + t = self.catalog.get_table("default.rj_ev_a") + wb = t.new_batch_write_builder() + w = wb.new_write() + w.write_arrow(pa.Table.from_pydict( + {"k": [datetime.date(2020, 1, 1), datetime.date(2020, 1, 2)]}, schema=a_date)) + wb.new_commit().commit(w.prepare_commit()) + w.close() + self.catalog.alter_table( + "default.rj_ev_a", + [SchemaChange.update_column_type("k", AtomicType("TIMESTAMP(6)"))], False) + t = self.catalog.get_table("default.rj_ev_a") + a_ts = pa.schema([("k", pa.timestamp("us"))]) + wb = t.new_batch_write_builder() + w = wb.new_write() + w.write_arrow(pa.Table.from_pydict( + {"k": [datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 2)]}, schema=a_ts)) + wb.new_commit().commit(w.prepare_commit()) + w.close() + + b = pa.schema([("bk", pa.timestamp("us")), ("val", pa.string())]) + self.catalog.create_table("default.rj_ev_b", Schema.from_pyarrow_schema(b), False) + t = self.catalog.get_table("default.rj_ev_b") + wb = t.new_batch_write_builder() + w = wb.new_write() + w.write_arrow(pa.Table.from_pydict( + {"bk": [datetime.datetime(2020, 1, 1), datetime.datetime(2020, 6, 1)], + "val": ["jan1", "jun1"]}, schema=b)) + wb.new_commit().commit(w.prepare_commit()) + w.close() + + ds = range_join("default.rj_ev_a", "default.rj_ev_b", self.catalog_options, + left_on="k", right_on="bk", num_ranges=3) + got = sorted((str(r["k"]), r["val"]) for r in ds.take_all()) + self.assertEqual(got, [("2020-01-01 00:00:00", "jan1"), + ("2020-06-01 00:00:00", "jun1")]) + + def test_int_to_string_schema_evolution_no_dropped_rows(self): + # INT->STRING isn't order-preserving ('10' < '2'), so an old INT file's footer + # bounds are invalid under the new string order. Such files must be treated as + # unknown (join every range), not pruned, or rows are silently dropped. + from pypaimon.schema.data_types import AtomicType + from pypaimon.schema.schema_change import SchemaChange + + a_int = pa.schema([("k", pa.int32())]) + self.catalog.create_table( + "default.rj_is_a", Schema.from_pyarrow_schema(a_int), False) + t = self.catalog.get_table("default.rj_is_a") + wb = t.new_batch_write_builder() + w = wb.new_write() + # int order 5<42<100, but as strings '100'<'42'<'5'. + w.write_arrow(pa.Table.from_pydict({"k": [5, 42, 100]}, schema=a_int)) + wb.new_commit().commit(w.prepare_commit()) + w.close() + self.catalog.alter_table( + "default.rj_is_a", + [SchemaChange.update_column_type("k", AtomicType("STRING"))], False) + + b = pa.schema([("bk", pa.string()), ("val", pa.string())]) + self.catalog.create_table("default.rj_is_b", Schema.from_pyarrow_schema(b), False) + t = self.catalog.get_table("default.rj_is_b") + wb = t.new_batch_write_builder() + w = wb.new_write() + w.write_arrow(pa.Table.from_pydict( + {"bk": ["5", "42", "100"], "val": ["v5", "v42", "v100"]}, schema=b)) + wb.new_commit().commit(w.prepare_commit()) + w.close() + + for num_ranges in (1, 3): + ds = range_join("default.rj_is_a", "default.rj_is_b", self.catalog_options, + left_on="k", right_on="bk", num_ranges=num_ranges) + got = sorted((r["k"], r["val"]) for r in ds.take_all()) + self.assertEqual(got, [("100", "v100"), ("42", "v42"), ("5", "v5")]) + + def test_reread_budget_bounds_wide_and_unknown_splits(self): + Split = collections.namedtuple("Split", "files") + File = collections.namedtuple("File", "row_count file_size") + + def rng(lo, hi, rows=100, size=100): + return (Split([File(rows, size)]), lo, hi) + + # _total_reads = bytes x ranges a split overlaps. + self.assertEqual(rjmod._total_reads([rng(0, 10)], [], [(None, 5), (5, None)]), 200) + # Budget is bytes, not rows: a wide split (few rows, large files) counts its bytes. + wide_row = [rng(0, 10, rows=1, size=1000)] + self.assertEqual(rjmod._total_reads(wide_row, [], [(None, 5), (5, None)]), 2000) + # All-unknown collapses to a single range. + unknown = [(Split([File(100, 100)]), None, None)] + self.assertEqual(len(rjmod._bounded_ranges(unknown, unknown, 8)), 1) + # Wide known splits (each overlaps many ranges) are bounded by the budget. + wide = [rng(0, 100), rng(0, 100), rng(0, 100), rng(0, 100)] + ranges = rjmod._bounded_ranges(wide, wide, 16) + budget = rjmod._REREAD_BUDGET * (8 * 100) + self.assertTrue(len(ranges) == 1 + or rjmod._total_reads(wide, wide, ranges) <= budget) + # Clustered (disjoint) splits keep at least as much parallelism as wide ones. + clustered = [rng(0, 9), rng(10, 19), rng(20, 29), rng(30, 39)] + self.assertGreaterEqual(len(rjmod._bounded_ranges(clustered, clustered, 4)), + len(rjmod._bounded_ranges(wide, wide, 4))) + + def test_split_key_range_reads_stats(self): + # 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)]) + 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_stats_mode_none_still_correct(self): + # metadata.stats-mode=none only drops manifest stats; the parquet footer still + # carries min/max (range_join's actual source), so ranges still work. The + # unknown-split fallback itself is covered by the planning-logic tests. + no_stats = {"metadata.stats-mode": "none"} + loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_ns_loc", loc, [ + pa.Table.from_pydict({"k": list(range(0, 100)), + "row_id": list(range(0, 100))}, schema=loc), + pa.Table.from_pydict({"k": list(range(100, 200)), + "row_id": list(range(100, 200))}, schema=loc), + ], options=no_stats) + self._table("default.rj_ns_in", ins, [ + pa.Table.from_pydict({"k": list(range(50, 150))}, schema=ins)], + options=no_stats) + ds = range_join( + "default.rj_ns_in", "default.rj_ns_loc", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k", "row_id"], num_ranges=4) + got = sorted((r["k"], r["row_id"]) for r in ds.take_all()) + self.assertEqual(got, [(i, i) for i in range(50, 150)]) + + def test_null_keys_dropped_independent_of_num_ranges(self): + loc = pa.schema([("k", pa.int64()), ("row_id", pa.int64())]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_null_loc", loc, [ + pa.Table.from_pydict({"k": [1, 2, None, 3], "row_id": [1, 2, 99, 3]}, schema=loc)]) + self._table("default.rj_null_in", ins, [ + pa.Table.from_pydict({"k": [1, None, 3, None]}, schema=ins)]) + expected = [(1, 1), (3, 3)] # null never matches; no duplicates + for num_ranges in (1, 5): + ds = range_join( + "default.rj_null_in", "default.rj_null_loc", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k", "row_id"], + num_ranges=num_ranges) + got = sorted((r["k"], r["row_id"]) for r in ds.take_all()) + self.assertEqual(got, expected) + + def test_pk_nonkey_range_col_untrusted(self): + # A PK table's non-PK column may be rewritten by merge (aggregation/partial-update) + # beyond the footer min/max, so its bounds are untrusted -> every split unknown. + schema = pa.schema([("id", pa.int64()), ("g", pa.int64())]) + self._table("default.rj_agg", schema, [ + pa.Table.from_pydict({"id": [1, 2], "g": [10, 20]}, schema=schema)], + primary_keys=["id"], options={"bucket": "1"}) + ranged, _ = rjmod._plan_ranged_splits( + "default.rj_agg", self.catalog_options, None, "g") + self.assertTrue(ranged) + self.assertTrue(all(lo is None and hi is None for _, lo, hi in ranged)) + + def test_partition_key_validation(self): + loc = pa.schema([("k", pa.int64()), ("v", pa.string())]) + self._table("default.rj_pv", loc, [ + pa.Table.from_pydict({"k": [1], "v": ["a"]}, schema=loc)]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_pv_in", ins, [ + pa.Table.from_pydict({"k": [1]}, schema=ins)]) + with self.assertRaises(ValueError): # not a partition column + range_join("default.rj_pv_in", "default.rj_pv", self.catalog_options, + on="k", left_projection=["k"], right_projection=["k", "v"], + right_partitions={"nope": "a"}) + + def test_num_ranges_validation(self): + loc = pa.schema([("k", pa.int64()), ("v", pa.string())]) + self._table("default.rj_nr", loc, [ + pa.Table.from_pydict({"k": [1, 2], "v": ["a", "b"]}, schema=loc)]) + ins = pa.schema([("k", pa.int64())]) + self._table("default.rj_nr_in", ins, [ + pa.Table.from_pydict({"k": [1]}, schema=ins)]) + for bad in (0, -1, "5", 2.0): + with self.assertRaises(ValueError): + range_join("default.rj_nr_in", "default.rj_nr", self.catalog_options, + on="k", left_projection=["k"], + right_projection=["k", "v"], num_ranges=bad) + + def test_rejects_unrangeable_key_types(self): + # Rejected at the driver (not inside a worker): nested and tz-aware keys. + arr = pa.schema([("k", pa.list_(pa.int64())), ("v", pa.string())]) + self._table("default.rj_arr_a", arr, []) + self._table("default.rj_arr_b", arr, []) + with self.assertRaisesRegex(ValueError, "must not be"): + range_join("default.rj_arr_a", "default.rj_arr_b", self.catalog_options, on="k") + ltz = pa.schema([("k", pa.timestamp("us", tz="UTC")), ("v", pa.string())]) + self._table("default.rj_ltz_a", ltz, []) + self._table("default.rj_ltz_b", ltz, []) + with self.assertRaisesRegex(ValueError, "must not be"): + range_join("default.rj_ltz_a", "default.rj_ltz_b", self.catalog_options, on="k") + # A nested SECOND key is rejected too (every key is validated, not just the range key). + multi = pa.schema([("k", pa.int64()), ("k2", pa.list_(pa.int64())), ("v", pa.string())]) + self._table("default.rj_mk_a", multi, []) + self._table("default.rj_mk_b", multi, []) + with self.assertRaisesRegex(ValueError, "nested/complex"): + range_join("default.rj_mk_a", "default.rj_mk_b", self.catalog_options, on=["k", "k2"]) + + +if __name__ == "__main__": + unittest.main()