| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| """Shuffle-free Ray join for tables clustered by the first join key. | |
| Ranges use manifest/key stats with a Parquet-footer fallback. Missing stats are | |
| safe: affected splits join every overlapping range and are filtered in memory. | |
| """ | |
| import logging | |
| import threading | |
| from typing import Any, Dict, List, Optional | |
| from pypaimon.ray.join_common import ( | |
| OnSpec, | |
| get_table, | |
| key_type, | |
| norm_on, | |
| pin_latest_snapshot, | |
| read_splits, | |
| ) | |
| __all__ = ["range_join"] | |
| _LOG = logging.getLogger(__name__) | |
| _MAX_RANGES = 512 | |
| # Cap total re-read to this many full scans of the read bytes (see _bounded_ranges). | |
| _REREAD_BUDGET = 2 | |
| def _stats_range(stats, field_id, key_type): | |
| """Read matching field-id/type bounds from manifest stats.""" | |
| from pypaimon.schema.data_types import PyarrowFieldParser | |
| min_row, max_row = stats.min_values, stats.max_values | |
| min_fields = getattr(min_row, "fields", []) or [] | |
| max_fields = getattr(max_row, "fields", []) or [] | |
| idx = next((i for i, field in enumerate(min_fields) if field.id == field_id), None) | |
| if idx is None or idx >= len(max_fields) or max_fields[idx].id != field_id: | |
| return None | |
| try: | |
| stored_type = PyarrowFieldParser.from_paimon_type(min_fields[idx].type) | |
| max_type = PyarrowFieldParser.from_paimon_type(max_fields[idx].type) | |
| if stored_type != key_type or max_type != key_type: | |
| return None | |
| lo, hi = min_row.get_field(idx), max_row.get_field(idx) | |
| except Exception: | |
| return None | |
| return None if lo is None or hi is None else (lo, hi) | |
| def _file_stats_range(file, field_id, key_type): | |
| """Prefer key stats, then value stats.""" | |
| return (_stats_range(file.key_stats, field_id, key_type) | |
| or _stats_range(file.value_stats, field_id, key_type)) | |
| def _file_identity(file): | |
| # External paths disambiguate imported files. | |
| return file.external_path or file.file_name | |
| def _manifest_stats_by_file(table, field_id, key_type, snapshot_id): | |
| """Read active-file stats for the pinned snapshot.""" | |
| from pypaimon.manifest.manifest_file_manager import ManifestFileManager | |
| from pypaimon.manifest.manifest_list_manager import ManifestListManager | |
| try: | |
| snapshot = table.snapshot_manager().get_snapshot_by_id(snapshot_id) | |
| manifests = ManifestListManager(table).read_all(snapshot) | |
| entries = ManifestFileManager(table).read_entries_parallel( | |
| manifests, drop_stats=False) | |
| except Exception as e: | |
| # Stats are optional. | |
| _LOG.warning( | |
| "range_join: manifest stats read failed (%s); falling back to file footers", e) | |
| return {} | |
| by_file, ambiguous = {}, set() | |
| for entry in entries: | |
| file = entry.file | |
| rng = _file_stats_range(file, field_id, key_type) | |
| if rng is None: | |
| continue | |
| identity = _file_identity(file) | |
| previous = by_file.get(identity) | |
| if previous is not None and previous != rng: | |
| ambiguous.add(identity) | |
| else: | |
| by_file[identity] = rng | |
| for identity in ambiguous: | |
| by_file.pop(identity, None) | |
| return by_file | |
| def _parquet_col_range(metadata, col): | |
| """Min/max of ``col`` across a parquet file's row groups; None when a row group | |
| lacks usable stats for ``col``.""" | |
| lo, hi = None, None | |
| for i in range(metadata.num_row_groups): | |
| rg = metadata.row_group(i) | |
| stats = None | |
| for j in range(rg.num_columns): | |
| if rg.column(j).path_in_schema == col: | |
| stats = rg.column(j).statistics | |
| break | |
| if stats is None or not stats.has_min_max: | |
| return None | |
| lo = stats.min if lo is None else min(lo, stats.min) | |
| hi = stats.max if hi is None else max(hi, stats.max) | |
| return None if lo is None else (lo, hi) | |
| def _footer_col_type(metadata, col): | |
| """The arrow type ``col`` is stored as in this parquet file; None if unavailable.""" | |
| try: | |
| return metadata.schema.to_arrow_schema().field(col).type | |
| except Exception: | |
| return None | |
| def _split_key_range( | |
| split, name_for_schema, field_id, key_type, file_io, manifest_stats): | |
| """Split bounds from key/manifest stats, then Parquet footers.""" | |
| import pyarrow.parquet as pq | |
| lo, hi = None, None | |
| for f in split.files: | |
| rng = (_file_stats_range(f, field_id, key_type) | |
| or manifest_stats.get(_file_identity(f))) | |
| if rng is not None: | |
| lo = rng[0] if lo is None else min(lo, rng[0]) | |
| hi = rng[1] if hi is None else max(hi, rng[1]) | |
| continue | |
| col = name_for_schema(f.schema_id) | |
| if col is None: | |
| return None, None | |
| path = f.external_path if f.external_path else f.file_path | |
| if path is None or not path.endswith(".parquet"): | |
| return None, None | |
| try: | |
| stream = file_io.new_input_stream(path) | |
| try: | |
| metadata = pq.read_metadata(stream) | |
| finally: | |
| stream.close() | |
| except Exception as e: | |
| # Footer read can fail (e.g. local-cache streams aren't seekable): degrade to | |
| # unknown, but warn -- the fallback would otherwise be silent. | |
| _LOG.warning("range_join: parquet footer read failed for %s (%s); treating " | |
| "its range as unknown", path, e) | |
| return None, None | |
| if _footer_col_type(metadata, col) != key_type: | |
| return None, None | |
| rng = _parquet_col_range(metadata, col) | |
| if rng is None: | |
| return None, None | |
| lo = rng[0] if lo is None else min(lo, rng[0]) | |
| hi = rng[1] if hi is None else max(hi, rng[1]) | |
| if lo is None: | |
| return None, None | |
| return lo, hi | |
| def _range_stats_trusted(table, splits, range_col): | |
| """Whether stored bounds still enclose the read value.""" | |
| from pypaimon.read.query_auth_split import QueryAuthSplit | |
| masked = any(isinstance(s, QueryAuthSplit) and s.auth_result.column_masking | |
| and range_col in s.auth_result.column_masking for s in splits) | |
| # Merge engines may rewrite non-PK values. | |
| return not masked and not ( | |
| table.primary_keys and range_col not in table.primary_keys) | |
| def _plan_ranged_splits(table_id, catalog_options, projection, range_col, partitions=None): | |
| """Plan ``(split, min, max)`` entries on the driver.""" | |
| import os | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pypaimon.common.predicate_builder import PredicateBuilder | |
| from pypaimon.schema.data_types import PyarrowFieldParser | |
| table = get_table(table_id, catalog_options, None, "range_join") | |
| schema_id = table.table_schema.id | |
| snapshot_id = pin_latest_snapshot(table) | |
| if snapshot_id is None: | |
| return [], schema_id | |
| file_io = table.file_io | |
| key_type = PyarrowFieldParser.from_paimon_schema( | |
| table.table_schema.fields).field(range_col).type | |
| # Range key's physical name in a file's schema, by field id (rename/swap safe). | |
| # Cached; current-schema files skip the schema load. | |
| key_field_id = next(f.id for f in table.table_schema.fields if f.name == range_col) | |
| name_cache, cache_lock = {schema_id: range_col}, threading.Lock() | |
| def name_for_schema(sid): | |
| with cache_lock: | |
| if sid in name_cache: | |
| return name_cache[sid] | |
| try: | |
| fields = table.schema_manager.get_schema(sid).fields | |
| name = next((f.name for f in fields if f.id == key_field_id), None) | |
| except Exception: | |
| name = None | |
| with cache_lock: | |
| name_cache[sid] = name | |
| return name | |
| rb = table.new_read_builder() | |
| if partitions: | |
| # Build the partition predicate before projection, so its field list still has | |
| # the partition columns. None means the null partition (is_null, not = None). | |
| pb = rb.new_predicate_builder() | |
| rb = rb.with_partition_filter(PredicateBuilder.and_predicates( | |
| [pb.is_null(c) if v is None else pb.equal(c, v) | |
| for c, v in partitions.items()])) | |
| if projection is not None: | |
| rb = rb.with_projection(projection) | |
| splits = list(rb.new_scan().plan().splits()) | |
| # Merges or masking can move values beyond stored bounds. | |
| if not _range_stats_trusted(table, splits, range_col): | |
| return [(s, None, None) for s in splits], schema_id | |
| manifest_stats = _manifest_stats_by_file( | |
| table, key_field_id, key_type, snapshot_id) | |
| workers = min(16, (os.cpu_count() or 4) * 4, len(splits) or 1) | |
| with ThreadPoolExecutor(max_workers=workers) as pool: | |
| bounds = pool.map( | |
| lambda s: _split_key_range( | |
| s, name_for_schema, key_field_id, key_type, file_io, manifest_stats), | |
| splits) | |
| return [(s, lo, hi) for s, (lo, hi) in zip(splits, bounds)], schema_id | |
| def _cut_points(ranged_sides, num_ranges): | |
| """Pick ``num_ranges - 1`` cut values from row-count-weighted file boundaries.""" | |
| points = [] | |
| for ranged in ranged_sides: | |
| for split, lo, hi in ranged: | |
| if lo is None: | |
| continue | |
| rows = sum(f.row_count for f in split.files) | |
| points.append((lo, rows / 2.0)) | |
| points.append((hi, rows / 2.0)) | |
| if not points: | |
| return [] | |
| points.sort(key=lambda p: p[0]) | |
| total = sum(w for _, w in points) | |
| cuts, acc, k = [], 0.0, 1 | |
| for value, weight in points: | |
| acc += weight | |
| if k >= num_ranges: | |
| break | |
| if acc >= total * k / num_ranges: | |
| if not cuts or value > cuts[-1]: # strictly increasing | |
| cuts.append(value) | |
| k += 1 | |
| return cuts | |
| def _split_rows(split): | |
| return sum(f.row_count for f in split.files) | |
| def _split_bytes(split): | |
| # Re-read cost is bytes, not rows: a wide-row split is cheap by rows, costly by I/O. | |
| return sum(f.file_size for f in split.files) | |
| def _total_reads(l_ranged, r_ranged, ranges): | |
| """Bytes physically read = each split's bytes times the ranges it overlaps (each range | |
| reads the whole split and clips). Counts unknown-stats and wide known splits alike.""" | |
| reads = 0 | |
| for ranged in (l_ranged, r_ranged): | |
| for split, lo, hi in ranged: | |
| spans = sum(1 for r_lo, r_hi in ranges if _overlaps(lo, hi, r_lo, r_hi)) | |
| reads += _split_bytes(split) * spans | |
| return reads | |
| def _bounded_ranges(l_ranged, r_ranged, num_ranges): | |
| """Cut into ``num_ranges`` ranges, halving until total re-read <= _REREAD_BUDGET full | |
| scans of bytes, so poorly clustered input can't cost far more than one scan.""" | |
| total_bytes = sum(_split_bytes(s) | |
| for ranged in (l_ranged, r_ranged) for s, _, _ in ranged) | |
| budget = _REREAD_BUDGET * max(1, total_bytes) | |
| while True: | |
| ranges = _ranges_from_cuts(_cut_points((l_ranged, r_ranged), num_ranges)) | |
| if num_ranges <= 1 or _total_reads(l_ranged, r_ranged, ranges) <= budget: | |
| return ranges | |
| num_ranges = max(1, num_ranges // 2) | |
| def _ranges_from_cuts(cuts): | |
| # Half-open [lo, hi); None = unbounded end. | |
| bounds = [None] + cuts + [None] | |
| return [(bounds[i], bounds[i + 1]) for i in range(len(bounds) - 1)] | |
| def _overlaps(lo, hi, r_lo, r_hi): | |
| if lo is None: # unknown split range: belongs to every range | |
| return True | |
| return (r_lo is None or hi >= r_lo) and (r_hi is None or lo < r_hi) | |
| def _restrict_to_range(arrow_table, col, lo, hi): | |
| """Keep rows with ``lo <= col < hi``. Null keys are always dropped (an inner | |
| join never matches them), which also keeps the result independent of num_ranges.""" | |
| import pyarrow.compute as pc | |
| mask = pc.is_valid(arrow_table[col]) | |
| if lo is not None: | |
| mask = pc.and_(mask, pc.greater_equal(arrow_table[col], lo)) | |
| if hi is not None: | |
| mask = pc.and_(mask, pc.less(arrow_table[col], hi)) | |
| return arrow_table.filter(mask) | |
| def range_join( | |
| left: str, | |
| right: str, | |
| catalog_options: Dict[str, str], | |
| *, | |
| on: Optional[OnSpec] = None, | |
| left_on: Optional[OnSpec] = None, | |
| right_on: Optional[OnSpec] = None, | |
| num_ranges: Optional[int] = None, | |
| left_projection: Optional[List[str]] = None, | |
| right_projection: Optional[List[str]] = None, | |
| left_partitions: Optional[Dict[str, Any]] = None, | |
| right_partitions: Optional[Dict[str, Any]] = None, | |
| join_type: str = "inner", | |
| ray_remote_args: Optional[Dict[str, Any]] = None, | |
| ) -> "ray.data.Dataset": | |
| """Join two tables clustered by the first join key with no global shuffle. | |
| ``on`` when both sides use the same column names, or ``left_on``/``right_on`` | |
| when they differ (positionally paired). The first pair is the range key used to | |
| cut the key space. ``left_partitions``/``right_partitions`` ({column: value} dicts | |
| on partition columns) prune each side to those partitions first. Sides must not | |
| share column names other than ``on`` keys. Returns a ``ray.data.Dataset``. | |
| """ | |
| import ray | |
| if not hasattr(ray.data, "from_arrow_refs"): | |
| raise RuntimeError( | |
| "range_join needs a Ray version with ray.data.from_arrow_refs; " | |
| f"installed ray is {ray.__version__}.") | |
| if (on is None) == (left_on is None and right_on is None): | |
| raise ValueError("range_join requires exactly one of on= or left_on=/right_on=.") | |
| if on is not None: | |
| lkeys = rkeys = norm_on(on) | |
| else: | |
| if left_on is None or right_on is None: | |
| raise ValueError("range_join requires both left_on= and right_on=.") | |
| lkeys, rkeys = norm_on(left_on), norm_on(right_on) | |
| if len(lkeys) != len(rkeys) or not lkeys: | |
| raise ValueError( | |
| f"range_join join keys must pair up non-empty; got left_on={lkeys}, right_on={rkeys}.") | |
| if join_type != "inner": | |
| # Outer joins would need every unmatched row emitted exactly once across | |
| # ranges plus null-key handling; only inner is supported for now. | |
| raise ValueError(f"range_join currently supports only join_type='inner'; got {join_type!r}.") | |
| ltable = get_table(left, catalog_options, None, "range_join") | |
| rtable = get_table(right, catalog_options, None, "range_join") | |
| # Partition filters must name partition columns, else they'd be silently ignored. | |
| for name, tbl, parts in (("left_partitions", ltable, left_partitions), | |
| ("right_partitions", rtable, right_partitions)): | |
| bad = sorted(set(parts) - set(tbl.partition_keys)) if parts else [] | |
| if bad: | |
| raise ValueError( | |
| f"range_join {name} keys {bad} are not partition columns; " | |
| f"partition columns are {list(tbl.partition_keys)}.") | |
| missing = [c for c in lkeys if c not in ltable.field_dict] \ | |
| + [c for c in rkeys if c not in rtable.field_dict] | |
| if missing: | |
| raise ValueError(f"range_join keys not found in table schema: {missing}.") | |
| type_mismatch = [ | |
| (lc, rc, key_type(ltable, lc), key_type(rtable, rc)) | |
| for lc, rc in zip(lkeys, rkeys) | |
| if key_type(ltable, lc) != key_type(rtable, rc) | |
| ] | |
| if type_mismatch: | |
| raise ValueError( | |
| "range_join key columns must have the same type on both sides; " | |
| f"mismatched (left, right, left type, right type): {type_mismatch}.") | |
| # Reject unsupported key types up front (not inside a worker as ArrowInvalid). Every | |
| # join key must be hashable; nested (ARRAY<>/MAP<>/ROW<>/...) and VARIANT are not. | |
| for c in lkeys: | |
| t = key_type(ltable, c).upper() | |
| if "<" in t or t.startswith("VARIANT"): | |
| raise ValueError( | |
| f"range_join join key {c!r} must not be a nested/complex type; got {t}.") | |
| # The range key (first pair) additionally must be range-partitionable. | |
| range_key_type = key_type(ltable, lkeys[0]).upper() | |
| reason = None | |
| if range_key_type.startswith(("FLOAT", "DOUBLE")): | |
| # NaN falls out of every range while the hash join still matches it -> drops rows. | |
| reason = "FLOAT/DOUBLE" | |
| elif "LOCAL TIME ZONE" in range_key_type or "TIMESTAMP_LTZ" in range_key_type: | |
| # Footer stats decode to naive datetimes; a tz-aware column can't compare to them. | |
| reason = "TIMESTAMP WITH LOCAL TIME ZONE" | |
| if reason: | |
| raise ValueError( | |
| f"range_join range key {lkeys[0]!r} must not be {reason}; " | |
| "use an integer/string/date/timestamp key.") | |
| # The join keys must survive projection, or the local join has no key. | |
| if left_projection is not None and not set(lkeys) <= set(left_projection): | |
| raise ValueError( | |
| f"left_projection must include the join keys {lkeys}; got {left_projection}.") | |
| if right_projection is not None and not set(rkeys) <= set(right_projection): | |
| raise ValueError( | |
| f"right_projection must include the join keys {rkeys}; got {right_projection}.") | |
| # pyarrow drops the right keys (coalesced into the left), so the output keeps the LEFT | |
| # key names. A right non-key column sharing a left column name collides -> reject it. | |
| lcols = left_projection if left_projection is not None else ltable.field_names | |
| rcols = right_projection if right_projection is not None else rtable.field_names | |
| collisions = sorted(set(lcols) & (set(rcols) - set(rkeys))) | |
| if collisions: | |
| raise ValueError( | |
| f"range_join output columns collide: {collisions}. The output keeps the left " | |
| "key names and the right non-key columns; project or rename the overlap away.") | |
| l_range_col, r_range_col = lkeys[0], rkeys[0] | |
| l_ranged, l_schema_id = _plan_ranged_splits( | |
| left, catalog_options, left_projection, l_range_col, left_partitions) | |
| r_ranged, r_schema_id = _plan_ranged_splits( | |
| right, catalog_options, right_projection, r_range_col, right_partitions) | |
| def _empty(): | |
| empty = read_splits( | |
| left, catalog_options, left_projection, [], l_schema_id, "range_join").join( | |
| read_splits(right, catalog_options, right_projection, [], r_schema_id, "range_join"), | |
| keys=lkeys, right_keys=rkeys, join_type=join_type) | |
| return ray.data.from_arrow(empty) | |
| if not l_ranged or not r_ranged: # inner join: one empty side, empty result | |
| return _empty() | |
| if num_ranges is None: | |
| num_ranges = max(len(l_ranged), len(r_ranged)) | |
| elif not isinstance(num_ranges, int) or num_ranges < 1: | |
| raise ValueError(f"num_ranges must be an int >= 1; got {num_ranges!r}.") | |
| num_ranges = max(1, min(_MAX_RANGES, num_ranges)) # cap tasks even when explicit | |
| # Reduce ranges until total re-read stays bounded -- see _bounded_ranges. | |
| ranges = _bounded_ranges(l_ranged, r_ranged, num_ranges) | |
| def _join_range(left_splits, right_splits, lo, hi): | |
| # No predicate pushdown: the range key may be schema-evolved (e.g. a file stored | |
| # as INT read as STRING), which the reader can't compare against a new-type bound. | |
| # The in-memory clip below does the exact, evolution-safe filtering. | |
| lt = _restrict_to_range( | |
| read_splits(left, catalog_options, left_projection, left_splits, | |
| l_schema_id, "range_join"), | |
| l_range_col, lo, hi) | |
| rt = _restrict_to_range( | |
| read_splits(right, catalog_options, right_projection, right_splits, | |
| r_schema_id, "range_join"), | |
| r_range_col, lo, hi) | |
| return lt.join(rt, keys=lkeys, right_keys=rkeys, join_type=join_type) | |
| # ``@ray.remote()`` (empty parens) is rejected by Ray, so wrap conditionally. | |
| remote_fn = ray.remote(**ray_remote_args)(_join_range) if ray_remote_args else ray.remote(_join_range) | |
| refs = [] | |
| for r_lo, r_hi in ranges: | |
| ls = [s for s, lo, hi in l_ranged if _overlaps(lo, hi, r_lo, r_hi)] | |
| rs = [s for s, lo, hi in r_ranged if _overlaps(lo, hi, r_lo, r_hi)] | |
| if not ls or not rs: # inner join: a one-sided range can't match | |
| continue | |
| refs.append(remote_fn.remote(ls, rs, r_lo, r_hi)) | |
| if not refs: | |
| return _empty() | |
| # Keep each range's result as a distributed object ref -- never pulled into the driver. | |
| return ray.data.from_arrow_refs(refs) |