Optimize len, nanlen

ml31415 · Apr 26, 2023 · 16ad9a2 · 16ad9a2
1 parent d9355bd
commit 16ad9a2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/numpy_groupies/aggregate_numpy.py b/numpy_groupies/aggregate_numpy.py
@@ -301,6 +301,7 @@ def _aggregate_base(
     dtype=None,
     axis=None,
     _impl_dict=_impl_dict,
+    is_pandas=False,
     **kwargs
 ):
     iv = input_validation(group_idx, a, size=size, order=order, axis=axis, func=func)
@@ -326,7 +327,9 @@ def _aggregate_base(
                     kwargs["_nansqueeze"] = True
                 else:
                     good = ~np.isnan(a)
-                    a = a[good]
+                    if "len" not in func or is_pandas:
+                        # a is not needed for len, nanlen!
+                        a = a[good]
                     group_idx = group_idx[good]
 
         dtype = check_dtype(dtype, func, a, flat_size)

diff --git a/numpy_groupies/aggregate_pandas.py b/numpy_groupies/aggregate_pandas.py
@@ -69,6 +69,7 @@ def aggregate(
         func=func,
         axis=axis,
         _impl_dict=_impl_dict,
+        is_pandas=True,
         **kwargs
     )