Repository: dask-contrib/dask-expr
Pull request #906 (2024-02-29 09:21:10)

Implement more efficient GroupBy.mean
This actually made a bigger difference than expected, it cuts query 17 from 5 minutes to 3:20, shuffling twice is just expensive

I will reactor the layout a little bit in a follow-up, but didn't want to touch too many things at once


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index 2d697a7..1b5807a 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -2428,9 +2428,7 @@ class DataFrame(FrameBase):
 
             elif isinstance(v, (Scalar, Series)):
                 if isinstance(v, Series):
-                    if not expr.are_co_aligned(
-                        result.expr, v.expr, allow_broadcast=False
-                    ):
+                    if not expr.are_co_aligned(result.expr, v.expr):
                         if len(args) > 0:
                             result = expr.Assign(result, *args)
                             args = []
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index 83682f1..cc529d1 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -2836,16 +2836,24 @@ def is_broadcastable(dfs, s):
     )
 
 
-def non_blockwise_ancestors(expr):
-    """Traverse through tree to find ancestors that are not blockwise or are IO"""
+def are_co_aligned(*exprs):
+    """Do inputs come from the same parents, modulo blockwise?"""
+
     from dask_expr._cumulative import CumulativeAggregations
     from dask_expr._reductions import Reduction
 
-    stack = [expr]
+    seen = set()
+    # Scalars can always be broadcasted
+    stack = [e for e in exprs if e.ndim > 0]
+    ancestors = []
     while stack:
         e = stack.pop()
+        if e._name in seen:
+            continue
+        seen.add(e._name)
+
         if isinstance(e, IO):
-            yield e
+            ancestors.append(e)
         elif e.ndim == 0:
             # Scalars are valid ancestors that are always broadcastable,
             # so don't walk through them
@@ -2853,22 +2861,16 @@ def non_blockwise_ancestors(expr):
         elif isinstance(e, (Blockwise, CumulativeAggregations, Reduction)):
             # TODO: Capture this in inheritance logic
             dependencies = e.dependencies()
-            stack.extend([expr for expr in dependencies])
+            stack.extend(dependencies)
         elif isinstance(e, _DelayedExpr):
             continue
         else:
-            yield e
-
+            ancestors.append(e)
 
-def are_co_aligned(*exprs, allow_broadcast=True):
-    """Do inputs come from the same parents, modulo blockwise?"""
-    # Scalars can always be broadcasted
-    exprs = [e for e in exprs if e.ndim > 0]
-    ancestors = [set(non_blockwise_ancestors(e)) for e in exprs]
     unique_ancestors = {
         # Account for column projection within IO expressions
-        _tokenize_partial(item, ["columns", "_series"])
-        for item in flatten(ancestors, container=set)
+        _tokenize_partial(item, ["columns", "_series", "_dataset_info_cache"])
+        for item in ancestors
     }
     # Don't check divisions or npartitions at all
     return len(unique_ancestors) <= 1
diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py
index 23dfddd..5754542 100644
--- a/dask_expr/_groupby.py
+++ b/dask_expr/_groupby.py
@@ -31,6 +31,7 @@ from dask.dataframe.groupby import (
     _groupby_aggregate_spec,
     _groupby_apply_funcs,
     _groupby_get_group,
+    _groupby_raise_unaligned,
     _groupby_slice_apply,
     _groupby_slice_shift,
     _groupby_slice_transform,
@@ -769,26 +770,62 @@ class Std(SingleAggregation):
         )
 
 
-class Mean(SingleAggregation):
+def _mean_chunk(df, *by, observed=None, dropna=None):
+    if is_series_like(df):
+        df = df.to_frame()
+
+    g = _groupby_raise_unaligned(df, by=by, observed=observed, dropna=dropna)
+    x = g.sum(numeric_only=True)
+    n = g[x.columns].count().rename(columns=lambda c: c + "-count")
+    return concat([x, n], axis=1)
+
+
+def _mean_combine(g, levels, sort=False):
+    return g.groupby(level=levels, sort=sort).sum()
+
+
+def _mean_agg(g, levels, sort=False, observed=False, dropna=True):
+    result = g.groupby(level=levels, sort=sort, observed=observed, dropna=dropna).sum()
+    s = result[result.columns[: len(result.columns) // 2]]
+    c = result[result.columns[len(result.columns) // 2 :]]
+    c.columns = s.columns
+    return s / c
+
+
+class Mean(GroupByReduction):
+    _parameters = SingleAggregation._parameters
+    _defaults = SingleAggregation._defaults
+    reduction_aggregate = staticmethod(_mean_agg)
+    reduction_combine = staticmethod(_mean_combine)
+    chunk = staticmethod(_mean_chunk)
+
     @functools.cached_property
-    def _meta(self):
-        return self._lower()._meta
+    def aggregate_kwargs(self):
+        return {
+            "levels": self.levels,
+            "sort": self.sort,
+            "observed": self.observed,
+            "dropna": self.dropna,
+        }
 
-    def _lower(self):
-        s = Sum(*self.operands)
-        # Drop chunk/aggregate_kwargs for count
-        c = Count(
-            *[
-                self.operand(param)
-                if param not in ("chunk_kwargs", "aggregate_kwargs")
-                else {}
-                for param in self._parameters
-            ],
-            *self.by,
-        )
-        if is_dataframe_like(s._meta):
-            c = c[s.columns]
-        return s / c
+    @functools.cached_property
+    def chunk_kwargs(self):
+        return {"observed": self.observed, "dropna": self.dropna}
+
+    @functools.cached_property
+    def combine_kwargs(self):
+        return {"levels": self.levels}
+
+    def _divisions(self):
+        if self.sort:
+            return (None, None)
+        split_out = self.split_out
+        if split_out is True:
+            split_out = self.frame.npartitions
+        return (None,) * (split_out + 1)
+
+    def _simplify_up(self, parent, dependents):
+        return groupby_projection(self, parent, dependents)
 
 
 def nunique_df_combine(dfs, *args, **kwargs):
@@ -1629,13 +1666,28 @@ class GroupBy:
         return len(set(post_group_columns) - set(numerics.columns)) == 0
 
     @derived_from(pd.core.groupby.GroupBy)
-    def mean(self, numeric_only=False, **kwargs):
+    def mean(self, numeric_only=False, split_out=1, **kwargs):
         if not numeric_only and not self._all_numeric():
             raise NotImplementedError(
                 "'numeric_only=False' is not implemented in Dask."
             )
         numeric_kwargs = self._numeric_only_kwargs(numeric_only)
-        return self._single_agg(Mean, **kwargs, **numeric_kwargs)
+        result = self._single_agg(Mean, split_out=split_out, **kwargs, **numeric_kwargs)
+        return self._postprocess_series_squeeze(result)
+
+    def _postprocess_series_squeeze(self, result):
+        if (
+            isinstance(self.obj, Series)
+            or is_scalar(self._slice)
+            and self._slice is not None
+        ):
+            if len(result.columns) < 1:
+                raise NotImplementedError(
+                    "Cannot call `SeriesGroupBy.var` or `SeriesGroupBy.mean` on the key "
+                    "column. Please use `aggregate` if you really need to do this."
+                )
+            result = result[result.columns[0]]
+        return result
 
     @derived_from(pd.core.groupby.GroupBy)
     def min(self, numeric_only=False, **kwargs):
@@ -1803,18 +1855,7 @@ class GroupBy:
                 *self.by,
             )
         )
-        if (
-            isinstance(self.obj, Series)
-            or is_scalar(self._slice)
-            and self._slice is not None
-        ):
-            if len(result.columns) < 1:
-                raise NotImplementedError(
-                    "Cannot call `SeriesGroupBy.var` on the key column. "
-                    "Please use `aggregate` if you really need to do this."
-                )
-            result = result[result.columns[0]]
-        return result
+        return self._postprocess_series_squeeze(result)
 
     @derived_from(pd.core.groupby.GroupBy)
     def std(
@@ -1843,18 +1884,7 @@ class GroupBy:
                 *self.by,
             )
         )
-        if (
-            isinstance(self.obj, Series)
-            or is_scalar(self._slice)
-            and self._slice is not None
-        ):
-            if len(result.columns) < 1:
-                raise NotImplementedError(
-                    "Cannot call `SeriesGroupBy.std` on the key column. "
-                    "Please use `aggregate` if you really need to do this."
-                )
-            result = result[result.columns[0]]
-        return result
+        return self._postprocess_series_squeeze(result)
 
     @_aggregate_docstring(based_on="pd.core.groupby.DataFrameGroupBy.agg")
     def aggregate(
diff --git a/dask_expr/_merge.py b/dask_expr/_merge.py
index 3e47e10..6431b92 100644
--- a/dask_expr/_merge.py
+++ b/dask_expr/_merge.py
@@ -642,6 +642,7 @@ class HashJoinP2P(Merge, PartitionsFiltered):
                 self.suffixes,
                 self.left_index,
                 self.right_index,
+                self.indicator,
             )
         return dsk
 
diff --git a/dask_expr/io/parquet.py b/dask_expr/io/parquet.py
index a3e0a21..2845bad 100644
--- a/dask_expr/io/parquet.py
+++ b/dask_expr/io/parquet.py
@@ -48,7 +48,6 @@ from dask_expr._expr import (
     Literal,
     Or,
     Projection,
-    are_co_aligned,
     determine_column_projection,
 )
 from dask_expr._reductions import Len
@@ -617,7 +616,7 @@ class ReadParquetPyarrowFS(ReadParquet):
 
     @cached_property
     def normalized_path(self):
-        return pa_fs.FileSystem.from_uri(self.path)[1]
+        return _normalize_and_strip_protocol(self.path)
 
     @cached_property
     def fs(self):
@@ -639,21 +638,24 @@ class ReadParquetPyarrowFS(ReadParquet):
         dataset_info = {}
 
         path_normalized = self.normalized_path
-        # At this point we will post a couple of listbucket operations which
-        # includes the same data as a HEAD request.
-        # The information included here (see pyarrow FileInfo) are size, type,
-        # path and modified since timestamps
-        # This isn't free but realtively cheap (200-300ms or less for ~1k files)
-        finfo = self.fs.get_file_info(path_normalized)
-        if finfo.type == pa.fs.FileType.Directory:
+        # We'll first treat the path as if it was a directory since this is the
+        # most common case. Only if this fails, we'll treat it as a file. This
+        # way, the happy path performs one remote request instead of two if we
+        # were to check the type of the path first.
+        try:
+            # At this point we will post a listbucket request which includes the
+            # same data as a HEAD request. The information included here (see
+            # pyarrow FileInfo) are size, type, path and modified since
+            # timestamps This isn't free but realtively cheap (200-300ms or less
+            # for ~1k files)
             dataset_selector = pa_fs.FileSelector(path_normalized, recursive=True)
             all_files = [
                 finfo
                 for finfo in self.fs.get_file_info(dataset_selector)
                 if finfo.type == pa.fs.FileType.File
             ]
-        else:
-            all_files = [finfo]
+        except (NotADirectoryError, FileNotFoundError):
+            all_files = [self.fs.get_file_info(path_normalized)]
         # TODO: At this point we could verify if we're dealing with a very
         # inhomogeneous datasets already without reading any further data
 
@@ -1173,15 +1175,19 @@ class _DNF:
     def extract_pq_filters(cls, pq_expr: ReadParquet, predicate_expr: Expr) -> _DNF:
         _filters = None
         if isinstance(predicate_expr, (LE, GE, LT, GT, EQ, NE)):
-            if are_co_aligned(pq_expr, predicate_expr.left) and not isinstance(
-                predicate_expr.right, Expr
+            if (
+                not isinstance(predicate_expr.right, Expr)
+                and isinstance(predicate_expr.left, Projection)
+                and predicate_expr.left.frame._name == pq_expr._name
             ):
                 op = predicate_expr._operator_repr
                 column = predicate_expr.left.columns[0]
                 value = predicate_expr.right
                 _filters = (column, op, value)
-            elif are_co_aligned(pq_expr, predicate_expr.right) and not isinstance(
-                predicate_expr.left, Expr
+            elif (
+                not isinstance(predicate_expr.left, Expr)
+                and isinstance(predicate_expr.left, Projection)
+                and predicate_expr.left.frame._name == pq_expr._name
             ):
                 # Simple dict to make sure field comes first in filter
                 flip = {LE: GE, LT: GT, GE: LE, GT: LT}
@@ -1321,3 +1327,14 @@ def _read_partition_stats_group(parts, fs, columns=None):
 
     # Helper function used by _extract_statistics
     return [_read_partition_stats(part, fs, columns=columns) for part in parts]
+
+
+def _normalize_and_strip_protocol(path):
+    protocol_separators = ["://", "::"]
+    for sep in protocol_separators:
+        split = path.split(sep, 1)
+        if len(split) > 1:
+            path = split[1]
+            break
+    path = path.rstrip("/")
+    return path


Test patch:
diff --git a/dask_expr/io/tests/test_parquet.py b/dask_expr/io/tests/test_parquet.py
index 791a79d..9ba0aaf 100644
--- a/dask_expr/io/tests/test_parquet.py
+++ b/dask_expr/io/tests/test_parquet.py
@@ -6,9 +6,9 @@ from dask.dataframe.utils import assert_eq
 from pyarrow import fs
 
 from dask_expr import from_pandas, read_parquet
-from dask_expr._expr import Lengths, Literal
+from dask_expr._expr import Filter, Lengths, Literal
 from dask_expr._reductions import Len
-from dask_expr.io import ReadParquet
+from dask_expr.io import FusedIO, ReadParquet
 
 
 def _make_file(dir, df=None):
@@ -121,3 +121,92 @@ def test_partition_pruning(tmpdir):
         # FIXME ?
         check_names=False,
     )
+
+
+def test_predicate_pushdown(tmpdir):
+    original = pd.DataFrame(
+        {
+            "a": [1, 2, 3, 4, 5] * 10,
+            "b": [0, 1, 2, 3, 4] * 10,
+            "c": range(50),
+            "d": [6, 7] * 25,
+            "e": [8, 9] * 25,
+        }
+    )
+    fn = _make_file(tmpdir, df=original)
+    df = read_parquet(fn, filesystem="arrow")
+    assert_eq(df, original)
+    x = df[df.a == 5][df.c > 20]["b"]
+    y = x.optimize(fuse=False)
+    assert isinstance(y.expr.frame, FusedIO)
+    assert ("a", "==", 5) in y.expr.frame.operands[0].operand("filters")[0]
+    assert ("c", ">", 20) in y.expr.frame.operands[0].operand("filters")[0]
+    assert list(y.columns) == ["b"]
+
+    # Check computed result
+    y_result = y.compute()
+    assert y_result.name == "b"
+    assert len(y_result) == 6
+    assert (y_result == 4).all()
+
+    # Don't push down if replace is in there
+    x = df[df.replace(5, 50).a == 5]["b"]
+    y = x.optimize(fuse=False)
+    assert isinstance(y.expr, Filter)
+    assert len(y.compute()) == 0
+
+
+def test_predicate_pushdown_compound(tmpdir):
+    pdf = pd.DataFrame(
+        {
+            "a": [1, 2, 3, 4, 5] * 10,
+            "b": [0, 1, 2, 3, 4] * 10,
+            "c": range(50),
+            "d": [6, 7] * 25,
+            "e": [8, 9] * 25,
+        }
+    )
+    fn = _make_file(tmpdir, df=pdf)
+    df = read_parquet(fn, filesystem="arrow")
+
+    # Test AND
+    x = df[(df.a == 5) & (df.c > 20)]["b"]
+    y = x.optimize(fuse=False)
+    assert isinstance(y.expr.frame, FusedIO)
+    assert {("c", ">", 20), ("a", "==", 5)} == set(y.expr.frame.operands[0].filters[0])
+    assert_eq(
+        y,
+        pdf[(pdf.a == 5) & (pdf.c > 20)]["b"],
+        check_index=False,
+    )
+
+    # Test OR
+    x = df[(df.a == 5) | (df.c > 20)]
+    x = x[x.b != 0]["b"]
+    y = x.optimize(fuse=False)
+    assert isinstance(y.expr.frame, FusedIO)
+    filters = [
+        set(y.expr.frame.operands[0].filters[0]),
+        set(y.expr.frame.operands[0].filters[1]),
+    ]
+    assert {("c", ">", 20), ("b", "!=", 0)} in filters
+    assert {("a", "==", 5), ("b", "!=", 0)} in filters
+    expect = pdf[(pdf.a == 5) | (pdf.c > 20)]
+    expect = expect[expect.b != 0]["b"]
+    assert_eq(
+        y,
+        expect,
+        check_index=False,
+    )
+
+    # Test OR and AND
+    x = df[((df.a == 5) | (df.c > 20)) & (df.b != 0)]["b"]
+    z = x.optimize(fuse=False)
+    assert isinstance(z.expr.frame, FusedIO)
+    filters = [
+        set(z.expr.frame.operands[0].filters[0]),
+        set(z.expr.frame.operands[0].filters[1]),
+    ]
+    assert {("c", ">", 20), ("b", "!=", 0)} in filters
+    assert {("a", "==", 5), ("b", "!=", 0)} in filters
+    assert_eq(y, z)
diff --git a/dask_expr/tests/test_distributed.py b/dask_expr/tests/test_distributed.py
index 2cf9aa9..af9b321 100644
--- a/dask_expr/tests/test_distributed.py
+++ b/dask_expr/tests/test_distributed.py
@@ -354,3 +354,22 @@ async def test_future_in_map_partitions(c, s, a, b):
     result = await result
     expected = pd.DataFrame({"a": [4951, 4952, 4953, 4954]})
     pd.testing.assert_frame_equal(result, expected)
+
+
+@gen_cluster(client=True)
+async def test_merge_indicator(c, s, a, b):
+    data = {
+        "id": ["101-a", "102-a", "103-a"],
+        "test": ["val101a", "val102a", "val103a"],
+    }
+    pdf = pd.DataFrame(data)
+    df = 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),
+    )