Fix bug with type to pyarrow type

Eventual-Inc · Sep 9, 2024 · f741a64 · f741a64
1 parent 8add10e
commit f741a64
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/src/daft-core/src/ffi.rs b/src/daft-core/src/ffi.rs
@@ -69,7 +69,7 @@ pub fn field_to_py(
     Ok(field.to_object(py))
 }
 
-pub fn to_py_schema(
+pub fn dtype_to_py(
     dtype: &arrow2::datatypes::DataType,
     py: Python,
     pyarrow: &PyModule,
@@ -81,8 +81,9 @@ pub fn to_py_schema(
         pyo3::intern!(py, "_import_from_c"),
         (schema_ptr as Py_uintptr_t,),
     )?;
+    let dtype = field.getattr(pyo3::intern!(py, "type"))?.to_object(py);
 
-    Ok(field.to_object(py))
+    Ok(dtype.to_object(py))
 }
 
 fn fix_child_array_slice_offsets(array: ArrayRef) -> ArrayRef {

diff --git a/src/daft-core/src/python/datatype.rs b/src/daft-core/src/python/datatype.rs
@@ -331,11 +331,10 @@ impl PyDataType {
                 } else {
                     // Fall back to default Daft super extension representation if installed pyarrow doesn't have the
                     // canonical tensor extension type.
-                    ffi::to_py_schema(&self.dtype.to_arrow()?, py, pyarrow)?
+                    ffi::dtype_to_py(&self.dtype.to_arrow()?, py, pyarrow)?
                 },
             ),
-            _ => ffi::to_py_schema(&self.dtype.to_arrow()?, py, pyarrow)?
-                .getattr(py, pyo3::intern!(py, "type")),
+            _ => ffi::dtype_to_py(&self.dtype.to_arrow()?, py, pyarrow),
         }
     }