diff --git a/src/daft-core/src/python/schema.rs b/src/daft-core/src/python/schema.rs index 33fce18df0..eaa2ff78ac 100644 --- a/src/daft-core/src/python/schema.rs +++ b/src/daft-core/src/python/schema.rs @@ -7,7 +7,6 @@ use serde::{Deserialize, Serialize}; use super::datatype::PyDataType; use super::field::PyField; use crate::datatypes; -use crate::ffi::field_to_py; use crate::schema; use common_py_serde::impl_bincode_py_state_serialization; @@ -29,7 +28,16 @@ impl PySchema { .schema .fields .iter() - .map(|(_, f)| field_to_py(&f.to_arrow()?, py, pyarrow)) + .map(|(_, f)| { + // NOTE: Use PyDataType::to_arrow because we need to dip into Python to get + // the registered Arrow extension types + let py_dtype: PyDataType = f.dtype.clone().into(); + let py_arrow_dtype = py_dtype.to_arrow(py)?; + pyarrow + .getattr(pyo3::intern!(py, "field")) + .unwrap() + .call1((f.name.clone(), py_arrow_dtype)) + }) .collect::>>()?; pyarrow .getattr(pyo3::intern!(py, "schema"))