Add PyCapsule support for Arrow import and export (#825)

Author: timsaucer

docs/source/user-guide/io/arrow.rst

@@ -0,0 +1,73 @@
+.. 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.
+
+Arrow
+=====
+
+DataFusion implements the
+`Apache Arrow PyCapsule interface <https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html>`_
+for importing and exporting DataFrames with zero copy. With this feature, any Python
+project that implements this interface can share data back and forth with DataFusion
+with zero copy.
+
+We can demonstrate using `pyarrow <https://arrow.apache.org/docs/python/index.html>`_.
+
+Importing to DataFusion
+-----------------------
+
+Here we will create an Arrow table and import it to DataFusion.
+
+To import an Arrow table, use :py:func:`datafusion.context.SessionContext.from_arrow`.
+This will accept any Python object that implements
+`__arrow_c_stream__ <https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowstream-export>`_
+or `__arrow_c_array__ <https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowarray-export>`_
+and returns a ``StructArray``. Common pyarrow sources you can use are:
+
+- `Array <https://arrow.apache.org/docs/python/generated/pyarrow.Array.html>`_ (but it must return a Struct Array)
+- `Record Batch <https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatch.html>`_
+- `Record Batch Reader <https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatchReader.html>`_
+- `Table <https://arrow.apache.org/docs/python/generated/pyarrow.Table.html>`_
+
+.. ipython:: python
+
+    from datafusion import SessionContext
+    import pyarrow as pa
+
+    data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+    table = pa.Table.from_pydict(data)
+
+    ctx = SessionContext()
+    df = ctx.from_arrow(table)
+    df
+
+Exporting from DataFusion
+-------------------------
+
+DataFusion DataFrames implement ``__arrow_c_stream__`` PyCapsule interface, so any
+Python library that accepts these can import a DataFusion DataFrame directly.
+
+.. warning::
+    It is important to note that this will cause the DataFrame execution to happen, which may be
+    a time consuming task. That is, you will cause a
+    :py:func:`datafusion.dataframe.DataFrame.collect` operation call to occur.
+
+
+.. ipython:: python
+
+    df = df.select((col("a") * lit(1.5)).alias("c"), lit("df").alias("d"))
+    pa.table(df)
+

docs/source/user-guide/io/index.rst

@@ -21,8 +21,8 @@ IO
 .. toctree::
    :maxdepth: 2
 
+   arrow
+   avro
    csv
-   parquet
    json
-   avro
-
+   parquet

examples/import.py

@@ -54,5 +54,5 @@
 
 # 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)
+df = ctx.from_arrow(arrow_table)
 assert type(df) == datafusion.DataFrame

python/datafusion/context.py

@@ -586,19 +586,31 @@ def from_pydict(
         """
         return DataFrame(self.ctx.from_pydict(data, name))
 
-    def from_arrow_table(
-        self, data: pyarrow.Table, name: str | None = None
-    ) -> DataFrame:
-        """Create a :py:class:`~datafusion.dataframe.DataFrame` from an Arrow table.
+    def from_arrow(self, data: Any, name: str | None = None) -> DataFrame:
+        """Create a :py:class:`~datafusion.dataframe.DataFrame` from an Arrow source.
+
+        The Arrow data source can be any object that implements either
+        ``__arrow_c_stream__`` or ``__arrow_c_array__``. For the latter, it must return
+        a struct array. Common examples of sources from pyarrow include
 
         Args:
-            data: Arrow table.
+            data: Arrow data source.
             name: Name of the DataFrame.
 
         Returns:
             DataFrame representation of the Arrow table.
         """
-        return DataFrame(self.ctx.from_arrow_table(data, name))
+        return DataFrame(self.ctx.from_arrow(data, name))
+
+    @deprecated("Use ``from_arrow`` instead.")
+    def from_arrow_table(
+        self, data: pyarrow.Table, name: str | None = None
+    ) -> DataFrame:
+        """Create a :py:class:`~datafusion.dataframe.DataFrame` from an Arrow table.
+
+        This is an alias for :py:func:`from_arrow`.
+        """
+        return self.from_arrow(data, name)
 
     def from_pandas(self, data: pandas.DataFrame, name: str | None = None) -> DataFrame:
         """Create a :py:class:`~datafusion.dataframe.DataFrame` from a Pandas DataFrame.

python/datafusion/dataframe.py

@@ -524,3 +524,19 @@ def unnest_columns(self, *columns: str, preserve_nulls: bool = True) -> DataFram
         """
         columns = [c for c in columns]
         return DataFrame(self.df.unnest_columns(columns, preserve_nulls=preserve_nulls))
+
+    def __arrow_c_stream__(self, requested_schema: pa.Schema) -> Any:
+        """Export an Arrow PyCapsule Stream.
+
+        This will execute and collect the DataFrame. We will attempt to respect the
+        requested schema, but only trivial transformations will be applied such as only
+        returning the fields listed in the requested schema if their data types match
+        those in the DataFrame.
+
+        Args:
+            requested_schema: Attempt to provide the DataFrame using this schema.
+
+        Returns:
+            Arrow PyCapsule object.
+        """
+        return self.df.__arrow_c_stream__(requested_schema)

python/datafusion/tests/test_context.py

@@ -156,7 +156,7 @@ def test_from_arrow_table(ctx):
     table = pa.Table.from_pydict(data)
 
     # convert to DataFrame
-    df = ctx.from_arrow_table(table)
+    df = ctx.from_arrow(table)
     tables = list(ctx.catalog().database().names())
 
     assert df
@@ -166,13 +166,42 @@ def test_from_arrow_table(ctx):
     assert df.collect()[0].num_rows == 3
 
 
+def record_batch_generator(num_batches: int):
+    schema = pa.schema([("a", pa.int64()), ("b", pa.int64())])
+    for i in range(num_batches):
+        yield pa.RecordBatch.from_arrays(
+            [pa.array([1, 2, 3]), pa.array([4, 5, 6])], schema=schema
+        )
+
+
+@pytest.mark.parametrize(
+    "source",
+    [
+        # __arrow_c_array__ sources
+        pa.array([{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}]),
+        # __arrow_c_stream__ sources
+        pa.RecordBatch.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}),
+        pa.RecordBatchReader.from_batches(
+            pa.schema([("a", pa.int64()), ("b", pa.int64())]), record_batch_generator(1)
+        ),
+        pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}),
+    ],
+)
+def test_from_arrow_sources(ctx, source) -> None:
+    df = ctx.from_arrow(source)
+    assert df
+    assert isinstance(df, DataFrame)
+    assert df.schema().names == ["a", "b"]
+    assert df.count() == 3
+
+
 def test_from_arrow_table_with_name(ctx):
     # create a PyArrow table
     data = {"a": [1, 2, 3], "b": [4, 5, 6]}
     table = pa.Table.from_pydict(data)
 
     # convert to DataFrame with optional name
-    df = ctx.from_arrow_table(table, name="tbl")
+    df = ctx.from_arrow(table, name="tbl")
     tables = list(ctx.catalog().database().names())
 
     assert df
@@ -185,7 +214,7 @@ def test_from_arrow_table_empty(ctx):
     table = pa.Table.from_pydict(data, schema=schema)
 
     # convert to DataFrame
-    df = ctx.from_arrow_table(table)
+    df = ctx.from_arrow(table)
     tables = list(ctx.catalog().database().names())
 
     assert df
@@ -200,7 +229,7 @@ def test_from_arrow_table_empty_no_schema(ctx):
     table = pa.Table.from_pydict(data)
 
     # convert to DataFrame
-    df = ctx.from_arrow_table(table)
+    df = ctx.from_arrow(table)
     tables = list(ctx.catalog().database().names())
 
     assert df

python/datafusion/tests/test_dataframe.py

@@ -47,7 +47,7 @@ def df():
         names=["a", "b", "c"],
     )
 
-    return ctx.create_dataframe([[batch]])
+    return ctx.from_arrow(batch)
 
 
 @pytest.fixture
@@ -835,13 +835,42 @@ def test_write_compressed_parquet_missing_compression_level(df, tmp_path, compre
         df.write_parquet(str(path), compression=compression)
 
 
-# ctx = SessionContext()
-
-# # create a RecordBatch and a new DataFrame from it
-# batch = pa.RecordBatch.from_arrays(
-#     [pa.array([1, 2, 3]), pa.array([4, 5, 6]), pa.array([8, 5, 8])],
-#     names=["a", "b", "c"],
-# )
-
-# df = ctx.create_dataframe([[batch]])
-# test_execute_stream(df)
+def test_dataframe_export(df) -> None:
+    # Guarantees that we have the canonical implementation
+    # reading our dataframe export
+    table = pa.table(df)
+    assert table.num_columns == 3
+    assert table.num_rows == 3
+
+    desired_schema = pa.schema([("a", pa.int64())])
+
+    # Verify we can request a schema
+    table = pa.table(df, schema=desired_schema)
+    assert table.num_columns == 1
+    assert table.num_rows == 3
+
+    # Expect a table of nulls if the schema don't overlap
+    desired_schema = pa.schema([("g", pa.string())])
+    table = pa.table(df, schema=desired_schema)
+    assert table.num_columns == 1
+    assert table.num_rows == 3
+    for i in range(0, 3):
+        assert table[0][i].as_py() is None
+
+    # Expect an error when we cannot convert schema
+    desired_schema = pa.schema([("a", pa.float32())])
+    failed_convert = False
+    try:
+        table = pa.table(df, schema=desired_schema)
+    except Exception:
+        failed_convert = True
+    assert failed_convert
+
+    # Expect an error when we have a not set non-nullable
+    desired_schema = pa.schema([("g", pa.string(), False)])
+    failed_convert = False
+    try:
+        table = pa.table(df, schema=desired_schema)
+    except Exception:
+        failed_convert = True
+    assert failed_convert

src/context.rs

@@ -20,12 +20,15 @@ use std::path::PathBuf;
 use std::str::FromStr;
 use std::sync::Arc;
 
+use arrow::array::RecordBatchReader;
+use arrow::ffi_stream::ArrowArrayStreamReader;
+use arrow::pyarrow::FromPyArrow;
 use datafusion::execution::session_state::SessionStateBuilder;
 use object_store::ObjectStore;
 use url::Url;
 use uuid::Uuid;
 
-use pyo3::exceptions::{PyKeyError, PyValueError};
+use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
 use pyo3::prelude::*;
 
 use crate::catalog::{PyCatalog, PyTable};
@@ -444,7 +447,7 @@ impl PySessionContext {
         let table = table_class.call_method1("from_pylist", args)?;
 
         // Convert Arrow Table to datafusion DataFrame
-        let df = self.from_arrow_table(table, name, py)?;
+        let df = self.from_arrow(table, name, py)?;
         Ok(df)
     }
 
@@ -463,29 +466,42 @@ impl PySessionContext {
         let table = table_class.call_method1("from_pydict", args)?;
 
         // Convert Arrow Table to datafusion DataFrame
-        let df = self.from_arrow_table(table, name, py)?;
+        let df = self.from_arrow(table, name, py)?;
         Ok(df)
     }
 
     /// Construct datafusion dataframe from Arrow Table
-    pub fn from_arrow_table(
+    pub fn from_arrow(
         &mut self,
         data: Bound<'_, PyAny>,
         name: Option<&str>,
         py: Python,
     ) -> PyResult<PyDataFrame> {
-        // Instantiate pyarrow Table object & convert to batches
-        let table = data.call_method0("to_batches")?;
+        let (schema, batches) =
+            if let Ok(stream_reader) = ArrowArrayStreamReader::from_pyarrow_bound(&data) {
+                // Works for any object that implements __arrow_c_stream__ in pycapsule.
+
+                let schema = stream_reader.schema().as_ref().to_owned();
+                let batches = stream_reader
+                    .collect::<Result<Vec<RecordBatch>, arrow::error::ArrowError>>()
+                    .map_err(DataFusionError::from)?;
+
+                (schema, batches)
+            } else if let Ok(array) = RecordBatch::from_pyarrow_bound(&data) {
+                // While this says RecordBatch, it will work for any object that implements
+                // __arrow_c_array__ and returns a StructArray.
+
+                (array.schema().as_ref().to_owned(), vec![array])
+            } else {
+                return Err(PyTypeError::new_err(
+                    "Expected either a Arrow Array or Arrow Stream in from_arrow().",
+                ));
+            };
 
-        let schema = data.getattr("schema")?;
-        let schema = schema.extract::<PyArrowType<Schema>>()?;
-
-        // Cast PyAny 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>>>()?;
-        let list_of_batches = PyArrowType::from(vec![batches.0]);
-        self.create_dataframe(list_of_batches, name, Some(schema), py)
+        let list_of_batches = PyArrowType::from(vec![batches]);
+        self.create_dataframe(list_of_batches, name, Some(schema.into()), py)
     }
 
     /// Construct datafusion dataframe from pandas
@@ -504,7 +520,7 @@ impl PySessionContext {
         let table = table_class.call_method1("from_pandas", args)?;
 
         // Convert Arrow Table to datafusion DataFrame
-        let df = self.from_arrow_table(table, name, py)?;
+        let df = self.from_arrow(table, name, py)?;
         Ok(df)
     }
 
@@ -518,7 +534,7 @@ impl PySessionContext {
         let table = data.call_method0("to_arrow")?;
 
         // Convert Arrow Table to datafusion DataFrame
-        let df = self.from_arrow_table(table, name, data.py())?;
+        let df = self.from_arrow(table, name, data.py())?;
         Ok(df)
     }