Merge pull request #247 from lincc-frameworks/reduce_overhaul

disallow non-column arguments for reduce

Author: dougbrn

src/nested_pandas/nestedframe/core.py

@@ -1007,8 +1007,9 @@ def reduce(self, func, *args, infer_nesting=True, **kwargs) -> NestedFrame:  # t
             columns to apply the function to. See the Notes for recommendations
             on writing func outputs.
         args : positional arguments
-            Positional arguments to pass to the function, the first *args should be the names of the
-            columns to apply the function to.
+            A list of string column names to pull from the NestedFrame to pass along
+            to the function. If the function has additional arguments, pass them as
+            keyword arguments (e.g. `arg_name=value`).
         infer_nesting : bool, default True
             If True, the function will pack output columns into nested
             structures based on column names adhering to a nested naming
@@ -1083,10 +1084,20 @@ def reduce(self, func, *args, infer_nesting=True, **kwargs) -> NestedFrame:  # t
             # Stop when we reach an argument that is not a valid column, as we assume
             # that the remaining args are extra arguments to the function
             if not isinstance(arg, str):
-                break
+                raise TypeError(
+                    f"Received an argument '{arg}' that is not a string. "
+                    "All arguments to `reduce` must be strings corresponding to"
+                    " column names to pass along to the function. If your function"
+                    " has additional arguments, pass them as kwargs (arg_name=value)."
+                )
             components = self._parse_hierarchical_components(arg)
             if not self._is_known_column(components):
-                break
+                raise ValueError(
+                    f"Received a string argument '{arg}' that was not found in the columns list. "
+                    "All arguments to `reduce` must be strings corresponding to"
+                    " column names to pass along to the function. If your function"
+                    " has additional arguments, pass them as kwargs (arg_name=value)."
+                )
             layer = "base" if len(components) < 2 else components[0]
             col = components[-1]
             requested_columns.append((layer, col))

tests/nested_pandas/nestedframe/test_nestedframe.py

@@ -966,7 +966,7 @@ def offset_avg(offset, col_to_avg, column_names):
         sum([7, 10, 7]) / 3.0,
     ]
 
-    result = nf.reduce(offset_avg, "b", "packed.c", ["offset_avg"])
+    result = nf.reduce(offset_avg, "b", "packed.c", column_names=["offset_avg"])
     assert len(result) == len(nf)
     assert isinstance(result, NestedFrame)
     assert result.index.name == "idx"
@@ -978,7 +978,7 @@ def offset_avg(offset, col_to_avg, column_names):
     def make_id(col1, prefix_str):
         return f"{prefix_str}{col1}"
 
-    result = nf.reduce(make_id, "b", "some_id_")
+    result = nf.reduce(make_id, "b", prefix_str="some_id_")
     assert result[0][1] == "some_id_4"
 
 
@@ -1095,6 +1095,27 @@ def complex_output(flux):
     assert list(result.lc.nest.fields) == ["flux_quantiles", "labels"]
 
 
+def test_reduce_arg_errors():
+    """Test that reduce errors based on non-column args trigger as expected"""
+
+    ndf = generate_data(10, 10, seed=1)
+
+    def func(a, flux, add):
+        """a function that takes a scalar, a column, and a boolean"""
+        if add:
+            return {"nested2.flux": flux + a}
+        return {"nested2.flux": flux + a}
+
+    with pytest.raises(TypeError):
+        ndf.reduce(func, "a", "nested.flux", True)
+
+    with pytest.raises(ValueError):
+        ndf.reduce(func, "ab", "nested.flux", add=True)
+
+    # this should work
+    ndf.reduce(func, "a", "nested.flux", add=True)
+
+
 def test_scientific_notation():
     """
     Test that NestedFrame.query handles constants that are written in scientific notation.