feat: add tests for registering DataFrame as a view and handling TableProvider

Author: kosiew

python/tests/test_context.py

@@ -27,6 +27,7 @@
     SessionConfig,
     SessionContext,
     SQLOptions,
+    TableProvider,
     column,
     literal,
 )
@@ -330,6 +331,35 @@ def test_deregister_table(ctx, database):
     assert public.names() == {"csv1", "csv2"}
 
 
+def test_register_table_from_dataframe_into_view(ctx):
+    df = ctx.from_pydict({"a": [1, 2]})
+    provider = df.into_view()
+    ctx.register_table("view_tbl", provider)
+    result = ctx.sql("SELECT * FROM view_tbl").collect()
+    assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
+
+
+def test_table_provider_from_capsule(ctx):
+    df = ctx.from_pydict({"a": [1, 2]})
+    provider = df.into_view()
+    capsule = provider.__datafusion_table_provider__()
+    provider2 = TableProvider.from_capsule(capsule)
+    ctx.register_table("capsule_tbl", provider2)
+    result = ctx.sql("SELECT * FROM capsule_tbl").collect()
+    assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
+
+
+def test_table_provider_from_capsule_invalid():
+    with pytest.raises(Exception):  # noqa: B017
+        TableProvider.from_capsule(object())
+
+
+def test_register_table_with_dataframe_errors(ctx):
+    df = ctx.from_pydict({"a": [1]})
+    with pytest.raises(Exception):  # noqa: B017
+        ctx.register_table("bad", df)
+
+
 def test_register_dataset(ctx):
     # create a RecordBatch and register it as a pyarrow.dataset.Dataset
     batch = pa.RecordBatch.from_arrays(

tests/dataframe_into_view.rs

@@ -0,0 +1,39 @@
+use std::sync::Arc;
+
+use datafusion::arrow::array::Int32Array;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::datasource::MemTable;
+use datafusion::prelude::SessionContext;
+use datafusion_python::dataframe::PyDataFrame;
+
+#[test]
+fn dataframe_into_view_returns_table_provider() {
+    // Create an in-memory table with one Int32 column.
+    let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+    let batch = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+    )
+    .unwrap();
+    let table = MemTable::try_new(schema, vec![vec![batch]]).unwrap();
+
+    // Build a DataFrame from the table and convert it into a view.
+    let ctx = SessionContext::new();
+    let df = ctx.read_table(Arc::new(table)).unwrap();
+    let py_df = PyDataFrame::new(df);
+    let provider = py_df.into_view().unwrap();
+
+    // Register the view in a new context and ensure it can be queried.
+    let ctx = SessionContext::new();
+    ctx.register_table("view", provider.as_table().table())
+        .unwrap();
+
+    let rt = tokio::runtime::Runtime::new().unwrap();
+    let batches = rt
+        .block_on(ctx.sql("SELECT * FROM view").unwrap().collect())
+        .unwrap();
+
+    assert_eq!(batches.len(), 1);
+    assert_eq!(batches[0].num_rows(), 3);
+}