feat: Support args in dataframe apply method

bigframes/dataframe.py

@@ -4828,11 +4828,19 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
                 #   3. The order of the columns in the dataframe must correspond
                 #      to the order of the input params in the function.
                 udf_input_dtypes = func.udf_def.signature.bf_input_types
-                if len(udf_input_dtypes) != len(self.columns) + len(args):
+                if not args and len(udf_input_dtypes) != len(self.columns):
                     raise ValueError(
                         f"Parameter count mismatch: BigFrames BigQuery function"
-                        f" (including the args) expected {len(udf_input_dtypes)}"
-                        f" but received {len(self.columns) + len(args)}."
+                        f" expected {len(udf_input_dtypes)} parameters but"
+                        f" received {len(self.columns)} DataFrame columns."
+                    )
+                if args and len(udf_input_dtypes) != len(self.columns) + len(args):
+                    raise ValueError(
+                        f"Parameter count mismatch: BigFrames BigQuery function"
+                        f" expected {len(udf_input_dtypes)} parameters but"
+                        f" received {len(self.columns) + len(args)} values"
+                        f" ({len(self.columns)} DataFrame coulmns and"
+                        f" {len(args)} args)."
                     )
                 end_slice = -len(args) if args else None
                 if udf_input_dtypes[:end_slice] != tuple(self.dtypes.to_list()):

tests/system/large/functions/test_managed_function.py

@@ -467,12 +467,14 @@ def foo(x, y, z):
 
         # Fails to apply on dataframe with incompatible number of columns.
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 3 but received 2."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 3 parameters but received 2 DataFrame.*",
         ):
             bf_df[["Id", "Age"]].apply(foo, axis=1)
 
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 3 but received 4."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 3 parameters but received 4 DataFrame.*",
         ):
             bf_df.assign(Country="lalaland").apply(foo, axis=1)
 
@@ -983,7 +985,8 @@ def the_sum(s1, s2, x):
 
         # Fails to apply on dataframe with incompatible number of columns.
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 3 but received 4."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 3 parameters but received 4 values.*",
         ):
             scalars_df[columns + ["float64_col"]].apply(the_sum_mf, axis=1, args=args1)
 

tests/system/large/functions/test_remote_function.py

@@ -1958,7 +1958,8 @@ def the_sum(s1, s2, x):
 
         # Fails to apply on dataframe with incompatible number of columns.
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 3 but received 4."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 3 parameters but received 4 values.*",
         ):
             scalars_df[columns].apply(
                 the_sum_mf,
@@ -2306,11 +2307,13 @@ def foo(x, y, z):
 
         # Fails to apply on dataframe with incompatible number of columns
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 3 but received 2."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 3 parameters but received 2 DataFrame.*",
         ):
             bf_df[["Id", "Age"]].apply(foo, axis=1)
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 3 but received 4."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 3 parameters but received 4 DataFrame.*",
         ):
             bf_df.assign(Country="lalaland").apply(foo, axis=1)
 
@@ -2388,11 +2391,13 @@ def foo(x, y, z):
 
         # Fails to apply on dataframe with incompatible number of columns
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 3 but received 2."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 3 parameters but received 2 DataFrame.*",
         ):
             bf_df[["Id", "Age"]].apply(foo, axis=1)
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 3 but received 4."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 3 parameters but received 4 DataFrame.*",
         ):
             bf_df.assign(Country="lalaland").apply(foo, axis=1)
 
@@ -2460,11 +2465,13 @@ def foo(x):
 
         # Fails to apply on dataframe with incompatible number of columns
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 1 but received 0."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 1 parameters but received 0 DataFrame.*",
         ):
             bf_df[[]].apply(foo, axis=1)
         with pytest.raises(
-            ValueError, match="^Parameter count mismatch:.* expected 1 but received 2."
+            ValueError,
+            match="^Parameter count mismatch:.* expected 1 parameters but received 2 DataFrame.*",
         ):
             bf_df.assign(Country="lalaland").apply(foo, axis=1)