[python] Add range_join: shuffle-free Ray join for range-clustered tables

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",

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(

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