feature: Create dataframe from pandas, polars, dictionary, list or pyarrow Table

datafusion/tests/test_context.py

@@ -26,6 +26,7 @@
     SessionContext,
     SessionConfig,
     RuntimeConfig,
+    DataFrame,
 )
 import pytest
 
@@ -95,6 +96,86 @@ def test_create_dataframe_registers_unique_table_name(ctx):
         assert c in "0123456789abcdef"
 
 
+def test_from_arrow_table(ctx):
+    # create a PyArrow table
+    data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+    table = pa.Table.from_pydict(data)
+
+    # convert to DataFrame
+    df = ctx.from_arrow_table(table)
+    tables = list(ctx.tables())
+
+    assert df
+    assert len(tables) == 1
+    assert type(df) == DataFrame
+    assert set(df.schema().names) == {"a", "b"}
+    assert df.collect()[0].num_rows == 3
+
+
+def test_from_pylist(ctx):
+    # create a dataframe from Python list
+    data = [
+        {"a": 1, "b": 4},
+        {"a": 2, "b": 5},
+        {"a": 3, "b": 6},
+    ]
+
+    df = ctx.from_pylist(data)
+    tables = list(ctx.tables())
+
+    assert df
+    assert len(tables) == 1
+    assert type(df) == DataFrame
+    assert set(df.schema().names) == {"a", "b"}
+    assert df.collect()[0].num_rows == 3
+
+
+def test_from_pydict(ctx):
+    # create a dataframe from Python dictionary
+    data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+
+    df = ctx.from_pydict(data)
+    tables = list(ctx.tables())
+
+    assert df
+    assert len(tables) == 1
+    assert type(df) == DataFrame
+    assert set(df.schema().names) == {"a", "b"}
+    assert df.collect()[0].num_rows == 3
+
+
+def test_from_pandas(ctx):
+    # create a dataframe from pandas dataframe
+    pd = pytest.importorskip("pandas")
+    data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+    pandas_df = pd.DataFrame(data)
+
+    df = ctx.from_pandas(pandas_df)
+    tables = list(ctx.tables())
+
+    assert df
+    assert len(tables) == 1
+    assert type(df) == DataFrame
+    assert set(df.schema().names) == {"a", "b"}
+    assert df.collect()[0].num_rows == 3
+
+
+def test_from_polars(ctx):
+    # create a dataframe from Polars dataframe
+    pd = pytest.importorskip("polars")
+    data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+    polars_df = pd.DataFrame(data)
+
+    df = ctx.from_polars(polars_df)
+    tables = list(ctx.tables())
+
+    assert df
+    assert len(tables) == 1
+    assert type(df) == DataFrame
+    assert set(df.schema().names) == {"a", "b"}
+    assert df.collect()[0].num_rows == 3
+
+
 def test_register_table(ctx, database):
     default = ctx.catalog()
     public = default.database("public")

examples/export.py

@@ -16,18 +16,13 @@
 # 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]])
+# create a new datafusion DataFrame
+df = ctx.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
 # Dataframe:
 # +---+---+
 # | a | b |

examples/import.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 as pa
+import pandas as pd
+import polars as pl
+
+
+# Create a context
+ctx = datafusion.SessionContext()
+
+# Create a datafusion DataFrame from a Python dictionary
+# The dictionary keys represent column names and the dictionary values
+# represent column values
+df = ctx.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
+assert type(df) == datafusion.DataFrame
+# Dataframe:
+# +---+---+
+# | a | b |
+# +---+---+
+# | 1 | 4 |
+# | 2 | 5 |
+# | 3 | 6 |
+# +---+---+
+
+# Create a datafusion DataFrame from a Python list of rows
+df = ctx.from_pylist([{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}])
+assert type(df) == datafusion.DataFrame
+
+# Convert pandas DataFrame to datafusion DataFrame
+pandas_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+df = ctx.from_pandas(pandas_df)
+assert type(df) == datafusion.DataFrame
+
+# Convert polars DataFrame to datafusion DataFrame
+polars_df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+df = ctx.from_polars(polars_df)
+assert type(df) == datafusion.DataFrame
+
+# Convert Arrow Table to datafusion DataFrame
+arrow_table = pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
+df = ctx.from_arrow_table(arrow_table)
+assert type(df) == datafusion.DataFrame

src/context.rs

@@ -50,6 +50,7 @@ use datafusion::prelude::{
     AvroReadOptions, CsvReadOptions, DataFrame, NdJsonReadOptions, ParquetReadOptions,
 };
 use datafusion_common::ScalarValue;
+use pyo3::types::PyTuple;
 use tokio::runtime::Runtime;
 use tokio::task::JoinHandle;
 
@@ -302,6 +303,80 @@ impl PySessionContext {
         PyDataFrame::new(DataFrame::new(self.ctx.state(), plan.plan.as_ref().clone()))
     }
 
+    /// Construct datafusion dataframe from Python list
+    #[allow(clippy::wrong_self_convention)]
+    fn from_pylist(&mut self, data: PyObject, _py: Python) -> PyResult<PyDataFrame> {
+        Python::with_gil(|py| {
+            // Instantiate pyarrow Table object & convert to Arrow Table
+            let table_class = py.import("pyarrow")?.getattr("Table")?;
+            let args = PyTuple::new(py, &[data]);
+            let table = table_class.call_method1("from_pylist", args)?.into();
+
+            // Convert Arrow Table to datafusion DataFrame
+            let df = self.from_arrow_table(table, py)?;
+            Ok(df)
+        })
+    }
+
+    /// Construct datafusion dataframe from Python dictionary
+    #[allow(clippy::wrong_self_convention)]
+    fn from_pydict(&mut self, data: PyObject, _py: Python) -> PyResult<PyDataFrame> {
+        Python::with_gil(|py| {
+            // Instantiate pyarrow Table object & convert to Arrow Table
+            let table_class = py.import("pyarrow")?.getattr("Table")?;
+            let args = PyTuple::new(py, &[data]);
+            let table = table_class.call_method1("from_pydict", args)?.into();
+
+            // Convert Arrow Table to datafusion DataFrame
+            let df = self.from_arrow_table(table, py)?;
+            Ok(df)
+        })
+    }
+
+    /// Construct datafusion dataframe from Arrow Table
+    #[allow(clippy::wrong_self_convention)]
+    fn from_arrow_table(&mut self, data: PyObject, _py: Python) -> PyResult<PyDataFrame> {
+        Python::with_gil(|py| {
+            // Instantiate pyarrow Table object & convert to batches
+            let table = data.call_method0(py, "to_batches")?;
+
+            // Cast PyObject to RecordBatch type
+            // Because create_dataframe() expects a vector of vectors of record batches
+            // here we need to wrap the vector of record batches in an additional vector
+            let batches = table.extract::<PyArrowType<Vec<RecordBatch>>>(py)?;
+            let list_of_batches = PyArrowType::try_from(vec![batches.0])?;
+            self.create_dataframe(list_of_batches, py)
+        })
+    }
+
+    /// Construct datafusion dataframe from pandas
+    #[allow(clippy::wrong_self_convention)]
+    fn from_pandas(&mut self, data: PyObject, _py: Python) -> PyResult<PyDataFrame> {
+        Python::with_gil(|py| {
+            // Instantiate pyarrow Table object & convert to Arrow Table
+            let table_class = py.import("pyarrow")?.getattr("Table")?;
+            let args = PyTuple::new(py, &[data]);
+            let table = table_class.call_method1("from_pandas", args)?.into();
+
+            // Convert Arrow Table to datafusion DataFrame
+            let df = self.from_arrow_table(table, py)?;
+            Ok(df)
+        })
+    }
+
+    /// Construct datafusion dataframe from polars
+    #[allow(clippy::wrong_self_convention)]
+    fn from_polars(&mut self, data: PyObject, _py: Python) -> PyResult<PyDataFrame> {
+        Python::with_gil(|py| {
+            // Convert Polars dataframe to Arrow Table
+            let table = data.call_method0(py, "to_arrow")?;
+
+            // Convert Arrow Table to datafusion DataFrame
+            let df = self.from_arrow_table(table, py)?;
+            Ok(df)
+        })
+    }
+
     fn register_table(&mut self, name: &str, table: &PyTable) -> PyResult<()> {
         self.ctx
             .register_table(name, table.table())