Add examples and tests for registering and using Python UDTFs

Author: kosiew

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

@@ -254,6 +254,26 @@ 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`.
 
+.. code-block:: python
+
+    from datafusion import SessionContext, Table, udtf
+
+    ctx = SessionContext()
+
+    @udtf("table_from_sql")
+    def table_from_sql_udtf() -> Table:
+        return Table.from_dataframe(ctx.sql("SELECT 1 AS value"))
+
+    ctx.register_udtf(table_from_sql_udtf)
+
+    ctx.sql("SELECT * FROM table_from_sql()").show()
+
+.. note::
+
+    The ``TABLE(...)`` wrapper syntax is not currently implemented for
+    invoking table functions from SQL. Call the function directly instead,
+    such as ``SELECT * FROM my_table_function()``.
+
 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``.

examples/python-udtf-table-capsule-regression.py

@@ -19,10 +19,10 @@ def table_from_sql_udtf() -> Table:
 
     ctx.register_udtf(table_from_sql_udtf)
 
-    try:
-        ctx.sql("SELECT * FROM table(table_from_sql())").collect()
-    except NotImplementedError as err:
-        print("Collecting from table_from_sql() failed:", err)
+    result = ctx.sql("SELECT * FROM table_from_sql()").collect()
+    as_pydict = [batch.to_pydict() for batch in result]
+    print("table_from_sql() returned:", as_pydict)
+    assert as_pydict == [{"value": [1]}]
 
     ctx.register_table("numbers", Table(ctx.sql("SELECT 1 AS value")))
 

python/tests/test_context.py

@@ -31,6 +31,7 @@
     Table,
     column,
     literal,
+    udtf,
 )
 
 
@@ -115,6 +116,21 @@ def test_register_record_batches(ctx):
     assert result[0].column(1) == pa.array([-3, -3, -3])
 
 
+def test_register_python_table_function(ctx):
+    @udtf("table_from_sql")
+    def table_from_sql_udtf() -> Table:
+        return Table.from_dataframe(
+            ctx.sql("SELECT 1 AS value UNION ALL SELECT 2 AS value")
+        )
+
+    ctx.register_udtf(table_from_sql_udtf)
+
+    result = ctx.sql("SELECT * FROM table_from_sql() ORDER BY value").collect()
+    table = pa.Table.from_batches(result)
+
+    assert table.to_pydict() == {"value": [1, 2]}
+
+
 def test_create_dataframe_registers_unique_table_name(ctx):
     # create a RecordBatch and register it as memtable
     batch = pa.RecordBatch.from_arrays(