Repository: scverse/scanpy

Issue:
Return type of aggregate is inconsistent
```python
import scanpy as sc
adata = sc.datasets.pbmc68k_reduced()
sc.get.aggregate(adata, by="louvain", func="mean")
```

```
AnnData object with n_obs × n_vars = 11 × 765
    obs: 'louvain'
    var: 'n_counts', 'means', 'dispersions', 'dispersions_norm', 'highly_variable'
    layers: 'mean'
```

```python
sc.get.aggregate(adata.obsm["X_umap"], by=adata.obs["louvain"].array, func="mean")
```

```
{'mean': array([[ -6.18019123,  -6.12846152],
        [ -3.10995685,   8.4991954 ],
        [  6.30307056,  -2.15245383],
        [ -4.72268065,  -3.24033642],
        [-11.94002487,  -5.39480163],
        [ -1.39242794,   6.6239316 ],
        [  4.3991326 ,  -0.16749119],
        [  4.847834  ,  -9.30549509],
        [-10.41891144,  -1.15700949],
        [ -7.91249486,  -4.06782072],
        [  1.12418592,  -6.94506866]])}
```

So it returns an `AnnData` when an `AnnData` is passed, but a dict when a less structured object is passed.

This is probably because it's `singledispatched` under the hood, but IDK that this behaviour is great.

I think it could make more sense for this to either:

* Always return an `AnnData`
* Throw an error if something other than an AnnData is passed in

A third option is that we document this behaviour, but I generally don't love it.

There are other places that we do something like this, i.e. return a different type depending on the input. However, I feel like there's more of a loss of information here and less of an obvious return type.

Maybe in future this could get a `return_type: type[AnnData] | type[Dict] | type[xr.Dataset] = AnnData` argument that controls what is returned?

WDYT @ilan-gold @Intron7?

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/scverse/scanpy/pull/2950?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) Report
Attention: Patch coverage is `90.00000%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 73.14%. Comparing base [(`e6c7251`)](https://app.codecov.io/gh/scverse/scanpy/commit/e6c7251d66eae3983baad37575a6b0bba8fe1318?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) to head [(`3633712`)](https://app.codecov.io/gh/scverse/scanpy/pull/2950?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #2950      +/-   ##
==========================================
- Coverage   75.48%   73.14%   -2.35%     
==========================================
  Files         116      116              
  Lines       12904    12872      -32     
==========================================
- Hits         9741     9415     -326     
- Misses       3163     3457     +294     
```


| [Files](https://app.codecov.io/gh/scverse/scanpy/pull/2950?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) | Coverage Δ | |
|---|---|---|
| [scanpy/get/\_aggregated.py](https://app.codecov.io/gh/scverse/scanpy/pull/2950?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c2NhbnB5L2dldC9fYWdncmVnYXRlZC5weQ==) | `94.73% <90.00%> (-0.62%)` | :arrow_down: |

... and [27 files with indirect coverage changes](https://app.codecov.io/gh/scverse/scanpy/pull/2950/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse)

</details>

Patch:
diff --git a/scanpy/get/_aggregated.py b/scanpy/get/_aggregated.py
index 6d603287..5530059c 100644
--- a/scanpy/get/_aggregated.py
+++ b/scanpy/get/_aggregated.py
@@ -156,7 +156,6 @@ def _power(X: Array, power: float | int) -> Array:
     return X**power if isinstance(X, np.ndarray) else X.power(power)
 
 
-@singledispatch
 def aggregate(
     adata: AnnData,
     by: str | Collection[str],
@@ -232,6 +231,11 @@ def aggregate(
 
     Note that this filters out any combination of groups that wasn't present in the original data.
     """
+    if not isinstance(adata, AnnData):
+        raise NotImplementedError(
+            "sc.get.aggregate is currently only implemented for AnnData input, "
+            f"was passed {type(adata)}."
+        )
     if axis is None:
         axis = 1 if varm else 0
     axis, axis_name = _resolve_axis(axis)
@@ -260,7 +264,7 @@ def aggregate(
     dim_df = getattr(adata, axis_name)
     categorical, new_label_df = _combine_categories(dim_df, by)
     # Actual computation
-    layers = aggregate(
+    layers = _aggregate(
         data,
         by=categorical,
         func=func,
@@ -288,13 +292,25 @@ def aggregate(
         return result
 
 
-@aggregate.register(pd.DataFrame)
+@singledispatch
+def _aggregate(
+    data,
+    by: pd.Categorical,
+    func: AggType | Iterable[AggType],
+    *,
+    mask: NDArray[np.bool_] | None = None,
+    dof: int = 1,
+):
+    raise NotImplementedError(f"Data type {type(data)} not supported for aggregation")
+
+
+@_aggregate.register(pd.DataFrame)
 def aggregate_df(data, by, func, *, mask=None, dof=1):
-    return aggregate(data.values, by, func, mask=mask, dof=dof)
+    return _aggregate(data.values, by, func, mask=mask, dof=dof)
 
 
-@aggregate.register(np.ndarray)
-@aggregate.register(sparse.spmatrix)
+@_aggregate.register(np.ndarray)
+@_aggregate.register(sparse.spmatrix)
 def aggregate_array(
     data,
     by: pd.Categorical,


Test patch:
diff --git a/scanpy/tests/test_aggregated.py b/scanpy/tests/test_aggregated.py
index d0e01bb6..99d36947 100644
--- a/scanpy/tests/test_aggregated.py
+++ b/scanpy/tests/test_aggregated.py
@@ -454,3 +454,9 @@ def test_aggregate_obsm_labels():
     )
     result = sc.get.aggregate(adata, by="labels", func="sum", obsm="entry")
     assert_equal(expected, result)
+
+
+def test_dispatch_not_implemented():
+    adata = sc.datasets.blobs()
+    with pytest.raises(NotImplementedError):
+        sc.get.aggregate(adata.X, adata.obs["blobs"], "sum")