implement Datum for PyScalar and PyArray (#173)

kylebarron · Aug 27, 2024 · e388617 · e388617
1 parent 3b1f8d9
commit e388617
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/pyo3-arrow/src/array.rs b/pyo3-arrow/src/array.rs
@@ -8,7 +8,7 @@ use arrow::datatypes::{
     UInt64Type, UInt8Type,
 };
 use arrow_array::{
-    Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, LargeBinaryArray,
+    Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Datum, LargeBinaryArray,
     LargeStringArray, PrimitiveArray, StringArray, StringViewArray,
 };
 use arrow_buffer::ScalarBuffer;
@@ -139,6 +139,12 @@ impl Display for PyArray {
     }
 }
 
+impl Datum for PyArray {
+    fn get(&self) -> (&dyn Array, bool) {
+        (self.array.as_ref(), false)
+    }
+}
+
 #[pymethods]
 impl PyArray {
     #[new]

diff --git a/pyo3-arrow/src/scalar.rs b/pyo3-arrow/src/scalar.rs
@@ -5,7 +5,7 @@ use std::sync::Arc;
 use arrow::array::AsArray;
 use arrow::datatypes::*;
 use arrow_array::timezone::Tz;
-use arrow_array::{Array, ArrayRef, UnionArray};
+use arrow_array::{Array, ArrayRef, Datum, UnionArray};
 use arrow_schema::{ArrowError, DataType, FieldRef};
 use indexmap::IndexMap;
 use pyo3::prelude::*;
@@ -65,6 +65,12 @@ impl Display for PyScalar {
     }
 }
 
+impl Datum for PyScalar {
+    fn get(&self) -> (&dyn Array, bool) {
+        (self.array.as_ref(), true)
+    }
+}
+
 #[pymethods]
 impl PyScalar {
     fn __repr__(&self) -> String {