Fix min/max with datetimes with timezones

arro3-compute/src/aggregate.rs

@@ -6,7 +6,7 @@ use arrow::array::{
 };
 use arrow::{compute, downcast_primitive_array};
 use arrow_array::{
-    ArrayRef, ArrowPrimitiveType, BinaryViewArray, BooleanArray, GenericBinaryArray,
+    Array, ArrayRef, ArrowPrimitiveType, BinaryViewArray, BooleanArray, GenericBinaryArray,
     GenericStringArray, OffsetSizeTrait, PrimitiveArray, StringViewArray,
 };
 use arrow_schema::{ArrowError, DataType};
@@ -71,7 +71,9 @@ fn max_array(array: ArrayRef) -> Result<ArrayRef, ArrowError> {
 fn max_primitive<T: ArrowPrimitiveType>(array: &PrimitiveArray<T>) -> ArrayRef {
     let mut builder = PrimitiveBuilder::<T>::with_capacity(1);
     builder.append_option(compute::max(array));
-    Arc::new(builder.finish())
+    // Need to append the original data type, because PrimitiveBuilder::<T> will sometimes lose the
+    // exact data type. E.g. It will lose the time zone for datetime types.
+    Arc::new(builder.finish().with_data_type(array.data_type().clone()))
 }
 
 #[inline(never)]
@@ -164,7 +166,9 @@ fn min_array(array: ArrayRef) -> Result<ArrayRef, ArrowError> {
 fn min_primitive<T: ArrowPrimitiveType>(array: &PrimitiveArray<T>) -> ArrayRef {
     let mut builder = PrimitiveBuilder::<T>::with_capacity(1);
     builder.append_option(compute::min(array));
-    Arc::new(builder.finish())
+    // Need to append the original data type, because PrimitiveBuilder::<T> will sometimes lose the
+    // exact data type. E.g. It will lose the time zone for datetime types.
+    Arc::new(builder.finish().with_data_type(array.data_type().clone()))
 }
 
 #[inline(never)]

pyo3-arrow/src/array.rs

@@ -50,8 +50,7 @@ impl PyArray {
     ///
     /// This will panic if the array's data type does not match the field's data type.
     pub fn new(array: ArrayRef, field: FieldRef) -> Self {
-        assert_eq!(array.data_type(), field.data_type());
-        Self { array, field }
+        Self::try_new(array, field).unwrap()
     }
 
     /// Create a new Python Array from an [ArrayRef] and a [FieldRef].
@@ -61,7 +60,7 @@ impl PyArray {
         // providing.
         if array.data_type() != field.data_type() {
             return Err(ArrowError::SchemaError(
-                "Array DataType must match Field DataType".to_string(),
+                format!("Array DataType must match Field DataType. Array DataType is {}; field DataType is {}", array.data_type(), field.data_type())
             ));
         }
         Ok(Self { array, field })

tests/compute/test_aggregate.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timezone
 
 import arro3.compute as ac
 import pyarrow as pa
@@ -53,3 +53,16 @@ def test_min_max_datetime():
     arro3_arr = Array(pa_arr)
     assert ac.min(arro3_arr).as_py() == dt1
     assert ac.max(arro3_arr).as_py() == dt3
+
+
+def test_min_max_datetime_with_timezone():
+    dt1 = datetime.now(timezone.utc)
+    dt2 = datetime.now(timezone.utc)
+    dt3 = datetime.now(timezone.utc)
+    arr = pa.array([dt1, dt2, dt3])
+    assert arr.type.tz == "UTC"
+
+    assert ac.min(arr).as_py() == dt1
+    assert ac.min(arr).type.tz == "UTC"
+    assert ac.max(arr).as_py() == dt3
+    assert ac.max(arr).type.tz == "UTC"