Skip to content

Commit a55969c

Browse files
committed
Add documentation
1 parent c3caf8e commit a55969c

File tree

2 files changed

+43
-7
lines changed

2 files changed

+43
-7
lines changed

docs/source/user-guide/common-operations/udf-and-udfa.rst

+32
Original file line numberDiff line numberDiff line change
@@ -242,3 +242,35 @@ determine which evaluate functions are called.
242242
})
243243
244244
df.select("a", exp_smooth(col("a")).alias("smooth_a")).show()
245+
246+
Table Functions
247+
---------------
248+
249+
User Defined Table Functions are slightly different than the other functions
250+
described here. These functions take any number of `Expr` arguments, but only
251+
literal expressions are supported. Table functions must return a Table
252+
Provider as described in the ref:`_io_custom_table_provider` page.
253+
254+
Once you have a table function, you can register it with the session context
255+
by using :py:func:`datafusion.context.SessionContext.register_udtf`.
256+
257+
There are examples of both rust backed and python based table functions in the
258+
examples folder of the repository. If you have a rust backed table function
259+
that you wish to expose via PyO3, you need to expose it as a ``PyCapsule``.
260+
261+
.. code-block:: rust
262+
263+
#[pymethods]
264+
impl MyTableFunction {
265+
fn __datafusion_table_function__<'py>(
266+
&self,
267+
py: Python<'py>,
268+
) -> PyResult<Bound<'py, PyCapsule>> {
269+
let name = cr"datafusion_table_function".into();
270+
271+
let func = self.clone();
272+
let provider = FFI_TableFunction::new(Arc::new(func), None);
273+
274+
PyCapsule::new(py, provider, Some(name))
275+
}
276+
}

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

+11-7
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from typing import TYPE_CHECKING
2121

2222
import pyarrow as pa
23-
from datafusion import SessionContext, udtf
23+
from datafusion import Expr, SessionContext, udtf
2424
from datafusion_ffi_example import MyTableFunction, MyTableProvider
2525

2626
if TYPE_CHECKING:
@@ -77,19 +77,23 @@ class PythonTableFunction:
7777
provider, and this function takes no arguments
7878
"""
7979

80-
def __init__(self) -> None:
81-
self.table_provider = MyTableProvider(3, 2, 4)
82-
83-
def __call__(self) -> TableProviderExportable:
84-
return self.table_provider
80+
def __call__(
81+
self, num_cols: Expr, num_rows: Expr, num_batches: Expr
82+
) -> TableProviderExportable:
83+
args = [
84+
num_cols.to_variant().value_i64(),
85+
num_rows.to_variant().value_i64(),
86+
num_batches.to_variant().value_i64(),
87+
]
88+
return MyTableProvider(*args)
8589

8690

8791
def test_python_table_function():
8892
ctx = SessionContext()
8993
table_func = PythonTableFunction()
9094
table_udtf = udtf(table_func, "my_table_func")
9195
ctx.register_udtf(table_udtf)
92-
result = ctx.sql("select * from my_table_func()").collect()
96+
result = ctx.sql("select * from my_table_func(3,2,4)").collect()
9397

9498
assert len(result) == 4
9599
assert result[0].num_columns == 3

0 commit comments

Comments
 (0)