From 16ad9a21a9fddabf5c0a9975b406889c81641107 Mon Sep 17 00:00:00 2001
From: dcherian <deepak@cherian.net>
Date: Tue, 25 Apr 2023 21:35:47 -0600
Subject: [PATCH] Optimize len, nanlen

---
 numpy_groupies/aggregate_numpy.py  | 5 ++++-
 numpy_groupies/aggregate_pandas.py | 1 +
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/numpy_groupies/aggregate_numpy.py b/numpy_groupies/aggregate_numpy.py
index 4863671..dd77fbb 100644
--- 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
index 5be2c97..4770fc0 100644
--- 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
     )