Require UDF return type and update docs (#283)

Require UDF return type and update docs

Author: brandon-b-miller

dask_sql/datacontainer.py

@@ -195,14 +195,20 @@ def __init__(self, func, row_udf: bool, return_type=None):
         """
         self.row_udf = row_udf
         self.func = func
+
+        if return_type is None:
+            # These UDFs go through apply and without providing
+            # a return type, dask will attempt to guess it, and
+            # dask might be wrong.
+            raise ValueError("Return type must be provided")
         self.meta = (None, return_type)
 
     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)
+            result = df.apply(self.func, axis=1, meta=self.meta).astype(self.meta[1])
         else:
             result = self.func(*args, **kwargs)
         return result

docs/pages/custom.rst

@@ -43,8 +43,7 @@ These functions may be registered as above and flagged as row UDFs using the `ro
     def f(row):
         return row['a'] + row['b']
 
-    # todo - fix the api
-    c.register_function(f, "f", [], None, row_udf=True)
+    c.register_function(f, "f", [("a", np.int64), ("b", np.int64)], np.int64, row_udf=True)
     c.sql("SELECT f(a, b) FROM data")
 
 ** Note: Row UDFs use `apply` which may have unpredictable performance characteristics, depending on the function and dataframe library **

tests/integration/test_function.py

@@ -38,6 +38,31 @@ def f(row):
     assert_frame_equal(return_df.reset_index(drop=True), df[["a"]] ** 2)
 
 
+@pytest.mark.parametrize(
+    "retty",
+    [None, np.float64, np.float32, np.int64, np.int32, np.int16, np.int8, np.bool_],
+)
+def test_custom_function_row_return_types(c, df, retty):
+    def f(row):
+        return row["a"] ** 2
+
+    if retty is None:
+        with pytest.raises(ValueError):
+            c.register_function(f, "f", [("x", np.float64)], retty, row_udf=True)
+        return
+
+    c.register_function(f, "f", [("x", np.float64)], retty, row_udf=True)
+    return_df = c.sql(
+        """
+        SELECT F(a) AS a
+        FROM df
+        """
+    )
+    return_df = return_df.compute()
+    expectation = (df[["a"]] ** 2).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