Enhance Table initialization to support DataFrame conversion and update PySessionContext methods to be immutable

Author: kosiew

python/datafusion/catalog.py

@@ -166,11 +166,16 @@ class Table:
 
     def __init__(
         self,
-        table: _InternalRawTable | _InternalTableProvider | Table,
+        table: _InternalRawTable | _InternalTableProvider | Table | Any,
     ) -> None:
-        """Wrap a low level table or table provider."""
+        """Wrap a low level table, table provider, or convertibles like DataFrame."""
         if isinstance(table, Table):
             table = table.table
+        else:
+            from datafusion.dataframe import DataFrame as DataFrameWrapper
+
+            if isinstance(table, DataFrameWrapper):
+                table = table.df.into_view()
 
         if not isinstance(table, (_InternalRawTable, _InternalTableProvider)):
             raise TypeError(EXPECTED_PROVIDER_MSG)

python/tests/test_catalog.py

@@ -196,6 +196,20 @@ def test_schema_register_table_with_dataframe_errors(ctx: SessionContext):
     assert str(exc_info.value) == EXPECTED_PROVIDER_MSG
 
 
+def test_table_wraps_dataframe(ctx: SessionContext):
+    df = ctx.sql("SELECT 1 AS value")
+
+    table = Table(df)
+    ctx.register_table("df_table", table)
+
+    try:
+        batches = ctx.sql("SELECT value FROM df_table").collect()
+        assert len(batches) == 1
+        assert batches[0].column(0).to_pylist() == [1]
+    finally:
+        ctx.deregister_table("df_table")
+
+
 def test_in_end_to_end_python_providers(ctx: SessionContext):
     """Test registering all python providers and running a query against them."""
 

src/context.rs

@@ -429,15 +429,15 @@ impl PySessionContext {
     }
 
     /// Returns a PyDataFrame whose plan corresponds to the SQL statement.
-    pub fn sql(&mut self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
+    pub fn sql(&self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
         let result = self.ctx.sql(query);
         let df = wait_for_future(py, result)??;
         Ok(PyDataFrame::new(df))
     }
 
     #[pyo3(signature = (query, options=None))]
     pub fn sql_with_options(
-        &mut self,
+        &self,
         query: &str,
         options: Option<PySQLOptions>,
         py: Python,

src/utils.rs

@@ -82,7 +82,7 @@ where
     const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(1_000);
 
     py.allow_threads(|| {
-        runtime.block_on(async {
+        let wait_future = || async {
             tokio::pin!(fut);
             loop {
                 tokio::select! {
@@ -92,7 +92,13 @@ where
                     }
                 }
             }
-        })
+        };
+
+        if tokio::runtime::Handle::try_current().is_ok() {
+            tokio::task::block_in_place(|| runtime.block_on(wait_future()))
+        } else {
+            runtime.block_on(wait_future())
+        }
     })
 }