nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: dask-contrib/dask-expr
Pull request #871 (2024-02-13 12:10:02)

Fix broadcast config when filters are pushed down


Patch:
diff --git a/dask_expr/_merge.py b/dask_expr/_merge.py
index 4091bbc..f1c76be 100644
--- a/dask_expr/_merge.py
+++ b/dask_expr/_merge.py
@@ -316,7 +316,6 @@ class Merge(Expr):
                 shuffle_right_on = right.index._meta.name
                 if shuffle_right_on is None:
                     shuffle_right_on = "_index"
-
             if self.is_broadcast_join:
                 left, right = self._bcast_left, self._bcast_right
 
@@ -445,17 +444,7 @@ class Merge(Expr):
             if new_right is self.right and new_left is self.left:
                 # don't drop the filter
                 return
-            return type(self)(
-                new_left,
-                new_right,
-                how=self.how,
-                left_on=self.left_on,
-                right_on=self.right_on,
-                left_index=self.left_index,
-                right_index=self.right_index,
-                suffixes=self.suffixes,
-                indicator=self.indicator,
-            )
+            return type(self)(new_left, new_right, *self.operands[2:])
         if isinstance(parent, (Projection, Index)):
             # Reorder the column projection to
             # occur before the Merge


Test patch:
diff --git a/dask_expr/tests/test_merge.py b/dask_expr/tests/test_merge.py
index 048ab3a..2a86d3d 100644
--- a/dask_expr/tests/test_merge.py
+++ b/dask_expr/tests/test_merge.py
@@ -5,6 +5,7 @@ import pytest
 
 from dask_expr import Merge, from_pandas, merge, repartition
 from dask_expr._expr import Filter, Projection
+from dask_expr._merge import BroadcastJoin
 from dask_expr._shuffle import Shuffle
 from dask_expr.tests._util import _backend_library, assert_eq
 
@@ -863,6 +864,17 @@ def test_isin_filter_pushdown(how):
     assert_eq(result, expected, check_index=False)
 
 
+def test_merge_filter_pushdown_broadcast():
+    pdf = pd.DataFrame({"a": [1, 2, 3], "b": 1})
+    pdf2 = pd.DataFrame({"c": [1, 2, 3], "b": 1})
+    df1 = from_pandas(pdf, npartitions=2)
+    df2 = from_pandas(pdf2, npartitions=2)
+    result = df1.merge(df2, broadcast=True, shuffle_method="tasks")
+    result = result[result.a > 1]
+    result.optimize().pprint()
+    assert len(list(result.optimize().find_operations(BroadcastJoin))) > 0
+
+
 def test_merge_scalar_comparison():
     pdf = pd.DataFrame({"a": [1, 2, 3], "b": 1})
     pdf2 = pd.DataFrame({"c": [1, 2, 3], "b": 1})
Repository: dask-contrib/dask-expr
Pull request #897 (2024-02-26 16:53:49)

Fix meta calculation error in groupby


Patch:
diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py
index 26cf4ba..1fe1557 100644
--- a/dask_expr/_groupby.py
+++ b/dask_expr/_groupby.py
@@ -393,18 +393,24 @@ class HolisticGroupbyAggregation(GroupbyAggregationBase):
 
     This class always calculates the aggregates by first collecting all the data for
     the groups and then aggregating at once.
+
+    We are always shuffling, so we will never call combine
     """
 
+    @functools.cached_property
+    def _meta(self):
+        meta = self._meta_chunk
+        aggregate = self.aggregate or (lambda x: x)
+        aggregate_kwargs = self.aggregate_kwargs
+        meta = aggregate([meta], **aggregate_kwargs)
+        return make_meta(meta)
+
     chunk = staticmethod(_non_agg_chunk)
 
     @property
     def should_shuffle(self):
         return True
 
-    @classmethod
-    def combine(cls, inputs, **kwargs):
-        return _groupby_aggregate_spec(_concat(inputs), **kwargs)
-
     @classmethod
     def aggregate(cls, inputs, **kwargs):
         return _groupby_aggregate_spec(_concat(inputs), **kwargs)
@@ -418,15 +424,6 @@ class HolisticGroupbyAggregation(GroupbyAggregationBase):
             **_as_dict("dropna", self.dropna),
         }
 
-    @property
-    def combine_kwargs(self) -> dict:
-        return {
-            "spec": self.arg,
-            "levels": _determine_levels(self.by),
-            **_as_dict("observed", self.observed),
-            **_as_dict("dropna", self.dropna),
-        }
-
     @property
     def aggregate_kwargs(self) -> dict:
         return {


Test patch:
diff --git a/dask_expr/tests/test_groupby.py b/dask_expr/tests/test_groupby.py
index 4d1a088..c544391 100644
--- a/dask_expr/tests/test_groupby.py
+++ b/dask_expr/tests/test_groupby.py
@@ -955,3 +955,9 @@ def test_groupby_size_drop_columns(df, pdf):
     result = df.groupby("x").size()
     assert_eq(result, pdf.groupby("x").size())
     assert result.simplify()._name == df[[]].groupby("x").size().simplify()._name
+
+
+def test_groupby_agg_meta_error(df, pdf):
+    result = df.groupby(["x"]).agg({"y": ["median", "std"]})
+    expected = pdf.groupby(["x"]).agg({"y": ["median", "std"]})
+    assert_eq(result, expected)
Repository: dask-contrib/dask-expr
Pull request #925 (2024-03-04 11:43:14)

Avoid culling for SetIndexBlockwise with divisions
closes https://github.com/dask/dask/issues/10974


Patch:
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index 291178b..fe0c3b6 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -2653,8 +2653,10 @@ class Partitions(Expr):
         return (self.frame._name, self.partitions[index])
 
     def _simplify_down(self):
+        from dask_expr import SetIndexBlockwise
+
         if isinstance(self.frame, Blockwise) and not isinstance(
-            self.frame, (BlockwiseIO, Fused)
+            self.frame, (BlockwiseIO, Fused, SetIndexBlockwise)
         ):
             operands = [
                 Partitions(op, self.partitions)


Test patch:
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index a1b9f1d..a0e9440 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -129,6 +129,11 @@ def test_setitem(pdf, df):
     assert_eq(df, pdf)
 
 
+def test_set_index_blockwise_culling(df, pdf):
+    result = df.set_index("x").partitions[0]
+    assert_eq(result, pdf.set_index("x").head(10))
+
+
 @xfail_gpu("https://github.com/rapidsai/cudf/issues/10271")
 def test_explode():
     with dask.config.set({"dataframe.convert-string": False}):
Repository: dask-contrib/dask-expr
Pull request #927 (2024-03-04 14:45:08)

Ensure that columns are copied in projection
closes https://github.com/dask/dask/issues/10957

This seems like a sensible thing to protect against for regular lists


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index 0e8a201..94dedd4 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -396,6 +396,8 @@ class FrameBase(DaskMethodsMixin):
                 return self.loc[other]
         if isinstance(other, np.ndarray) or is_series_like(other):
             other = list(other)
+        elif isinstance(other, list):
+            other = other.copy()
         return new_collection(self.expr.__getitem__(other))
 
     def __bool__(self):


Test patch:
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index a0e9440..cc25f3f 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -119,6 +119,13 @@ def test_info(df, verbose, buf, memory_usage):
         raise NotImplementedError(f"Case not covered for kwargs: {kwargs}")
 
 
+def test_column_projection_modify_list(df, pdf):
+    cols = ["x"]
+    result = df[cols]
+    cols.append("bla")
+    assert_eq(result, pdf[["x"]])
+
+
 def test_setitem(pdf, df):
     pdf = pdf.copy()
     pdf["z"] = pdf.x + pdf.y
Repository: dask-contrib/dask-expr
Pull request #940 (2024-03-06 08:57:16)

Adjust automatic split_out parameter
I tried avoiding fragmentation in the beginning, but I think we were too aggressive in squashing partitions during groupby. 

AB Test: https://github.com/coiled/benchmarks/actions/runs/8164285327

Query 10, 18 and 20 got faster (18 and 20 by a lot), everything else seems mostly noise except 13 where there might be a small slowdown, but not really noticeable.


Patch:
diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py
index 52d7a60..151cacf 100644
--- a/dask_expr/_groupby.py
+++ b/dask_expr/_groupby.py
@@ -80,7 +80,7 @@ def _as_dict(key, value):
 
 
 def _adjust_split_out_for_group_keys(npartitions, by):
-    return math.ceil(npartitions / (100 / (len(by) - 1)))
+    return math.ceil(npartitions / (20 / (len(by) - 1)))
 
 
 class Aggregation:


Test patch:
diff --git a/dask_expr/tests/test_groupby.py b/dask_expr/tests/test_groupby.py
index 1b96a32..62af29b 100644
--- a/dask_expr/tests/test_groupby.py
+++ b/dask_expr/tests/test_groupby.py
@@ -356,12 +356,12 @@ def test_split_out_automatically():
     assert_eq(q, expected)
 
     q = df.groupby(["a", "b"]).sum()
-    assert q.optimize().npartitions == 5
+    assert q.optimize().npartitions == 25
     expected = pdf.groupby(["a", "b"]).sum()
     assert_eq(q, expected)
 
     q = df.groupby(["a", "b", "c"]).sum()
-    assert q.optimize().npartitions == 10
+    assert q.optimize().npartitions == 50
     expected = pdf.groupby(["a", "b", "c"]).sum()
     assert_eq(q, expected)
Repository: dask-contrib/dask-expr
Pull request #947 (2024-03-07 13:53:30)

Warn if annotations are ignored
closes https://github.com/dask-contrib/dask-expr/issues/937


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index 85164f6..5a2c440 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -13,7 +13,7 @@ import dask.dataframe.methods as methods
 import numpy as np
 import pandas as pd
 import pyarrow as pa
-from dask import compute
+from dask import compute, get_annotations
 from dask.array import Array
 from dask.base import DaskMethodsMixin, is_dask_collection, named_schedulers
 from dask.core import flatten
@@ -276,6 +276,7 @@ def _wrap_unary_expr_op(self, op=None):
     return new_collection(getattr(self.expr, op)())
 
 
+_WARN_ANNOTATIONS = True
 #
 # Collection classes
 #
@@ -290,6 +291,12 @@ class FrameBase(DaskMethodsMixin):
     __dask_optimize__ = staticmethod(lambda dsk, keys, **kwargs: dsk)
 
     def __init__(self, expr):
+        global _WARN_ANNOTATIONS
+        if _WARN_ANNOTATIONS and (annot := get_annotations()):
+            _WARN_ANNOTATIONS = False
+            warnings.warn(
+                f"Dask annotations {annot} detected. Annotations will be ignored when using query-planning."
+            )
         self._expr = expr
 
     @property
diff --git a/dask_expr/io/parquet.py b/dask_expr/io/parquet.py
index c766cc8..9a6963e 100644
--- a/dask_expr/io/parquet.py
+++ b/dask_expr/io/parquet.py
@@ -149,6 +149,15 @@ class FragmentWrapper:
         self.unpack()
         return self._fragment
 
+    def __dask_tokenize__(self):
+        return type(self).__name__, normalize_token(
+            (
+                self.fragment,
+                self._fragment_packed,
+                self._file_size,
+            )
+        )
+
     def __reduce__(self):
         self.pack()
         return FragmentWrapper, (None, None, self._fragment_packed)


Test patch:
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index 9506384..66c6df9 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -2489,3 +2489,15 @@ def test_astype_filter_pushdown(df, pdf):
     expected = pdf.astype("float64")
     expected = expected[expected.x > 5.0]
     assert_eq(result, expected)
+
+
+@pytest.mark.filterwarnings("error")
+def test_warn_annotations():
+    from_pandas(pd.DataFrame({"a": [1, 2, 3]}), npartitions=2)
+
+    with pytest.warns(UserWarning, match="annotations.*retries.*ignore"):
+        with dask.annotate(retries=3):
+            from_pandas(pd.DataFrame({"a": [1, 2, 3]}), npartitions=2)
+    # Don't warn a second time
+    with dask.annotate(retries=3):
+        from_pandas(pd.DataFrame({"a": [1, 2, 3]}), npartitions=2)
Repository: dask-contrib/dask-expr
Pull request #963 (2024-03-11 08:29:00)

Fix 'Empty' prefix to non-empty Series repr
https://github.com/dask-contrib/dask-expr/pull/943#discussion_r1518975367


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index 7eb66b9..085dc43 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -418,7 +418,7 @@ class FrameBase(DaskMethodsMixin):
 {data}
 Dask Name: {name}, {n_expr}
 Expr={expr}"""
-        if len(self.columns) == 0:
+        if not isinstance(self, Series) and not len(self.columns):
             data = data.partition("\n")[-1].replace("Index", "Divisions")
             _str_fmt = f"Empty {_str_fmt}"
         n_expr = len({e._name for e in self.expr.walk()})


Test patch:
diff --git a/dask_expr/tests/test_format.py b/dask_expr/tests/test_format.py
index 8cf845b..a5e0a6d 100644
--- a/dask_expr/tests/test_format.py
+++ b/dask_expr/tests/test_format.py
@@ -103,7 +103,8 @@ def test_series_repr():
 
     exp = dedent(
         """\
-        Empty Dask Series Structure:
+        Dask Series Structure:
+        npartitions=3
         A    int64
         D      ...
         G      ...
@@ -113,6 +114,24 @@ def test_series_repr():
     )
     assert repr(ds) == exp
 
+    # Not a cheap way to determine if series is empty
+    # so does not prefix with "Empty" as we do w/ empty DataFrame
+    s = pd.Series([])
+    ds = from_pandas(s, 3)
+
+    exp = dedent(
+        """\
+        Dask Series Structure:
+        npartitions=3
+            string
+               ...
+               ...
+               ...
+        Dask Name: frompandas, 1 expression
+        Expr=df"""
+    )
+    assert repr(ds) == exp
+
 
 def test_df_repr():
     df = pd.DataFrame({"col1": range(10), "col2": map(float, range(10))})
@@ -132,6 +151,21 @@ def test_df_repr():
     )
     assert repr(ddf) == exp
 
+    df = pd.DataFrame()
+    ddf = from_pandas(df, 3)
+
+    exp = dedent(
+        """\
+        Empty Dask DataFrame Structure:
+        npartitions=3
+        0              int64  float64
+        4                ...      ...
+        7                ...      ...
+        9                ...      ...
+        Dask Name: frompandas, 1 expression
+        Expr=df"""
+    )
+
 
 def test_df_to_html():
     df = pd.DataFrame({"col1": range(10), "col2": map(float, range(10))})