🐛 Fix #65

pwwang · Oct 4, 2021 · ba8b3e7 · ba8b3e7
1 parent aa90055
commit ba8b3e7
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/datar/core/grouped.py b/datar/core/grouped.py
@@ -338,9 +338,15 @@ def _optimizable_func(
     """
     if not hasattr(numpy, name):
         return None
+
     if kwargs.get("na_rm", False):
+        # Only optimize when na_rm is True, because
+        # pandas will ignore NAs anyway
         return getattr(numpy, f"nan{name}")
-    return getattr(numpy, name)
+
+    # return getattr(numpy, name)
+    # refuse to optimize, as NAs get lost by pandas
+    return None
 
 
 def _optimizable(

diff --git a/tests/test_core_grouped.py b/tests/test_core_grouped.py
@@ -183,7 +183,9 @@ def test_agg():
     out = gf._datar_apply(
         None,
         _mappings=dict(
-            c=FastEvalFunction(mean, args=(f.b,), kwargs={}, dataarg=False)
+            c=FastEvalFunction(
+                mean, args=(f.b,), kwargs={"na_rm": True}, dataarg=False
+            )
         ),
         _method="agg",
     )

diff --git a/tests/test_dplyr_summarise.py b/tests/test_dplyr_summarise.py
@@ -257,7 +257,7 @@ def test_errors(caplog):
     with pytest.raises(ColumnNotExistingError):
         summarise(mtcars, a = mean(f.not_there))
 
-    with pytest.raises(KeyError):
+    with pytest.raises(ColumnNotExistingError):
         summarise(group_by(mtcars, f.cyl), a = mean(f.not_there))
 
     # Duplicate column names