feature: Additional export methods

datafusion/tests/test_dataframe.py

@@ -544,3 +544,70 @@ def test_to_pandas(df):
     assert type(pandas_df) == pd.DataFrame
     assert pandas_df.shape == (3, 3)
     assert set(pandas_df.columns) == {"a", "b", "c"}
+
+
+def test_empty_to_pandas(df):
+    # Skip test if pandas is not installed
+    pd = pytest.importorskip("pandas")
+
+    # Convert empty datafusion dataframe to pandas dataframe
+    pandas_df = df.limit(0).to_pandas()
+    assert type(pandas_df) == pd.DataFrame
+    assert pandas_df.shape == (0, 3)
+    assert set(pandas_df.columns) == {"a", "b", "c"}
+
+
+def test_to_polars(df):
+    # Skip test if polars is not installed
+    pl = pytest.importorskip("polars")
+
+    # Convert datafusion dataframe to polars dataframe
+    polars_df = df.to_polars()
+    assert type(polars_df) == pl.DataFrame
+    assert polars_df.shape == (3, 3)
+    assert set(polars_df.columns) == {"a", "b", "c"}
+
+
+def test_empty_to_polars(df):
+    # Skip test if polars is not installed
+    pl = pytest.importorskip("polars")
+
+    # Convert empty datafusion dataframe to polars dataframe
+    polars_df = df.limit(0).to_polars()
+    assert type(polars_df) == pl.DataFrame
+    assert polars_df.shape == (0, 3)
+    assert set(polars_df.columns) == {"a", "b", "c"}
+
+
+def test_to_arrow_table(df):
+    # Convert datafusion dataframe to pyarrow Table
+    pyarrow_table = df.to_arrow_table()
+    assert type(pyarrow_table) == pa.Table
+    assert pyarrow_table.shape == (3, 3)
+    assert set(pyarrow_table.column_names) == {"a", "b", "c"}
+
+
+def test_empty_to_arrow_table(df):
+    # Convert empty datafusion dataframe to pyarrow Table
+    pyarrow_table = df.limit(0).to_arrow_table()
+    assert type(pyarrow_table) == pa.Table
+    assert pyarrow_table.shape == (0, 3)
+    assert set(pyarrow_table.column_names) == {"a", "b", "c"}
+
+
+def test_to_pylist(df):
+    # Convert datafusion dataframe to Python list
+    pylist = df.to_pylist()
+    assert type(pylist) == list
+    assert pylist == [
+        {"a": 1, "b": 4, "c": 8},
+        {"a": 2, "b": 5, "c": 5},
+        {"a": 3, "b": 6, "c": 8},
+    ]
+
+
+def test_to_pydict(df):
+    # Convert datafusion dataframe to Python dictionary
+    pydict = df.to_pydict()
+    assert type(pydict) == dict
+    assert pydict == {"a": [1, 2, 3], "b": [4, 5, 6], "c": [8, 5, 8]}

examples/export.py

@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datafusion
+import pyarrow
+
+
+# create a context
+ctx = datafusion.SessionContext()
+
+# create a RecordBatch and a new datafusion DataFrame from it
+batch = pyarrow.RecordBatch.from_arrays(
+    [pyarrow.array([1, 2, 3]), pyarrow.array([4, 5, 6])],
+    names=["a", "b"],
+)
+df = ctx.create_dataframe([[batch]])
+# Dataframe:
+# +---+---+
+# | a | b |
+# +---+---+
+# | 1 | 4 |
+# | 2 | 5 |
+# | 3 | 6 |
+# +---+---+
+
+# export to pandas dataframe
+pandas_df = df.to_pandas()
+assert pandas_df.shape == (3, 2)
+
+# export to PyArrow table
+arrow_table = df.to_arrow_table()
+assert arrow_table.shape == (3, 2)
+
+# export to Polars dataframe
+polars_df = df.to_polars()
+assert polars_df.shape == (3, 2)
+
+# export to Python list of rows
+pylist = df.to_pylist()
+assert pylist == [{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}]
+
+# export to Pyton dictionary of columns
+pydict = df.to_pydict()
+assert pydict == {"a": [1, 2, 3], "b": [4, 5, 6]}

src/dataframe.rs

@@ -313,24 +313,70 @@ impl PyDataFrame {
         Ok(())
     }
 
-    /// Convert to pandas dataframe with pyarrow
-    /// Collect the batches, pass to Arrow Table & then convert to Pandas DataFrame
-    fn to_pandas(&self, py: Python) -> PyResult<PyObject> {
-        let batches = self.collect(py);
+    /// Convert to Arrow Table
+    /// Collect the batches and pass to Arrow Table
+    fn to_arrow_table(&self, py: Python) -> PyResult<PyObject> {
+        let batches = self.collect(py)?.to_object(py);
+        let schema: PyObject = self.schema().into_py(py);
 
         Python::with_gil(|py| {
             // 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 args = PyTuple::new(py, &[batches, schema]);
             let table: PyObject = table_class.call_method1("from_batches", args)?.into();
+            Ok(table)
+        })
+    }
+
+    /// Convert to pandas dataframe with pyarrow
+    /// Collect the batches, pass to Arrow Table & then convert to Pandas DataFrame
+    fn to_pandas(&self, py: Python) -> PyResult<PyObject> {
+        let table = self.to_arrow_table(py)?;
 
-            // Use Table.to_pandas() method to convert batches to pandas dataframe
+        Python::with_gil(|py| {
             // See also: https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table.to_pandas
             let result = table.call_method0(py, "to_pandas")?;
             Ok(result)
         })
     }
 
+    /// Convert to Python list using pyarrow
+    /// Each list item represents one row encoded as dictionary
+    fn to_pylist(&self, py: Python) -> PyResult<PyObject> {
+        let table = self.to_arrow_table(py)?;
+
+        Python::with_gil(|py| {
+            // See also: https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table.to_pylist
+            let result = table.call_method0(py, "to_pylist")?;
+            Ok(result)
+        })
+    }
+
+    /// Convert to Python dictionary using pyarrow
+    /// Each dictionary key is a column and the dictionary value represents the column values
+    fn to_pydict(&self, py: Python) -> PyResult<PyObject> {
+        let table = self.to_arrow_table(py)?;
+
+        Python::with_gil(|py| {
+            // See also: https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table.to_pydict
+            let result = table.call_method0(py, "to_pydict")?;
+            Ok(result)
+        })
+    }
+
+    /// Convert to polars dataframe with pyarrow
+    /// Collect the batches, pass to Arrow Table & then convert to polars DataFrame
+    fn to_polars(&self, py: Python) -> PyResult<PyObject> {
+        let table = self.to_arrow_table(py)?;
+
+        Python::with_gil(|py| {
+            let dataframe = py.import("polars")?.getattr("DataFrame")?;
+            let args = PyTuple::new(py, &[table]);
+            let result: PyObject = dataframe.call1(args)?.into();
+            Ok(result)
+        })
+    }
+
     // Executes this DataFrame to get the total number of rows.
     fn count(&self, py: Python) -> PyResult<usize> {
         Ok(wait_for_future(py, self.df.as_ref().clone().count())?)