Add Dask integration for a join operation (rapidsai#517)

In order to leverage the recent `AllGather` work within Dask-based execution in multi-GPU Polars, we need a new integration API in RapidsMPF.

Although this new API **could** simply specialize in broadcast joins, it seems reasonable to introduce a general join integration that handles any realistic combination of shuffling/broadcasting the child tables. Even if we are not performing a broadcast join, it may be beneficial to extract left- and right-hand partitions within the same task that is performing the local join. This way, we don't leave un-spillable data on the worker between the shuffle-extraction task(s) and the corresponding join task.

This PR adds a general integration API for join operations. It includes some of the pieces that will be needed for broadcast-based joins. However, it only implements the hash-join code path for now (to limit the size of the diff). Note that this PR does not implement "single-worker" execution either. Both broadcast joins and single-worker execution will be addressed in follow-up work.

Authors:
  - Richard (Rick) Zamora (https://github.com/rjzamora)

Approvers:
  - Tom Augspurger (https://github.com/TomAugspurger)

URL: rapidsai#517

Author: rjzamora

python/rapidsmpf/rapidsmpf/examples/dask.py

@@ -26,12 +26,17 @@
 from rapidsmpf.utils.cudf import cudf_to_pylibcudf_table
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
     from typing import Any
 
     import dask_cudf
 
     import cudf
 
+    from rapidsmpf.integrations.core import (
+        BCastJoinInfo,
+        ShufflerIntegration,
+    )
     from rapidsmpf.shuffler import Shuffler
 
 
@@ -153,6 +158,30 @@ def extract_partition(
         )
 
 
+def _get_cluster_kind(
+    cluster_kind: Literal["distributed", "single", "auto"],
+) -> Literal["distributed", "single"]:
+    """Validate and return the kind of cluster to use."""
+    if cluster_kind not in ("distributed", "single", "auto"):
+        raise ValueError(
+            f"Expected one of 'distributed', 'single', or 'auto'. Got {cluster_kind}"
+        )
+
+    if cluster_kind == "auto":
+        try:
+            from distributed import get_client
+
+            get_client()
+        except (ImportError, ValueError):
+            # Failed to import distributed/dask-cuda or find a Dask client.
+            # Use single shuffle instead.
+            cluster_kind = "single"
+        else:
+            cluster_kind = "distributed"
+
+    return cluster_kind
+
+
 def dask_cudf_shuffle(
     df: dask_cudf.DataFrame,
     on: list[str],
@@ -195,10 +224,10 @@ def dask_cudf_shuffle(
     This API is currently intended for demonstration and
     testing purposes only.
     """
-    if cluster_kind not in ("distributed", "single", "auto"):
-        raise ValueError(
-            f"Expected one of 'distributed', 'single', or 'auto'. Got {cluster_kind}"
-        )
+    if (cluster_kind := _get_cluster_kind(cluster_kind)) == "distributed":
+        shuffle = rapidsmpf.integrations.dask.rapidsmpf_shuffle_graph
+    else:
+        shuffle = rapidsmpf.integrations.single.rapidsmpf_shuffle_graph
 
     df0 = df.optimize()
     count_in = df0.npartitions
@@ -218,23 +247,6 @@ def dask_cudf_shuffle(
     else:
         sort_boundary_names = ()
 
-    if cluster_kind == "auto":
-        try:
-            from distributed import get_client
-
-            get_client()
-        except (ImportError, ValueError):
-            # Failed to import distributed/dask-cuda or find a Dask client.
-            # Use single shuffle instead.
-            cluster_kind = "single"
-        else:
-            cluster_kind = "distributed"
-
-    if cluster_kind == "distributed":
-        shuffle = rapidsmpf.integrations.dask.rapidsmpf_shuffle_graph
-    else:
-        shuffle = rapidsmpf.integrations.single.rapidsmpf_shuffle_graph
-
     shuffle_graph_args = (
         name_in,
         name_out,
@@ -269,3 +281,165 @@ def dask_cudf_shuffle(
         )
     else:
         return shuffled
+
+
+class DaskCudfJoinIntegration:
+    """Dask-cuDF protocol for unified join integration."""
+
+    @staticmethod
+    def get_shuffler_integration() -> ShufflerIntegration[cudf.DataFrame]:
+        """Return the shuffler integration."""
+        return DaskCudfIntegration()
+
+    @staticmethod
+    def join_partition(
+        left_input: Callable[[int], cudf.DataFrame],
+        right_input: Callable[[int], cudf.DataFrame],
+        bcast_info: BCastJoinInfo | None,
+        options: Any,
+    ) -> cudf.DataFrame:
+        """
+        Produce a joined DataFrame partition.
+
+        Parameters
+        ----------
+        left_input
+            A callable that produces chunks of the left partition.
+            The ``bcast_info.bcast_count`` parameter corresponds
+            to the number of chunks the callable can produce.
+        right_input
+            A callable that produces chunks of the right partition.
+            The ``bcast_info.bcast_count`` parameter corresponds
+            to the number of chunks the callable can produce.
+        bcast_info
+            The broadcast join information.
+            This should be None for a regular hash join.
+        options
+            Additional join options.
+
+        Returns
+        -------
+        A joined DataFrame partition.
+
+        Notes
+        -----
+        This method is used to produce a single joined table chunk.
+        """
+        join_kwargs = {
+            "left_on": options["left_on"],
+            "right_on": options["right_on"],
+            "how": options["how"],
+        }
+
+        if bcast_info is None:
+            return left_input(0).merge(right_input(0), **join_kwargs)
+        else:  # pragma: no cover
+            raise NotImplementedError("Broadcast join not implemented.")
+
+
+def dask_cudf_join(
+    left: dask_cudf.DataFrame,
+    right: dask_cudf.DataFrame,
+    left_on: list[str],
+    right_on: list[str],
+    *,
+    how: Literal["inner", "left", "right"] = "inner",
+    left_pre_shuffled: bool = False,
+    right_pre_shuffled: bool = False,
+    cluster_kind: Literal["distributed", "single", "auto"] = "auto",
+    config_options: Options = Options(),
+) -> dask_cudf.DataFrame:
+    """
+    Join two Dask-cuDF DataFrames with RapidsMPF.
+
+    Parameters
+    ----------
+    left
+        Left Dask-cuDF DataFrame.
+    right
+        Right Dask-cuDF DataFrame.
+    left_on
+        Left column names to join on.
+    right_on
+        Right column names to join on.
+    how
+        The type of join to perform.
+        Options are ``{'inner', 'left', 'right'}``.
+    left_pre_shuffled
+        Whether the left collection is already shuffled.
+    right_pre_shuffled
+        Whether the right collection is already shuffled.
+    cluster_kind
+        What kind of Dask cluster to use. Available
+        options are ``{'distributed', 'single', 'auto'}``.
+        If 'auto' (the default), 'distributed' will be
+        used if a global Dask client is found.
+        Note: Only ``'distributed'`` is supported for now.
+    config_options
+        RapidsMPF configuration options.
+
+    Returns
+    -------
+    A joined Dask-cuDF DataFrame collection.
+
+    Notes
+    -----
+    This API is currently intended for demonstration and
+    testing purposes only.
+    """
+    if (cluster_kind := _get_cluster_kind(cluster_kind)) == "distributed":
+        from rapidsmpf.integrations.dask.join import rapidsmpf_join_graph
+    else:  # pragma: no cover
+        # TODO: Support single-worker joins.
+        raise NotImplementedError("Single-worker join not implemented.")
+
+    left0 = left.optimize()
+    right0 = right.optimize()
+    left_partition_count_in = left0.npartitions
+    right_partition_count_in = right0.npartitions
+
+    token = tokenize(left0, right0, left_on, right_on, how)
+    left_name_in = left0._name
+    right_name_in = right0._name
+    name_out = f"unified-join-{token}"
+    graph = rapidsmpf_join_graph(
+        left_name_in,
+        right_name_in,
+        name_out,
+        left_partition_count_in,
+        right_partition_count_in,
+        DaskCudfJoinIntegration(),
+        # Options that may be used for shuffling, broadcasting,
+        # or repartitioning the left side.
+        {
+            "column_names": left0.columns,
+            "on": left_on,
+        },
+        # Options that may be used for shuffling, broadcasting,
+        # or repartitioning the right side.
+        {
+            "column_names": right0.columns,
+            "on": right_on,
+        },
+        # Options that may be used for joining.
+        {
+            "left_on": left_on,
+            "right_on": right_on,
+            "how": how,
+        },
+        left_pre_shuffled=left_pre_shuffled,
+        right_pre_shuffled=right_pre_shuffled,
+        config_options=config_options,
+    )
+    graph.update(left0.dask)
+    graph.update(right0.dask)
+
+    meta = left0.merge(right0, left_on=left_on, right_on=right_on, how=how)._meta
+    count_out = max(left_partition_count_in, right_partition_count_in)
+    return dd.from_graph(
+        graph,
+        meta,
+        (None,) * (count_out + 1),
+        [(name_out, pid) for pid in range(count_out)],
+        "rapidsmpf",
+    )