test: add tests for describe method output to PyList and Polars DataFrame

Author: kosiew

python/tests/test_dataframe.py

@@ -1902,6 +1902,21 @@ def test_describe_from_pydict(ctx):
     }
 
 
+def test_describe_to_pylist(df):
+    pylist = df.describe().to_pylist()
+    assert isinstance(pylist, list)
+    assert len(pylist) == 7
+    assert pylist[0]["describe"] == "count"
+
+
+def test_describe_to_polars(df):
+    pl = pytest.importorskip("polars")
+    polars_df = df.describe().to_polars()
+    assert isinstance(polars_df, pl.DataFrame)
+    assert polars_df.shape == (7, 4)
+    assert set(polars_df.columns) == {"describe", "a", "b", "c"}
+
+
 def test_describe_mixed_numeric_string():
     ctx = SessionContext()
     batch = pa.RecordBatch.from_arrays(

src/dataframe.rs

@@ -47,7 +47,7 @@ use crate::catalog::PyTable;
 use crate::errors::{py_datafusion_err, PyDataFusionError};
 use crate::expr::sort_expr::to_sort_expressions;
 use crate::physical_plan::PyExecutionPlan;
-use crate::record_batch::{poll_next_batch, PyRecordBatchStream};
+use crate::record_batch::{poll_next_batch, record_batches_to_pyarrow, PyRecordBatchStream};
 use crate::sql::logical::PyLogicalPlan;
 use crate::utils::{
     get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, spawn_future, validate_pycapsule,
@@ -323,11 +323,9 @@ impl PyDataFrame {
 
         let table_uuid = uuid::Uuid::new_v4().to_string();
 
-        // Convert record batches to PyObject list
-        let py_batches = batches
-            .iter()
-            .map(|rb| rb.to_pyarrow(py))
-            .collect::<PyResult<Vec<PyObject>>>()?;
+        // Convert record batches to PyArrow objects, keeping the underlying
+        // Rust data alive for the lifetime of the Python values.
+        let py_batches = record_batches_to_pyarrow(batches.clone(), py)?;
 
         let py_schema = self.schema().into_pyobject(py)?;
 
@@ -581,9 +579,9 @@ impl PyDataFrame {
     fn collect(&self, py: Python) -> PyResult<Vec<PyObject>> {
         let batches = wait_for_future(py, self.df.as_ref().clone().collect())?
             .map_err(PyDataFusionError::from)?;
-        // cannot use PyResult<Vec<RecordBatch>> return type due to
-        // https://github.com/PyO3/pyo3/issues/1813
-        batches.into_iter().map(|rb| rb.to_pyarrow(py)).collect()
+        // Convert to PyArrow, ensuring the Rust-backed memory outlives
+        // the Python objects.
+        record_batches_to_pyarrow(batches, py)
     }
 
     /// Cache DataFrame.
@@ -600,7 +598,7 @@ impl PyDataFrame {
 
         batches
             .into_iter()
-            .map(|rbs| rbs.into_iter().map(|rb| rb.to_pyarrow(py)).collect())
+            .map(|rbs| record_batches_to_pyarrow(rbs, py))
             .collect()
     }
 

src/record_batch.rs

@@ -36,6 +36,11 @@ pub struct PyRecordBatch {
 #[pymethods]
 impl PyRecordBatch {
     fn to_pyarrow(&self, py: Python) -> PyResult<PyObject> {
+        // Use `ToPyArrow` to create a Python `RecordBatch` that
+        // retains references to the underlying Rust data via the
+        // Arrow C Data Interface. The returned object holds `Arc`
+        // references so the Rust `RecordBatch` can be safely dropped
+        // once the Python object exists.
         self.batch.to_pyarrow(py)
     }
 }
@@ -46,6 +51,18 @@ impl From<RecordBatch> for PyRecordBatch {
     }
 }
 
+/// Convert a collection of [`RecordBatch`]es to Python objects while ensuring
+/// the Rust-backed memory remains alive for the lifetime of the Python
+/// objects.
+pub(crate) fn record_batches_to_pyarrow(
+    batches: Vec<RecordBatch>,
+    py: Python,
+) -> PyResult<Vec<PyObject>> {
+    // Each call to `to_pyarrow` clones `Arc` pointers to the underlying arrays,
+    // tying their lifetime to the resulting Python objects.
+    batches.into_iter().map(|rb| rb.to_pyarrow(py)).collect()
+}
+
 #[pyclass(name = "RecordBatchStream", module = "datafusion", subclass)]
 pub struct PyRecordBatchStream {
     stream: Arc<Mutex<SendableRecordBatchStream>>,