feat: enhance DataType and Expr to support objects implementing __arrow_c_schema__ without requiring PyArrow

Author: kosiew

docs/source/user-guide/common-operations/functions.rst

@@ -109,6 +109,10 @@ Casting
 
 Casting expressions to different data types using :py:func:`~datafusion.functions.arrow_cast`
 
+DataFusion's :class:`~datafusion.types.DataType` can be constructed from any
+object implementing ``__arrow_c_schema__`` and passed to ``arrow_cast`` without
+requiring :mod:`pyarrow`.
+
 .. ipython:: python
 
     df.select(

python/datafusion/expr.py

@@ -22,7 +22,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import pyarrow as pa
 
@@ -32,6 +32,7 @@
     from typing_extensions import deprecated  # Python 3.12
 
 from datafusion.common import DataTypeMap, NullTreatment, RexType
+from datafusion.types import DataType
 
 from ._internal import expr as expr_internal
 from ._internal import functions as functions_internal
@@ -515,21 +516,34 @@ def fill_null(self, value: Any | Expr | None = None) -> Expr:
             value = Expr.literal(value)
         return Expr(functions_internal.nvl(self.expr, value.expr))
 
-    _to_pyarrow_types: ClassVar[dict[type, pa.DataType]] = {
-        float: pa.float64(),
-        int: pa.int64(),
-        str: pa.string(),
-        bool: pa.bool_(),
-    }
+    def cast(self, to: DataType | type[float | int | str | bool] | object) -> Expr:
+        """Cast to a new data type.
 
-    def cast(self, to: pa.DataType[Any] | type[float | int | str | bool]) -> Expr:
-        """Cast to a new data type."""
-        if not isinstance(to, pa.DataType):
+        Args:
+            to: Target type as :class:`datafusion.types.DataType`, a Python
+                builtin (:class:`float`, :class:`int`, :class:`str`,
+                :class:`bool`), or any object implementing
+                ``__arrow_c_schema__``.
+        """
+        if isinstance(to, type):
+            mapping = {float: "Float64", int: "Int64", str: "Utf8", bool: "Boolean"}
             try:
-                to = self._to_pyarrow_types[to]
+                to = DataType.from_str(mapping[to])
             except KeyError as err:
-                error_msg = "Expected instance of pyarrow.DataType or builtins.type"
-                raise TypeError(error_msg) from err
+                msg = (
+                    "Expected DataType, builtin type (float, int, str, bool), "
+                    "or object implementing __arrow_c_schema__"
+                )
+                raise TypeError(msg) from err
+        elif not isinstance(to, DataType):
+            try:
+                to = DataType.from_arrow_c_schema(to)
+            except Exception as err:  # pragma: no cover - type check
+                msg = (
+                    "Expected DataType, builtin type (float, int, str, bool), "
+                    "or object implementing __arrow_c_schema__"
+                )
+                raise TypeError(msg) from err
 
         return Expr(self.expr.cast(to))
 

python/datafusion/types.py

@@ -1,15 +1,20 @@
-"""Internal Arrow type helpers with optional PyArrow conversion."""
+"""Arrow type helpers with optional PyArrow conversion.
+
+This module exposes :class:`datafusion.common.DataType`, a lightweight
+representation of Arrow types that can be created from any Python object
+implementing ``__arrow_c_schema__`` without requiring :mod:`pyarrow`.
+"""
 
 from __future__ import annotations
 
 from typing import Any
 
 try:  # pragma: no cover - optional dependency
     import pyarrow as pa
-except Exception:  # pragma: no cover - optional dependency
-    pa = None  # type: ignore
+except ModuleNotFoundError:  # pragma: no cover - optional dependency
+    pa = None  # type: ignore[assignment]
 
-from datafusion.common import DataTypeMap
+from datafusion.common import DataType, DataTypeMap
 
 _PYARROW_TYPE_FACTORIES = {
     "Null": lambda: pa.null() if pa else None,
@@ -31,18 +36,17 @@
 
 def pyarrow_available() -> bool:
     """Return ``True`` if :mod:`pyarrow` can be imported."""
-
     return pa is not None
 
 
-def to_pyarrow(data_type: DataTypeMap) -> "pa.DataType":
+def to_pyarrow(data_type: DataTypeMap) -> pa.DataType:
     """Convert a :class:`DataTypeMap` to a :mod:`pyarrow` data type.
 
     Raises ``ModuleNotFoundError`` if :mod:`pyarrow` is not installed.
     """
-
     if pa is None:  # pragma: no cover - optional dependency
-        raise ModuleNotFoundError("pyarrow is not installed")
+        msg = "pyarrow is not installed"
+        raise ModuleNotFoundError(msg)
     name = str(data_type.arrow_type)
     factory = _PYARROW_TYPE_FACTORIES.get(name)
     if factory is None:
@@ -51,14 +55,14 @@ def to_pyarrow(data_type: DataTypeMap) -> "pa.DataType":
     return factory()
 
 
-def from_pyarrow(pa_type: "pa.DataType") -> DataTypeMap:
+def from_pyarrow(pa_type: pa.DataType) -> DataTypeMap:
     """Convert a :mod:`pyarrow` data type to :class:`DataTypeMap`.
 
     Raises ``ModuleNotFoundError`` if :mod:`pyarrow` is not installed.
     """
-
     if pa is None:  # pragma: no cover - optional dependency
-        raise ModuleNotFoundError("pyarrow is not installed")
+        msg = "pyarrow is not installed"
+        raise ModuleNotFoundError(msg)
     return DataTypeMap.py_map_from_arrow_type_str(str(pa_type))
 
 
@@ -69,7 +73,16 @@ def ensure_pyarrow_type(value: DataTypeMap | Any) -> Any:
     it will be converted to the corresponding :mod:`pyarrow` data type.
     Otherwise ``value`` is returned unchanged.
     """
-
     if isinstance(value, DataTypeMap):
         return to_pyarrow(value) if pyarrow_available() else value
     return value
+
+
+__all__ = [
+    "DataType",
+    "DataTypeMap",
+    "ensure_pyarrow_type",
+    "from_pyarrow",
+    "pyarrow_available",
+    "to_pyarrow",
+]

src/common/data_type.rs

@@ -17,6 +17,7 @@
 
 use datafusion::arrow::array::Array;
 use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
+use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::common::{DataFusionError, ScalarValue};
 use datafusion::logical_expr::sqlparser::ast::NullTreatment as DFNullTreatment;
 use pyo3::{exceptions::PyValueError, prelude::*};
@@ -672,6 +673,23 @@ impl PyDataType {
     }
 }
 
+#[pymethods]
+impl PyDataType {
+    #[staticmethod]
+    pub fn from_arrow_c_schema(obj: PyArrowType<DataType>) -> Self {
+        PyDataType { data_type: obj.0 }
+    }
+
+    #[staticmethod]
+    pub fn from_str(arrow_str_type: &str) -> PyResult<Self> {
+        Self::py_map_from_arrow_type_str(arrow_str_type.to_string())
+    }
+
+    fn __repr__(&self) -> String {
+        format!("{:?}", self.data_type)
+    }
+}
+
 impl From<PyDataType> for DataType {
     fn from(data_type: PyDataType) -> DataType {
         data_type.data_type

src/expr.rs

@@ -28,7 +28,6 @@ use std::sync::Arc;
 use window::PyWindowFrame;
 
 use datafusion::arrow::datatypes::{DataType, Field};
-use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::functions::core::expr_ext::FieldAccessor;
 use datafusion::logical_expr::{
     col,
@@ -37,7 +36,7 @@ use datafusion::logical_expr::{
 };
 
 use crate::arrow_ffi::scalar_to_pyarrow;
-use crate::common::data_type::{DataTypeMap, NullTreatment, PyScalarValue, RexType};
+use crate::common::data_type::{DataTypeMap, NullTreatment, PyDataType, PyScalarValue, RexType};
 use crate::errors::{py_runtime_err, py_type_err, py_unsupported_variant_err, PyDataFusionResult};
 use crate::expr::aggregate_expr::PyAggregateFunction;
 use crate::expr::binary_expr::PyBinaryExpr;
@@ -313,10 +312,13 @@ impl PyExpr {
         self.expr.clone().is_not_null().into()
     }
 
-    pub fn cast(&self, to: PyArrowType<DataType>) -> PyExpr {
+    pub fn cast(&self, to: PyDataType) -> PyExpr {
         // self.expr.cast_to() requires DFSchema to validate that the cast
         // is supported, omit that for now
-        let expr = Expr::Cast(Cast::new(Box::new(self.expr.clone()), to.0));
+        let expr = Expr::Cast(Cast::new(
+            Box::new(self.expr.clone()),
+            to.data_type,
+        ));
         expr.into()
     }