Repository: dask-contrib/dask-expr
Pull request #931 (2024-03-05 09:54:17)

Support token for map_partitions
closes https://github.com/dask/dask/issues/10975


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index f40d7a6..db926c1 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -400,6 +400,25 @@ class FrameBase(DaskMethodsMixin):
             other = other.copy()
         return new_collection(self.expr.__getitem__(other))
 
+    def __repr__(self):
+        data = self._repr_data().to_string(max_rows=5)
+        _str_fmt = """Dask {klass} Structure:
+{data}
+Dask Name: {name}, {layers}
+Expr={expr}"""
+        if len(self.columns) == 0:
+            data = data.partition("\n")[-1].replace("Index", "Divisions")
+            _str_fmt = f"Empty {_str_fmt}"
+        layers = len({k for (k, *_) in self.dask.keys()})
+
+        return _str_fmt.format(
+            klass=self.__class__.__name__,
+            data=data,
+            name=key_split(self._name),
+            layers=maybe_pluralize(layers, "graph layer"),
+            expr=self.expr,
+        )
+
     def __bool__(self):
         raise ValueError(
             f"The truth value of a {self.__class__.__name__} is ambiguous. "
@@ -2414,6 +2433,9 @@ class DataFrame(FrameBase):
     def _ipython_key_completions_(self):
         return methods.tolist(self.columns)
 
+    def _repr_html_(self):
+        return self.to_html()
+
     @derived_from(pd.DataFrame)
     def assign(self, **pairs):
         result = self
@@ -2715,9 +2737,6 @@ class DataFrame(FrameBase):
             warnings.warn(meta_warning(meta))
         return new_collection(expr.Map(self, arg=func, na_action=na_action, meta=meta))
 
-    def __repr__(self):
-        return f"<dask_expr.expr.DataFrame: expr={self.expr}>"
-
     @derived_from(pd.DataFrame)
     def nlargest(self, n=5, columns=None, split_every=None):
         return new_collection(
@@ -3802,9 +3821,6 @@ class Series(FrameBase):
         axis = self._validate_axis(axis)
         return new_collection(self.expr.clip(lower, upper, axis))
 
-    def __repr__(self):
-        return f"<dask_expr.expr.Series: expr={self.expr}>"
-
     @derived_from(pd.Series)
     def to_frame(self, name=no_default):
         return new_collection(expr.ToFrame(self, name=name))
@@ -5239,6 +5255,7 @@ def map_partitions(
         clear_divisions,
         align_dataframes,
         parent_meta,
+        kwargs.pop("token", None),
         kwargs,
         *args[1:],
     )
@@ -5301,6 +5318,7 @@ def map_overlap(
                     transform_divisions,
                     clear_divisions,
                     align_dataframes,
+                    kwargs.pop("token", None),
                     kwargs,
                     *args,
                 )
@@ -5316,6 +5334,7 @@ def map_overlap(
         transform_divisions,
         clear_divisions,
         align_dataframes,
+        kwargs.pop("token", None),
         kwargs,
         *args,
     )
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index fe0c3b6..c17bef7 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -550,13 +550,33 @@ class MapPartitions(Blockwise):
         "clear_divisions",
         "align_dataframes",
         "parent_meta",
+        "token",
         "kwargs",
     ]
-    _defaults = {"kwargs": None, "align_dataframes": True, "parent_meta": None}
+    _defaults = {
+        "kwargs": None,
+        "align_dataframes": True,
+        "parent_meta": None,
+        "token": None,
+    }
+
+    @functools.cached_property
+    def token(self):
+        if "token" in self._parameters:
+            return self.operand("token")
+        return None
 
     def __str__(self):
         return f"MapPartitions({funcname(self.func)})"
 
+    @functools.cached_property
+    def _name(self):
+        if self.token is not None:
+            head = self.token
+        else:
+            head = funcname(self.func).lower()
+        return head + "-" + _tokenize_deterministic(*self.operands)
+
     def _broadcast_dep(self, dep: Expr):
         # Always broadcast single-partition dependencies in MapPartitions
         return dep.npartitions == 1
@@ -715,6 +735,7 @@ class MapOverlapAlign(Expr):
         "transform_divisions",
         "clear_divisions",
         "align_dataframes",
+        "token",
         "kwargs",
     ]
     _defaults = {
@@ -724,6 +745,7 @@ class MapOverlapAlign(Expr):
         "kwargs": None,
         "clear_divisions": False,
         "align_dataframes": False,
+        "token": None,
     }
 
     @functools.cached_property
@@ -759,6 +781,7 @@ class MapOverlapAlign(Expr):
             self.transform_divisions,
             self.clear_divisions,
             self.align_dataframes,
+            self.token,
             self.kwargs,
             *args[1:],
         )
@@ -775,6 +798,7 @@ class MapOverlap(MapPartitions):
         "transform_divisions",
         "clear_divisions",
         "align_dataframes",
+        "token",
         "kwargs",
     ]
     _defaults = {
@@ -784,6 +808,7 @@ class MapOverlap(MapPartitions):
         "kwargs": None,
         "clear_divisions": False,
         "align_dataframes": False,
+        "token": None,
     }
 
     @functools.cached_property
@@ -843,6 +868,7 @@ class MapOverlap(MapPartitions):
             self.clear_divisions,
             self.align_dataframes,
             None,
+            self.token,
             self._kwargs,
             *self.args[1:],
         )
diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py
index 76b5452..4695ace 100644
--- a/dask_expr/_groupby.py
+++ b/dask_expr/_groupby.py
@@ -1314,6 +1314,7 @@ class GroupByCumulative(Expr, GroupByBase):
             False,
             True,
             None,
+            None,
             {"chunk": self.chunk, "columns": columns, **dropna, **self.numeric_only},
             *self.by,
         )
@@ -1342,6 +1343,7 @@ class GroupByCumulative(Expr, GroupByBase):
             False,
             True,
             None,
+            None,
             {"chunk": M.last, "columns": columns, **dropna},
             *by,
         )
diff --git a/dask_expr/_merge_asof.py b/dask_expr/_merge_asof.py
index 03d7249..18ce6c7 100644
--- a/dask_expr/_merge_asof.py
+++ b/dask_expr/_merge_asof.py
@@ -130,6 +130,7 @@ class MergeAsof(Merge):
                 False,
                 True,
                 None,
+                None,
                 {"left_index": True, "right_index": True},
                 right,
             )


Test patch:
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index cc25f3f..dcd41ae 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -945,7 +945,7 @@ def test_isin_repr(df):
     result = df.isin([1, 2])
     # This was raising previously
     result = result.__repr__()
-    assert "<dask_expr.expr.DataFrame: expr=Isin(frame=df, values=" in result
+    assert "Expr=Isin(frame=df, values=" in result
 
 
 def test_round(pdf):
diff --git a/dask_expr/tests/test_concat.py b/dask_expr/tests/test_concat.py
index 6e55922..e75f63e 100644
--- a/dask_expr/tests/test_concat.py
+++ b/dask_expr/tests/test_concat.py
@@ -22,8 +22,8 @@ def df(pdf):
 
 def test_concat_str(df):
     result = str(concat([df, df], join="inner"))
-    expected = "<dask_expr.expr.DataFrame: expr=Concat(frames=[df, df], join=inner)>"
-    assert result == expected
+    expected = "Expr=Concat(frames=[df, df], join=inner)"
+    assert expected in result
 
 
 def test_concat(pdf, df):
diff --git a/dask_expr/tests/test_format.py b/dask_expr/tests/test_format.py
index 926d315..9b18ae1 100644
--- a/dask_expr/tests/test_format.py
+++ b/dask_expr/tests/test_format.py
@@ -95,3 +95,88 @@ def test_series_format():
     H      ..."""
     )
     assert ds.to_string() == exp
+
+
+def test_series_repr():
+    s = pd.Series([1, 2, 3, 4, 5, 6, 7, 8], index=list("ABCDEFGH"))
+    ds = from_pandas(s, 3)
+
+    exp = dedent(
+        """\
+        Empty Dask Series Structure:
+        A    int64
+        D      ...
+        G      ...
+        H      ...
+        Dask Name: frompandas, 1 graph layer
+        Expr=df"""
+    )
+    assert repr(ds) == exp
+
+
+def test_df_repr():
+    df = pd.DataFrame({"col1": range(10), "col2": map(float, range(10))})
+    ddf = from_pandas(df, 3)
+
+    exp = dedent(
+        """\
+        Dask DataFrame Structure:
+                        col1     col2
+        npartitions=3                
+        0              int64  float64
+        4                ...      ...
+        7                ...      ...
+        9                ...      ...
+        Dask Name: frompandas, 1 graph layer
+        Expr=df"""
+    )
+    assert repr(ddf) == exp
+
+
+def test_df_to_html():
+    df = pd.DataFrame({"col1": range(10), "col2": map(float, range(10))})
+    ddf = from_pandas(df, 3)
+
+    exp = dedent(
+        """\
+        <div><strong>Dask DataFrame Structure:</strong></div>
+        <table border="1" class="dataframe">
+          <thead>
+            <tr style="text-align: right;">
+              <th></th>
+              <th>col1</th>
+              <th>col2</th>
+            </tr>
+            <tr>
+              <th>npartitions=3</th>
+              <th></th>
+              <th></th>
+            </tr>
+          </thead>
+          <tbody>
+            <tr>
+              <th>0</th>
+              <td>int64</td>
+              <td>float64</td>
+            </tr>
+            <tr>
+              <th>4</th>
+              <td>...</td>
+              <td>...</td>
+            </tr>
+            <tr>
+              <th>7</th>
+              <td>...</td>
+              <td>...</td>
+            </tr>
+            <tr>
+              <th>9</th>
+              <td>...</td>
+              <td>...</td>
+            </tr>
+          </tbody>
+        </table>
+        <div>Dask Name: frompandas, 1 graph layer</div>"""
+    )
+    assert ddf.to_html() == exp
+    assert ddf._repr_html_() == exp  # for jupyter
diff --git a/dask_expr/tests/test_map_partitions_overlap.py b/dask_expr/tests/test_map_partitions_overlap.py
index 8ddc48f..8f607da 100644
--- a/dask_expr/tests/test_map_partitions_overlap.py
+++ b/dask_expr/tests/test_map_partitions_overlap.py
@@ -337,3 +337,16 @@ def test_map_partitions_propagates_index_metadata():
     res = ddf.map_partitions(lambda df: df.rename_axis("newindex"))
     sol = df.rename_axis("newindex")
     assert_eq(res, sol)
+
+
+def test_token_given(df, pdf):
+    def my_func(df):
+        return df + 1
+
+    result = df.map_partitions(my_func, token="token_given")
+    assert result._name.split("-")[0] == "token_given"
+    assert_eq(result, pdf + 1)
+
+    result = df.map_partitions(my_func)
+    assert result._name.split("-")[0] == "my_func"
+    assert_eq(result, pdf + 1)