Add table decorator and unit test

Author: timsaucer

examples/datafusion-ffi-example/python/tests/_test_table_function.py

@@ -88,12 +88,8 @@ def __call__(
         return MyTableProvider(*args)
 
 
-def test_python_table_function():
-    ctx = SessionContext()
-    table_func = PythonTableFunction()
-    table_udtf = udtf(table_func, "my_table_func")
-    ctx.register_udtf(table_udtf)
-    result = ctx.sql("select * from my_table_func(3,2,4)").collect()
+def common_table_function_test(test_ctx: SessionContext) -> None:
+    result = test_ctx.sql("select * from my_table_func(3,2,4)").collect()
 
     assert len(result) == 4
     assert result[0].num_columns == 3
@@ -108,3 +104,31 @@ def test_python_table_function():
     ]
 
     assert result == expected
+
+
+def test_python_table_function():
+    ctx = SessionContext()
+    table_func = PythonTableFunction()
+    table_udtf = udtf(table_func, "my_table_func")
+    ctx.register_udtf(table_udtf)
+
+    common_table_function_test(ctx)
+
+
+def test_python_table_function_decorator():
+    ctx = SessionContext()
+
+    @udtf("my_table_func")
+    def my_udtf(
+        num_cols: Expr, num_rows: Expr, num_batches: Expr
+    ) -> TableProviderExportable:
+        args = [
+            num_cols.to_variant().value_i64(),
+            num_rows.to_variant().value_i64(),
+            num_batches.to_variant().value_i64(),
+        ]
+        return MyTableProvider(*args)
+
+    ctx.register_udtf(my_udtf)
+
+    common_table_function_test(ctx)

python/datafusion/udf.py

@@ -804,8 +804,9 @@ def udtf(*args: Any, **kwargs: Any):
             # Case 1: Used as a function, require the first parameter to be callable
             return TableFunction._create_table_udf(*args, **kwargs)
         if args and hasattr(args[0], "__datafusion_table_function__"):
+            # Case 2: We have a datafusion FFI provided function
             return TableFunction(args[1], args[0])
-        # Case 2: Used as a decorator with parameters
+        # Case 3: Used as a decorator with parameters
         return TableFunction._create_table_udf_decorator(*args, **kwargs)
 
     @staticmethod
@@ -820,6 +821,17 @@ def _create_table_udf(
 
         return TableFunction(name, func)
 
+    @staticmethod
+    def _create_table_udf_decorator(
+        name: Optional[str] = None,
+    ) -> Callable[[Callable[[], WindowEvaluator]], Callable[..., Expr]]:
+        """Create a decorator for a WindowUDF."""
+
+        def decorator(func: Callable[[], WindowEvaluator]) -> Callable[..., Expr]:
+            return TableFunction._create_table_udf(func, name)
+
+        return decorator
+
     def __repr__(self) -> str:
         """User printable representation."""
         return self._udtf.__repr__()