feat: unify all Daft type to Python type conversions (#4972)

## Changes Made

We used to have five separate code paths for converting Daft types to
Python:
1. Literal -> Python
2. `Series.to_pylist`
3. `Series.__iter__`
4. the SQL planner's `lit_to_py_any`
5. Series casting logic

This PR makes all of them go through the literal to Python code path,
making the conversion logic consistent across our library.

Also added extension type to literal.

TODO:
- [x] add conversion mapping to docs

## Related Issues

<!-- Link to related GitHub issues, e.g., "Closes #123" -->

## Checklist

- [x] Documented in API Docs (if applicable)
- [x] Documented in User Guide (if applicable)
- [x] If adding a new documentation page, doc is added to
`docs/mkdocs.yml` navigation
- [x] Documentation builds and is formatted properly (tag @/ccmao1130
for docs review)

Author: kevinzwang

Cargo.lock

@@ -235,6 +235,7 @@ common-error = {path = "src/common/error", default-features = false}
 common-file-formats = {path = "src/common/file-formats"}
 common-image = {path = "src/common/image"}
 common-metrics = {path = "src/common/metrics"}
+common-ndarray = {path = "src/common/ndarray"}
 common-runtime = {path = "src/common/runtime", default-features = false}
 daft-algebra = {path = "src/daft-algebra"}
 daft-context = {path = "src/daft-context"}
@@ -273,9 +274,11 @@ jaq-core = "2.2.0"
 jaq-json = {version = "1.1.2", features = ["serde_json"]}
 jaq-std = "2.1.1"
 mur3 = "0.1.0"
+ndarray = "0.16.1"
 num-derive = "0.4.2"
 num-format = "0.4.4"
 num-traits = "0.2"
+numpy = "0.25.0"
 opentelemetry = {version = "0.29", features = ["trace", "metrics"]}
 opentelemetry-otlp = {version = "0.29", features = ["grpc-tonic"]}
 opentelemetry_sdk = "0.29"
@@ -348,7 +351,7 @@ features = ['async']
 path = "src/parquet2"
 
 [workspace.dependencies.pyo3]
-features = ["extension-module", "multiple-pymethods", "abi3-py39", "indexmap", "chrono"]
+features = ["extension-module", "multiple-pymethods", "abi3-py39", "indexmap", "chrono", "chrono-tz"]
 version = "0.25.1"
 
 [workspace.dependencies.pyo3-async-runtimes]

Cargo.toml

@@ -1327,6 +1327,8 @@ class PySeries:
     def from_pylist(name: str, pylist: list[Any], dtype: PyDataType) -> PySeries: ...
     def to_pylist(self) -> list[Any]: ...
     def to_arrow(self) -> pa.Array: ...
+    def __iter__(self) -> PySeriesIterator: ...
+    def __getitem__(self, index: int) -> Any: ...
     def __abs__(self) -> PySeries: ...
     def __add__(self, other: PySeries) -> PySeries: ...
     def __sub__(self, other: PySeries) -> PySeries: ...
@@ -1402,6 +1404,10 @@ class PySeries:
     @staticmethod
     def _debug_bincode_deserialize(b: bytes) -> PySeries: ...
 
+class PySeriesIterator:
+    def __next__(self) -> Any: ...
+    def __iter__(self) -> PySeriesIterator: ...
+
 class PyShowOptions:
     pass
 

daft/daft/__init__.pyi

@@ -5,51 +5,18 @@
 
 import daft.daft as native
 from daft.arrow_utils import ensure_array, ensure_chunked_array
-from daft.daft import CountMode, ImageFormat, ImageMode, PyRecordBatch, PySeries
+from daft.daft import CountMode, ImageFormat, ImageMode, PyRecordBatch, PySeries, PySeriesIterator
 from daft.datatype import DataType, TimeUnit, _ensure_registered_super_ext_type
 from daft.dependencies import np, pa, pd
 from daft.schema import Field
 from daft.utils import pyarrow_supports_fixed_shape_tensor
 
 if TYPE_CHECKING:
     import builtins
-    from collections.abc import Iterator
 
     from daft.daft import PyDataType
 
 
-class SeriesIterable:
-    """Iterable wrapper for Series that efficiently handles different data types."""
-
-    def __init__(self, series: Series):
-        self.series = series
-
-    def __iter__(self) -> Iterator[Any]:
-        dt = self.series.datatype()
-        if dt == DataType.python():
-
-            def yield_pylist() -> Iterator[Any]:
-                yield from self.series._series.to_pylist()
-
-            return yield_pylist()
-        elif dt._should_cast_to_python():
-
-            def yield_pylist() -> Iterator[Any]:
-                yield from self.series._series.cast(DataType.python()._dtype).to_pylist()
-
-            return yield_pylist()
-        else:
-
-            def arrow_to_py() -> Iterator[Any]:
-                # We directly call .to_arrow() on the internal PySeries object since the case
-                # above has already captured the fixed shape tensor case.
-                arrow_data = self.series._series.to_arrow()
-                for item in arrow_data:
-                    yield None if item is None else item.as_py()
-
-            return arrow_to_py()
-
-
 class Series:
     """A Daft Series is an array of data of a single type, and is usually a column in a DataFrame."""
 
@@ -58,10 +25,9 @@ class Series:
     def __init__(self) -> None:
         raise NotImplementedError("We do not support creating a Series via __init__ ")
 
-    def __iter__(self) -> Iterator[Any]:
+    def __iter__(self) -> PySeriesIterator:
         """Return an iterator over the elements of the Series."""
-        iterable = SeriesIterable(self)
-        return iterable.__iter__()
+        return self._series.__iter__()
 
     @staticmethod
     def _from_pyseries(pyseries: PySeries) -> Series:
@@ -227,19 +193,6 @@ def from_pandas(cls, data: pd.Series[Any], name: str = "pd_series", dtype: DataT
     def cast(self, dtype: DataType) -> Series:
         return Series._from_pyseries(self._series.cast(dtype._dtype))
 
-    def _cast_to_python(self) -> Series:
-        """Convert this Series into a Series of Python objects.
-
-        Call Series.to_pylist() and create a new Series from the raw Pylist directly.
-
-        This logic is needed by the Rust implementation of cast(),
-        but is written here (i.e. not in Rust) for conciseness.
-
-        Do not call this method directly in Python; call cast() instead.
-        """
-        pylist = self.to_pylist()
-        return Series.from_pylist(pylist, self.name(), dtype=DataType.python())
-
     def _pycast_to_pynative(self, typefn: type, dtype: PyDataType) -> Series:
         """Apply Python-level casting to this Series.
 
@@ -290,12 +243,7 @@ def to_arrow(self) -> pa.Array:
 
     def to_pylist(self) -> list[Any]:
         """Convert this Series to a Python list."""
-        if self.datatype().is_python():
-            return self._series.to_pylist()
-        elif self.datatype()._should_cast_to_python():
-            return self._series.cast(DataType.python()._dtype).to_pylist()
-        else:
-            return self._series.to_arrow().to_pylist()
+        return self._series.to_pylist()
 
     def filter(self, mask: Series) -> Series:
         if not isinstance(mask, Series):
@@ -339,9 +287,12 @@ def murmur3_32(self) -> Series:
     def __repr__(self) -> str:
         return repr(self._series)
 
+    def __getitem__(self, index: int) -> Any:
+        return self._series[index]
+
     def __bool__(self) -> bool:
         raise ValueError(
-            "Series don't have a truth value." "If you reached this error using `and` / `or`, use `&` / `|` instead."
+            "Series don't have a truth value. If you reached this error using `and` / `or`, use `&` / `|` instead."
         )
 
     def __len__(self) -> int:

daft/series.py

@@ -56,11 +56,11 @@
     * [Expressions](api/expressions.md)
     * [Functions](api/functions/)
     * [User-Defined Functions](api/udf.md)
+    * [Data Types](api/datatypes.md)
     * [Window](api/window.md)
     * [Sessions](api/sessions.md)
     * [Catalogs & Tables](api/catalogs_tables.md)
     * [Schema](api/schema.md)
-    * [Data Types](api/datatypes.md)
     * [Aggregations](api/aggregations.md)
     * [Series](api/series.md)
     * [Spark Connect](api/spark_connect.md)

docs/SUMMARY.md

@@ -1,5 +1,42 @@
 # DataTypes
 
+## Type Conversions
+
+### Daft to Python
+
+<!-- Note: the conversions here should match the behavior of the Rust `impl IntoPyObject for Literal`: `src/daft-core/src-lit/python.rs` -->
+
+This table shows the mapping from Daft DataTypes to Python types, as done in places such as [`Series.to_pylist`][daft.series.Series.to_pylist], [`Expression.cast`][daft.expressions.Expression.cast] to Python type, and arguments passed into functions decorated with `@daft.func`.
+
+| Daft DataType                                                        | Python Type                                                                         |
+|----------------------------------------------------------------------|-------------------------------------------------------------------------------------|
+| Null                                                                 | `None`                                                                              |
+| Boolean                                                              | `bool`                                                                              |
+| Utf8                                                                 | `str`                                                                               |
+| Binary, FixedSizeBinary                                              | `bytes`                                                                             |
+| Int8, Uint8, Int16, UInt16, Int32, UInt32, Int64, UInt64             | `int`                                                                               |
+| Timestamp                                                            | `datetime.datetime`                                                                 |
+| Date                                                                 | `datetime.date`                                                                     |
+| Time                                                                 | `datetime.time`                                                                     |
+| Duration                                                             | `datetime.timedelta`                                                                |
+| Interval                                                             | not supported                                                                       |
+| Float32, Float64                                                     | `float`                                                                             |
+| Decimal                                                              | `decimal.Decimal`                                                                   |
+| List[T], FixedSizeList[T]                                            | `list[T]`                                                                           |
+| Struct \{<field_1\>: T1, <field_2\>: T2, ...\}                       | `dict[str, T1 | T2 | ...]`                                                          |
+| Map[K, V]                                                            | `dict[K, V]`                                                                        |
+| Tensor[T], FixedShapeTensor[T]                                       | `numpy.typing.NDArray[T]`                                                           |
+| SparseTensor[T], FixedShapeSparseTensor[T]                           | `{`<br>`"values": T,`<br>`"indices": list[int],`<br>`"shape": list[int]`<br>`}`     |
+| Embedding[T]                                                         | `numpy.typing.NDArray[T]`                                                           |
+| Image                                                                | `numpy.typing.NDArray[numpy.uint8 | numpy.uint16 | numpy.float32]`                  |
+| Python                                                               | `Any`                                                                               |
+| Extension[T]                                                         | `T`                                                                                 |
+
+### Python to Daft
+TODO
+
+## daft.DataType
+
 Daft provides simple DataTypes that are ubiquituous in many DataFrames such as numbers, strings and dates - all the way up to more complex types like tensors and images. Learn more about [DataTypes](../core_concepts.md#datatypes) in Daft User Guide.
 
 ::: daft.datatype.DataType

docs/api/datatypes.md

@@ -2,6 +2,7 @@
 arrow2 = {workspace = true}
 bincode = {workspace = true}
 common-error = {path = "../error", default-features = false}
+common-ndarray = {workspace = true}
 daft-schema = {path = "../../daft-schema", default-features = false}
 image = {workspace = true, features = [
   "gif",
@@ -13,6 +14,7 @@ image = {workspace = true, features = [
   "bmp",
   "hdr"
 ]}
+ndarray = {workspace = true}
 serde = {workspace = true}
 
 [lints]

src/common/image/Cargo.toml

@@ -1,13 +1,47 @@
 use core::slice;
 use std::{hash::Hash, ops::Deref};
 
-use image::{DynamicImage, ImageBuffer};
+use common_ndarray::NdArray;
+use image::{flat::SampleLayout, DynamicImage, ImageBuffer, Pixel};
+use ndarray::{Array3, ShapeBuilder};
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 
 /// Wrapper around image::DynamicImage to implement certain traits
 #[derive(Debug, Clone, PartialEq)]
 pub struct Image(pub DynamicImage);
 
+impl Image {
+    pub fn into_ndarray(self) -> Box<dyn NdArray> {
+        fn into_ndarray3<P: Pixel>(buf: ImageBuffer<P, Vec<P::Subpixel>>) -> Array3<P::Subpixel> {
+            let SampleLayout {
+                channels,
+                channel_stride,
+                height,
+                height_stride,
+                width,
+                width_stride,
+            } = buf.sample_layout();
+            let shape = (height as usize, width as usize, channels as usize);
+            let strides = (height_stride, width_stride, channel_stride);
+            Array3::from_shape_vec(shape.strides(strides), buf.into_raw()).unwrap()
+        }
+
+        match self.0 {
+            DynamicImage::ImageLuma8(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageLumaA8(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageRgb8(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageRgba8(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageLuma16(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageLumaA16(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageRgb16(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageRgba16(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageRgb32F(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            DynamicImage::ImageRgba32F(buf) => Box::new(into_ndarray3(buf).into_dyn()),
+            _ => unimplemented!("unsupported DynamicImage variant"),
+        }
+    }
+}
+
 impl Deref for Image {
     type Target = DynamicImage;
 
@@ -51,7 +85,7 @@ impl Hash for Image {
                 };
                 buffer_slice.hash(state);
             }
-            _ => todo!(),
+            _ => unimplemented!("unsupported DynamicImage variant"),
         }
     }
 }

src/common/image/src/image.rs

@@ -0,0 +1,15 @@
+[dependencies]
+ndarray = {workspace = true}
+numpy = {workspace = true, optional = true}
+pyo3 = {workspace = true, optional = true}
+
+[features]
+python = ["dep:pyo3", "dep:numpy"]
+
+[lints]
+workspace = true
+
+[package]
+edition = {workspace = true}
+name = "common-ndarray"
+version = {workspace = true}

src/common/ndarray/Cargo.toml

@@ -0,0 +1,24 @@
+use ndarray::ArrayD;
+#[cfg(feature = "python")]
+use pyo3::{Bound, PyAny, Python};
+
+/// Trait to allow dynamic dispatch over ndarray arrays of any type.
+pub trait NdArray {
+    #[cfg(feature = "python")]
+    fn into_py(self: Box<Self>, py: Python) -> Bound<PyAny>;
+}
+
+#[cfg(not(feature = "python"))]
+impl<A> NdArray for ArrayD<A> {}
+
+#[cfg(feature = "python")]
+impl<A> NdArray for ArrayD<A>
+where
+    A: numpy::Element,
+{
+    fn into_py(self: Box<Self>, py: Python) -> Bound<PyAny> {
+        use numpy::IntoPyArray;
+
+        self.into_pyarray(py).into_any()
+    }
+}