Add documentation

timsaucer · timsaucer · commit a55969cfcfac · 2025-05-05T09:56:20.000-04:00
diff --git a/docs/source/user-guide/common-operations/udf-and-udfa.rst b/docs/source/user-guide/common-operations/udf-and-udfa.rst
@@ -242,3 +242,35 @@ determine which evaluate functions are called.
     })
 
     df.select("a", exp_smooth(col("a")).alias("smooth_a")).show()
+
+Table Functions
+---------------
+
+User Defined Table Functions are slightly different than the other functions
+described here. These functions take any number of `Expr` arguments, but only
+literal expressions are supported. Table functions must return a Table
+Provider as described in the ref:`_io_custom_table_provider` page.
+
+Once you have a table function, you can register it with the session context
+by using :py:func:`datafusion.context.SessionContext.register_udtf`.
+
+There are examples of both rust backed and python based table functions in the
+examples folder of the repository. If you have a rust backed table function
+that you wish to expose via PyO3, you need to expose it as a ``PyCapsule``.
+
+.. code-block:: rust
+
+    #[pymethods]
+    impl MyTableFunction {
+        fn __datafusion_table_function__<'py>(
+            &self,
+            py: Python<'py>,
+        ) -> PyResult<Bound<'py, PyCapsule>> {
+            let name = cr"datafusion_table_function".into();
+
+            let func = self.clone();
+            let provider = FFI_TableFunction::new(Arc::new(func), None);
+
+            PyCapsule::new(py, provider, Some(name))
+        }
+    }
diff --git a/examples/datafusion-ffi-example/python/tests/_test_table_function.py b/examples/datafusion-ffi-example/python/tests/_test_table_function.py
@@ -20,7 +20,7 @@
 from typing import TYPE_CHECKING
 
 import pyarrow as pa
-from datafusion import SessionContext, udtf
+from datafusion import Expr, SessionContext, udtf
 from datafusion_ffi_example import MyTableFunction, MyTableProvider
 
 if TYPE_CHECKING:
@@ -77,19 +77,23 @@ class PythonTableFunction:
     provider, and this function takes no arguments
     """
 
-    def __init__(self) -> None:
-        self.table_provider = MyTableProvider(3, 2, 4)
-
-    def __call__(self) -> TableProviderExportable:
-        return self.table_provider
+    def __call__(
+        self, 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)
 
 
 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()").collect()
+    result = ctx.sql("select * from my_table_func(3,2,4)").collect()
 
     assert len(result) == 4
     assert result[0].num_columns == 3
