Row UDF scalar arguments

dask_sql/datacontainer.py

@@ -205,10 +205,19 @@ def __init__(self, func, row_udf: bool, return_type=None):
 
     def __call__(self, *args, **kwargs):
         if self.row_udf:
-            df = args[0].to_frame()
-            for operand in args[1:]:
-                df[operand.name] = operand
-            result = df.apply(self.func, axis=1, meta=self.meta).astype(self.meta[1])
+            column_args = []
+            scalar_args = []
+            for operand in args:
+                if isinstance(operand, dd.Series):
+                    column_args.append(operand)
+                else:
+                    scalar_args.append(operand)
+                df = column_args[0].to_frame()
+                for col in column_args[1:]:
+                    df[col.name] = operand
+            result = df.apply(
+                self.func, axis=1, args=tuple(scalar_args), meta=self.meta
+            ).astype(self.meta[1])
         else:
             result = self.func(*args, **kwargs)
         return result

tests/integration/test_function.py

@@ -1,3 +1,5 @@
+import operator
+
 import dask.dataframe as dd
 import numpy as np
 import pytest
@@ -63,6 +65,30 @@ def f(row):
     assert_frame_equal(return_df.reset_index(drop=True), expectation)
 
 
+# Test row UDFs with two args
+@pytest.mark.parametrize("k", [1, 1.5, True])
+@pytest.mark.parametrize(
+    "op", [operator.add, operator.sub, operator.mul, operator.truediv]
+)
+@pytest.mark.parametrize("retty", [np.int64, np.float64, np.bool_])
+def test_custom_function_row_args(c, df, k, op, retty):
+    const_type = np.dtype(type(k)).type
+
+    def f(row, k):
+        return op(row["a"], k)
+
+    c.register_function(
+        f, "f", [("a", np.int64), ("k", const_type)], retty, row_udf=True
+    )
+
+    statement = f"SELECT F(a, {k}) as a from df"
+
+    return_df = c.sql(statement)
+    return_df = return_df.compute()
+    expectation = op(df[["a"]], k).astype(retty)
+    assert_frame_equal(return_df.reset_index(drop=True), expectation)
+
+
 def test_multiple_definitions(c, df_simple):
     def f(x):
         return x ** 2