Write unit test

datafusion/tests/test_dataframe.py

@@ -533,3 +533,14 @@ def test_cache(df):
 def test_count(df):
     # Get number of rows
     assert df.count() == 3
+
+
+def test_to_pandas(df):
+    # Skip test if pandas is not installed
+    pd = pytest.importorskip("pandas")
+
+    # Convert datafusion dataframe to pandas dataframe
+    pandas_df = df.to_pandas()
+    assert type(pandas_df) == pd.DataFrame
+    assert pandas_df.shape == (3, 3)
+    assert set(pandas_df.columns) == {"a", "b", "c"}

src/dataframe.rs

@@ -319,7 +319,7 @@ impl PyDataFrame {
         let batches = self.collect(py);
 
         Python::with_gil(|py| {
-            // Instantiate pyarrow Table class and use its from_batches method
+            // Instantiate pyarrow Table object and use its from_batches method
             let table_class = py.import("pyarrow")?.getattr("Table")?;
             let args = PyTuple::new(py, batches);
             let table: PyObject = table_class.call_method1("from_batches", args)?.into();