Add indicator support to merge (dask#8539)

jrbourbeau · Feb 28, 2024 · f86a031 · f86a031
1 parent 84e7c4d
commit f86a031
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/distributed/shuffle/_merge.py b/distributed/shuffle/_merge.py
@@ -172,6 +172,7 @@ def merge_unpack(
     suffixes: Suffixes,
     left_index: bool,
     right_index: bool,
+    indicator: bool = False,
 ):
     from dask.dataframe.multi import merge_chunk
 
@@ -193,6 +194,7 @@ def merge_unpack(
         suffixes=suffixes,
         left_index=left_index,
         right_index=right_index,
+        indicator=indicator,
     )
 
 
@@ -430,5 +432,6 @@ def _construct_graph(self) -> dict[tuple | str, tuple]:
                 self.suffixes,
                 self.left_index,
                 self.right_index,
+                self.indicator,
             )
         return dsk
diff --git a/distributed/shuffle/tests/test_merge.py b/distributed/shuffle/tests/test_merge.py
@@ -539,3 +539,22 @@ async def test_merge_does_not_deadlock_if_worker_joins(c, s, a):
         result = await result
     expected = pd.merge(pdf1, pdf2, left_on="a", right_on="x")
     assert_eq(result, expected, check_index=False)
+
+
+@gen_cluster(client=True)
+async def test_merge_indicator(c, s, a, b):
+    data = {
+        "id": [1, 2, 3],
+        "test": [4, 5, 6],
+    }
+    pdf = pd.DataFrame(data)
+    df = dd.from_pandas(pdf, npartitions=2)
+    result = df.merge(df, on="id", how="outer", indicator=True)
+    x = c.compute(result)
+    x = await x
+    expected = pdf.merge(pdf, on="id", how="outer", indicator=True)
+
+    pd.testing.assert_frame_equal(
+        x.sort_values("id", ignore_index=True),
+        expected.sort_values("id", ignore_index=True),
+    )