Reuse Table constructor to idenfity non-ffi tables when using udtf

Author: timsaucer

python/tests/test_catalog.py

@@ -20,7 +20,7 @@
 import pyarrow as pa
 import pyarrow.dataset as ds
 import pytest
-from datafusion import SessionContext, Table
+from datafusion import SessionContext, Table, udtf
 
 
 # Note we take in `database` as a variable even though we don't use
@@ -232,3 +232,19 @@ def test_in_end_to_end_python_providers(ctx: SessionContext):
             assert len(batches) == 1
             assert batches[0].column(0) == pa.array([1, 2, 3])
             assert batches[0].column(1) == pa.array([4, 5, 6])
+
+
+def test_register_python_function_as_udtf(ctx: SessionContext):
+    basic_table = Table(ctx.sql("SELECT 3 AS value"))
+
+    @udtf("my_table_function")
+    def my_table_function_udtf() -> Table:
+        return basic_table
+
+    ctx.register_udtf(my_table_function_udtf)
+
+    result = ctx.sql("SELECT * FROM my_table_function()").collect()
+    assert len(result) == 1
+    assert len(result[0]) == 1
+    assert len(result[0][0]) == 1
+    assert result[0][0][0].as_py() == 3

src/udtf.rs

@@ -21,12 +21,11 @@ use std::sync::Arc;
 use crate::errors::{py_datafusion_err, to_datafusion_err};
 use crate::expr::PyExpr;
 use crate::table::PyTable;
-use crate::utils::{table_provider_from_pycapsule, validate_pycapsule};
+use crate::utils::validate_pycapsule;
 use datafusion::catalog::{TableFunctionImpl, TableProvider};
 use datafusion::error::Result as DataFusionResult;
 use datafusion::logical_expr::Expr;
 use datafusion_ffi::udtf::{FFI_TableFunction, ForeignTableFunction};
-use pyo3::exceptions::PyNotImplementedError;
 use pyo3::types::{PyCapsule, PyTuple};
 
 /// Represents a user defined table function
@@ -98,11 +97,7 @@ fn call_python_table_function(
         let provider_obj = func.call1(py, py_args)?;
         let provider = provider_obj.bind(py);
 
-        table_provider_from_pycapsule(provider)?.ok_or_else(|| {
-            PyNotImplementedError::new_err(
-                "__datafusion_table_provider__ does not exist on Table Provider object.",
-            )
-        })
+        Ok::<Arc<dyn TableProvider>, PyErr>(PyTable::new(provider)?.table)
     })
     .map_err(to_datafusion_err)
 }