import logging import math from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Set, Tuple, Type from ray.data._internal.arrow_block import ArrowBlockAccessor from ray.data._internal.arrow_ops.transform_pyarrow import ( MIN_PYARROW_VERSION_RUN_END_ENCODED_TYPES, MIN_PYARROW_VERSION_VIEW_TYPES, ) from ray.data._internal.execution.interfaces import PhysicalOperator from ray.data._internal.execution.operators.hash_shuffle import ( HashShufflingOperatorBase, ShuffleAggregation, _combine, ) from ray.data._internal.execution.operators.shuffle_operators.shuffle_tasks import ( ReduceFn, ) from ray.data._internal.logical.operators import JoinType from ray.data._internal.util import GiB, MiB from ray.data._internal.utils.arrow_utils import get_pyarrow_version from ray.data._internal.utils.transform_pyarrow import _is_pa_extension_type from ray.data.block import Block from ray.data.context import DataContext if TYPE_CHECKING: import pyarrow as pa @dataclass(frozen=True) class _DatasetPreprocessingResult: """Result of join preprocessing containing split tables. Separates tables into supported (directly joinable) and unsupported (requires indexing) column projections. """ supported_projection: "pa.Table" unsupported_projection: "pa.Table" _JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP = { JoinType.INNER: "inner", JoinType.LEFT_OUTER: "left outer", JoinType.RIGHT_OUTER: "right outer", JoinType.FULL_OUTER: "full outer", JoinType.LEFT_SEMI: "left semi", JoinType.RIGHT_SEMI: "right semi", JoinType.LEFT_ANTI: "left anti", JoinType.RIGHT_ANTI: "right anti", } logger = logging.getLogger(__name__) class JoiningAggregation(ShuffleAggregation): """Stateless aggregation for distributed joining of 2 sequences. This implementation performs hash-based distributed joining by: - Accumulating identical keys from both sequences into the same partition - Performing join on individual partitions independently For actual joining, Pyarrow native joining functionality is utilised. """ def __init__( self, *, join_type: JoinType, left_key_col_names: Tuple[str, ...], right_key_col_names: Tuple[str, ...], left_columns_suffix: Optional[str] = None, right_columns_suffix: Optional[str] = None, data_context: DataContext, ): assert ( len(left_key_col_names) > 0 ), "At least 1 column to join on has to be provided" assert len(right_key_col_names) == len( left_key_col_names ), "Number of columns for both left and right join operands has to match" assert join_type in _JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP, ( f"Join type is not currently supported (got: {join_type}; " # noqa: C416 f"supported: {[jt for jt in JoinType]})" # noqa: C416 ) self._left_key_col_names: Tuple[str, ...] = left_key_col_names self._right_key_col_names: Tuple[str, ...] = right_key_col_names self._join_type: JoinType = join_type self._left_columns_suffix: Optional[str] = left_columns_suffix self._right_columns_suffix: Optional[str] = right_columns_suffix def finalize(self, partition_shards_map: Dict[int, List[Block]]) -> Iterator[Block]: """Performs join on blocks from left (seq 0) and right (seq 1) sequences.""" assert ( len(partition_shards_map) == 2 ), f"Two input-sequences are expected (got {len(partition_shards_map)})" left_partition_shards = partition_shards_map[0] right_partition_shards = partition_shards_map[1] left_table = _combine(left_partition_shards) right_table = _combine(right_partition_shards) yield join_tables( left_table, right_table, join_type=self._join_type, left_key_col_names=self._left_key_col_names, right_key_col_names=self._right_key_col_names, left_columns_suffix=self._left_columns_suffix, right_columns_suffix=self._right_columns_suffix, ) def _make_join_reduce_fn( *, join_type: JoinType, left_key_col_names: Tuple[str, ...], right_key_col_names: Tuple[str, ...], left_columns_suffix: Optional[str] = None, right_columns_suffix: Optional[str] = None, left_schema: Optional[Any] = None, right_schema: Optional[Any] = None, ) -> ReduceFn: """Build a V2-shuffle reduce fn that joins two co-partitioned inputs.""" import pyarrow as pa def _side_table(tables: List[Block], schema: Optional[Any]) -> Optional["pa.Table"]: if tables: return _combine(tables) if isinstance(schema, pa.Schema): return schema.empty_table() return None def _reduce( partition_id: int, tables_by_input: List[List[Block]] ) -> Iterator[Block]: assert ( len(tables_by_input) == 2 ), f"Join reduce expects two inputs (got {len(tables_by_input)})" left_table = _side_table(tables_by_input[0], left_schema) right_table = _side_table(tables_by_input[1], right_schema) if left_table is None or right_table is None: # TODO(you-cheng): A whole input side is empty AND its schema can't be inferred # (0 blocks + un-inferable schema, e.g. a map_batches side), so # _side_table returns None and we skip the partition. This silently # drops the preserved side's rows for preserving joins, left_outer/ # full_outer and left_anti/right_anti. return yield join_tables( left_table, right_table, join_type=join_type, left_key_col_names=left_key_col_names, right_key_col_names=right_key_col_names, left_columns_suffix=left_columns_suffix, right_columns_suffix=right_columns_suffix, ) return _reduce def join_tables( left_table: "pa.Table", right_table: "pa.Table", *, join_type: JoinType, left_key_col_names: Tuple[str, ...], right_key_col_names: Tuple[str, ...], left_columns_suffix: Optional[str] = None, right_columns_suffix: Optional[str] = None, ) -> "pa.Table": """Apply preprocess -> ``pa.Table.join`` -> postprocess to two input tables. Shared between the physical executor (``JoiningAggregation.finalize``) and plan-time schema inference (``Join.infer_schema``), which calls this with empty tables built from the input schemas. Plan-time and runtime schemas therefore agree by construction. """ left_on = list(left_key_col_names) right_on = list(right_key_col_names) # Eagerly validate suffix conflicts so callers get a clear error instead # of the opaque PyArrow schema-merge error ('Field X exists 2 times'). # Skip for semi/anti joins: only one side's columns appear in the result, # so overlapping non-key names between left and right are harmless. if join_type not in ( JoinType.LEFT_SEMI, JoinType.LEFT_ANTI, JoinType.RIGHT_SEMI, JoinType.RIGHT_ANTI, ): left_cols = set(left_table.schema.names) # PyArrow drops right key columns from output (coalescing them into # the left keys), so only right non-key columns can collide with # left columns. Subtracting only right_on (not left_on) correctly # handles asymmetric key names (left_on != right_on). right_output_cols = set(right_table.schema.names) - set(right_on) collisions = left_cols & right_output_cols if left_columns_suffix is None and right_columns_suffix is None and collisions: raise ValueError( "Left and right columns suffixes cannot be both None " f"(overlapping columns: {sorted(collisions)})" ) # Preprocess: split unsupported columns and add index columns if needed preprocess_result_l, preprocess_result_r = _preprocess( left_table, right_table, left_on, right_on, join_type ) # Perform the join on supported columns arrow_join_type = _JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP[join_type] supported = preprocess_result_l.supported_projection.join( preprocess_result_r.supported_projection, join_type=arrow_join_type, keys=left_on, right_keys=right_on, left_suffix=left_columns_suffix, right_suffix=right_columns_suffix, ) # Add back unsupported columns return _postprocess( supported, preprocess_result_l.unsupported_projection, preprocess_result_r.unsupported_projection, ) def _preprocess( left_table: "pa.Table", right_table: "pa.Table", left_on: List[str], right_on: List[str], join_type: JoinType, ) -> Tuple[_DatasetPreprocessingResult, _DatasetPreprocessingResult]: """Split inputs into supported/unsupported columns and add indices.""" supported_l, unsupported_l = _split_unsupported_columns(left_table) supported_r, unsupported_r = _split_unsupported_columns(right_table) # Handle joins on unsupported columns conflicting_columns: Set[str] = set(unsupported_l.column_names) & set(left_on) if conflicting_columns: raise ValueError( f"Cannot join on columns with unjoinable types. " f"Left join key columns {conflicting_columns} have unjoinable types " f"(map, union, list, struct, etc.) which cannot be used for join operations." ) conflicting_columns: Set[str] = set(unsupported_r.column_names) & set(right_on) if conflicting_columns: raise ValueError( f"Cannot join on columns with unjoinable types. " f"Right join key columns {conflicting_columns} have unjoinable types " f"(map, union, list, struct, etc.) which cannot be used for join operations." ) # Index if we have unsupported columns should_index_l = _should_index_side("left", supported_l, unsupported_l, join_type) should_index_r = _should_index_side("right", supported_r, unsupported_r, join_type) # Add index columns for back-referencing if we have unsupported columns if should_index_l: supported_l = _append_index_column( table=supported_l, col_name=_index_name("left") ) if should_index_r: supported_r = _append_index_column( table=supported_r, col_name=_index_name("right") ) left = _DatasetPreprocessingResult( supported_projection=supported_l, unsupported_projection=unsupported_l, ) right = _DatasetPreprocessingResult( supported_projection=supported_r, unsupported_projection=unsupported_r, ) return left, right def _postprocess( supported: "pa.Table", unsupported_l: "pa.Table", unsupported_r: "pa.Table", ) -> "pa.Table": """Re-attach unsupported columns to the joined table via the index column.""" should_index_l = _index_name("left") in supported.schema.names should_index_r = _index_name("right") in supported.schema.names # Add back unsupported columns (join type logic is in should_index_* variables) if should_index_l: supported = _add_back_unsupported_columns( joined_table=supported, unsupported_table=unsupported_l, index_col_name=_index_name("left"), ) if should_index_r: supported = _add_back_unsupported_columns( joined_table=supported, unsupported_table=unsupported_r, index_col_name=_index_name("right"), ) return supported def _index_name(suffix: str) -> str: return f"__rd_index_level_{suffix}__" def _should_index_side( side: str, supported_table: "pa.Table", unsupported_table: "pa.Table", join_type: JoinType, ) -> bool: """ Determine whether to create an index column for a given side of the join. A column is "supported" if it is "joinable", and "unsupported" otherwise. A supported_table is a table with only "supported" columns. Index columns are needed when we have both supported and unsupported columns in a table, and that table's columns will appear in the final result. Args: side: "left" or "right" to indicate which side of the join supported_table: Table containing ONLY joinable columns unsupported_table: Table containing ONLY unjoinable columns join_type: The join type, used to decide whether this side appears in the result (semi/anti joins drop one side). Returns: True if an index column should be created for this side """ # Must have both supported and unsupported columns to need indexing. # We cannot rely on row_count because it can return a non-zero row count # for an empty-schema. if len(supported_table.schema) == 0 or len(unsupported_table.schema) == 0: return False # For semi/anti joins, only index the side that appears in the result if side == "left": # Left side appears in result for all joins except right_semi/right_anti return join_type not in [JoinType.RIGHT_SEMI, JoinType.RIGHT_ANTI] else: # side == "right" # Right side appears in result for all joins except left_semi/left_anti return join_type not in [JoinType.LEFT_SEMI, JoinType.LEFT_ANTI] def _split_unsupported_columns( table: "pa.Table", ) -> Tuple["pa.Table", "pa.Table"]: """ Split a PyArrow table into two tables based on column joinability. Separates columns into supported types and unsupported types that cannot be directly joined on but should be preserved in results. Args: table: Input PyArrow table to split Returns: Tuple of (supported_table, unsupported_table) where: - supported_table contains columns with primitive/joinable types - unsupported_table contains columns with complex/unjoinable types """ supported, unsupported = [], [] for idx in range(len(table.columns)): col: "pa.ChunkedArray" = table.column(idx) col_type: "pa.DataType" = col.type if _is_pa_extension_type(col_type) or JoinOperator._is_pa_join_not_supported( col_type ): unsupported.append(idx) else: supported.append(idx) return table.select(supported), table.select(unsupported) def _add_back_unsupported_columns( joined_table: "pa.Table", unsupported_table: "pa.Table", index_col_name: str, ) -> "pa.Table": # Extract the index column array and drop the column from the joined table i = joined_table.schema.get_field_index(index_col_name) indices = joined_table.column(i) joined_table = joined_table.remove_column(i) # Project the unsupported columns using the indices and combine with joined table projected = ArrowBlockAccessor(unsupported_table).take(indices) return ArrowBlockAccessor(joined_table).hstack(projected) def _append_index_column(table: "pa.Table", col_name: str) -> "pa.Table": import numpy as np import pyarrow as pa index_col = pa.array(np.arange(table.num_rows)) return table.append_column(col_name, index_col) class JoinOperator(HashShufflingOperatorBase): def __init__( self, data_context: DataContext, left_input_op: PhysicalOperator, right_input_op: PhysicalOperator, left_key_columns: Tuple[str], right_key_columns: Tuple[str], join_type: JoinType, *, num_partitions: Optional[int] = None, left_columns_suffix: Optional[str] = None, right_columns_suffix: Optional[str] = None, partition_size_hint: Optional[int] = None, aggregator_ray_remote_args_override: Optional[Dict[str, Any]] = None, shuffle_aggregation_type: Optional[Type[ShuffleAggregation]] = None, ): # Use new stateless JoiningAggregation factory def _create_joining_aggregation() -> JoiningAggregation: if shuffle_aggregation_type is not None: if not issubclass(shuffle_aggregation_type, ShuffleAggregation): raise TypeError( f"shuffle_aggregation_type must be a subclass of {ShuffleAggregation}, " f"got {shuffle_aggregation_type}" ) aggregation_class = shuffle_aggregation_type or JoiningAggregation return aggregation_class( join_type=join_type, left_key_col_names=left_key_columns, right_key_col_names=right_key_columns, left_columns_suffix=left_columns_suffix, right_columns_suffix=right_columns_suffix, data_context=data_context, ) super().__init__( name_factory=( lambda num_partitions: f"Join(num_partitions={num_partitions})" ), input_ops=[left_input_op, right_input_op], data_context=data_context, key_columns=[left_key_columns, right_key_columns], num_input_seqs=2, num_partitions=num_partitions, partition_size_hint=partition_size_hint, partition_aggregation_factory=_create_joining_aggregation, aggregator_ray_remote_args_override=aggregator_ray_remote_args_override, shuffle_progress_bar_name="Shuffle", finalize_progress_bar_name="Join", ) @staticmethod def _is_pa_join_not_supported(dtype: "pa.DataType") -> bool: """ The latest pyarrow versions do not support joins where the tables contain the following types below (lists, structs, maps, unions, extension types, etc.) Args: dtype: The input type of column. Returns: True if the type cannot be present (non join-key) during joins. False if the type can be present. """ import pyarrow as pa pyarrow_version = get_pyarrow_version() is_v12 = pyarrow_version >= MIN_PYARROW_VERSION_RUN_END_ENCODED_TYPES is_v16 = pyarrow_version >= MIN_PYARROW_VERSION_VIEW_TYPES return ( pa.types.is_map(dtype) or pa.types.is_union(dtype) or pa.types.is_list(dtype) or pa.types.is_struct(dtype) or pa.types.is_null(dtype) or pa.types.is_large_list(dtype) or pa.types.is_fixed_size_list(dtype) or (is_v12 and pa.types.is_run_end_encoded(dtype)) or ( is_v16 and ( pa.types.is_binary_view(dtype) or pa.types.is_string_view(dtype) or pa.types.is_list_view(dtype) ) ) ) def _get_operator_num_cpus_override(self) -> float: return self.data_context.join_operator_actor_num_cpus_override @classmethod def _estimate_aggregator_memory_allocation( cls, *, num_aggregators: int, num_partitions: int, estimated_dataset_bytes: int, ) -> int: partition_byte_size_estimate = math.ceil( estimated_dataset_bytes / num_partitions ) # Estimate of object store memory required to accommodate all partitions # handled by a single aggregator aggregator_shuffle_object_store_memory_required: int = math.ceil( estimated_dataset_bytes / num_aggregators ) # Estimate of memory required to perform actual (in-memory) join # operation (inclusive of 50% overhead allocated for Pyarrow join # implementation) # # NOTE: # - 2x due to budgeted 100% overhead of Arrow's in-memory join join_memory_required: int = math.ceil(partition_byte_size_estimate * 2) # Estimate of memory required to accommodate single partition as an output # (inside Object Store) # # NOTE: x2 due to 2 sequences involved in joins output_object_store_memory_required: int = partition_byte_size_estimate aggregator_total_memory_required: int = ( # Inputs (object store) aggregator_shuffle_object_store_memory_required + # Join (heap) join_memory_required + # Output (object store) output_object_store_memory_required ) logger.info( f"Estimated memory requirement for joining aggregator " f"(partitions={num_partitions}, " f"aggregators={num_aggregators}, " f"dataset (estimate)={estimated_dataset_bytes / GiB:.1f}GiB): " f"shuffle={aggregator_shuffle_object_store_memory_required / MiB:.1f}MiB, " f"joining={join_memory_required / MiB:.1f}MiB, " f"output={output_object_store_memory_required / MiB:.1f}MiB, " f"total={aggregator_total_memory_required / MiB:.1f}MiB, " ) return aggregator_total_memory_required