Repository: ibis-project/ibis

Issue:
perf(dask): avoid calling compute for `ops.Covariance`, `ops.ArgMin` and `ops.ArgMax`
> dask has built in covariance and correlation methods, why are they not sufficient here/why do we need to compute?
>
> IMO anything that needs to compute like this shouldn't warn, we should just not support it unless we can do it cleanly.

_Originally posted by @jcrist in https://github.com/ibis-project/ibis/pull/8005#discussion_r1473095470_
            


Discussion:
See my comments [here](https://github.com/ibis-project/ibis/pull/8005#pullrequestreview-1854096466) and [here](https://github.com/ibis-project/ibis/pull/8005#issuecomment-1922235967). I'd strongly argue that we should either fix this before releasing, or drop support for these operations until they can be properly handled. IMO anywhere where we're explicitly calling `.compute()` outside of returning results to the end user is a bug in the dask backend.
@jcrist you mean before the 9.0 release, yes?
Yes.
Going to move this off the 9.0 milestone. I think Dask performance isn't a blocker.

Resolving patch:
diff --git a/ibis/backends/dask/executor.py b/ibis/backends/dask/executor.py
--- a/ibis/backends/dask/executor.py
+++ b/ibis/backends/dask/executor.py
@@ -27,7 +27,7 @@
     PandasWindowFunction,
     plan,
 )
-from ibis.common.exceptions import UnboundExpressionError
+from ibis.common.exceptions import UnboundExpressionError, UnsupportedOperationError
 from ibis.formats.pandas import PandasData, PandasType
 from ibis.util import gen_name
 
@@ -51,6 +51,14 @@ def limit_df(
     return df[df[col].between(offset, offset + n - 1)]
 
 
+def argminmax_chunk(df, keycol, valcol, method):
+    return df.iloc[getattr(df[keycol], method)()]
+
+
+def argminmax_aggregate(df, keycol, valcol, method):
+    return df[valcol].iloc[getattr(df[keycol], method)()]
+
+
 class DaskExecutor(PandasExecutor, DaskUtils):
     name = "dask"
     kernels = dask_kernels
@@ -184,66 +192,55 @@ def mapper(df):
 
     @classmethod
     def visit(cls, op: ops.ArgMin | ops.ArgMax, arg, key, where):
-        # TODO(kszucs): raise a warning about triggering compute()?
-        if isinstance(op, ops.ArgMin):
-            func = lambda x: x.idxmin()
-        else:
-            func = lambda x: x.idxmax()
-
-        if where is None:
+        method = "argmin" if isinstance(op, ops.ArgMin) else "argmax"
 
-            def agg(df):
-                indices = func(df[key.name])
-                if isinstance(indices, (dd.Series, dd.core.Scalar)):
-                    # to support both aggregating within a group and globally
-                    indices = indices.compute()
-                return df[arg.name].loc[indices]
-        else:
+        def agg(df):
+            if where is not None:
+                df = df.where(df[where.name])
 
-            def agg(df):
-                mask = df[where.name]
-                filtered = df[mask]
-                indices = func(filtered[key.name])
-                if isinstance(indices, (dd.Series, dd.core.Scalar)):
-                    # to support both aggregating within a group and globally
-                    indices = indices.compute()
-                return filtered[arg.name].loc[indices]
+            if isinstance(df, dd.DataFrame):
+                return df.reduction(
+                    chunk=argminmax_chunk,
+                    combine=argminmax_chunk,
+                    aggregate=argminmax_aggregate,
+                    meta=op.dtype.to_pandas(),
+                    token=method,
+                    keycol=key.name,
+                    valcol=arg.name,
+                    method=method,
+                )
+            else:
+                return argminmax_aggregate(df, key.name, arg.name, method)
 
         return agg
 
     @classmethod
     def visit(cls, op: ops.Correlation, left, right, where, how):
-        if where is None:
+        if how == "pop":
+            raise UnsupportedOperationError(
+                "Dask doesn't support `corr` with `how='pop'`"
+            )
 
-            def agg(df):
-                return df[left.name].corr(df[right.name])
-        else:
+        def agg(df):
+            if where is not None:
+                df = df.where(df[where.name])
 
-            def agg(df):
-                mask = df[where.name]
-                lhs = df[left.name][mask].compute()
-                rhs = df[right.name][mask].compute()
-                return lhs.corr(rhs)
+            return df[left.name].corr(df[right.name])
 
         return agg
 
     @classmethod
     def visit(cls, op: ops.Covariance, left, right, where, how):
-        # TODO(kszucs): raise a warning about triggering compute()?
-        ddof = {"pop": 0, "sample": 1}[how]
-        if where is None:
+        if how == "pop":
+            raise UnsupportedOperationError(
+                "Dask doesn't support `cov` with `how='pop'`"
+            )
 
-            def agg(df):
-                lhs = df[left.name].compute()
-                rhs = df[right.name].compute()
-                return lhs.cov(rhs, ddof=ddof)
-        else:
+        def agg(df):
+            if where is not None:
+                df = df.where(df[where.name])
 
-            def agg(df):
-                mask = df[where.name]
-                lhs = df[left.name][mask].compute()
-                rhs = df[right.name][mask].compute()
-                return lhs.cov(rhs, ddof=ddof)
+            return df[left.name].cov(df[right.name])
 
         return agg
 


Test patch:
diff --git a/ibis/backends/tests/test_aggregation.py b/ibis/backends/tests/test_aggregation.py
--- a/ibis/backends/tests/test_aggregation.py
+++ b/ibis/backends/tests/test_aggregation.py
@@ -874,6 +874,11 @@ def test_quantile(
             lambda t, where: t.G[where].cov(t.RBI[where], ddof=0),
             id="covar_pop",
             marks=[
+                pytest.mark.notyet(
+                    ["dask"],
+                    reason="dask doesn't support `cov(ddof=0)` yet",
+                    raises=com.UnsupportedOperationError,
+                ),
                 pytest.mark.notimpl(
                     ["polars", "druid"],
                     raises=com.OperationNotDefinedError,
@@ -914,6 +919,11 @@ def test_quantile(
             lambda t, where: t.G[where].corr(t.RBI[where]),
             id="corr_pop",
             marks=[
+                pytest.mark.notyet(
+                    ["dask"],
+                    raises=com.UnsupportedOperationError,
+                    reason="dask doesn't support `corr(ddof=0)` yet",
+                ),
                 pytest.mark.notimpl(
                     ["druid"],
                     raises=com.OperationNotDefinedError,
@@ -978,6 +988,11 @@ def test_quantile(
             lambda t, where: (t.G[where] > 34.0).cov(t.G[where] <= 34.0, ddof=0),
             id="covar_pop_bool",
             marks=[
+                pytest.mark.notyet(
+                    ["dask"],
+                    raises=com.UnsupportedOperationError,
+                    reason="dask doesn't support `cov(ddof=0)` yet",
+                ),
                 pytest.mark.notimpl(
                     ["polars", "druid"],
                     raises=com.OperationNotDefinedError,
@@ -1002,6 +1017,11 @@ def test_quantile(
             lambda t, where: (t.G[where] > 34.0).corr(t.G[where] <= 34.0),
             id="corr_pop_bool",
             marks=[
+                pytest.mark.notyet(
+                    ["dask"],
+                    raises=com.UnsupportedOperationError,
+                    reason="dask doesn't support `corr(ddof=0)` yet",
+                ),
                 pytest.mark.notimpl(
                     ["druid"],
                     raises=com.OperationNotDefinedError,