Add AsArray trait for more ergonomic downcasting

arrow-arith/src/aggregate.rs

@@ -1221,7 +1221,7 @@ mod tests {
             .into_iter()
             .collect();
         let sliced_input = sliced_input.slice(4, 2);
-        let sliced_input = as_primitive_array::<Float64Type>(&sliced_input);
+        let sliced_input = sliced_input.as_primitive::<Float64Type>();
 
         assert_eq!(sliced_input, &input);
 
@@ -1244,7 +1244,7 @@ mod tests {
             .into_iter()
             .collect();
         let sliced_input = sliced_input.slice(4, 2);
-        let sliced_input = as_boolean_array(&sliced_input);
+        let sliced_input = sliced_input.as_boolean();
 
         assert_eq!(sliced_input, &input);
 
@@ -1267,7 +1267,7 @@ mod tests {
             .into_iter()
             .collect();
         let sliced_input = sliced_input.slice(4, 2);
-        let sliced_input = as_string_array(&sliced_input);
+        let sliced_input = sliced_input.as_string::<i32>();
 
         assert_eq!(sliced_input, &input);
 
@@ -1290,7 +1290,7 @@ mod tests {
             .into_iter()
             .collect();
         let sliced_input = sliced_input.slice(4, 2);
-        let sliced_input = as_generic_binary_array::<i32>(&sliced_input);
+        let sliced_input = sliced_input.as_binary::<i32>();
 
         assert_eq!(sliced_input, &input);
 

arrow-arith/src/arithmetic.rs

@@ -728,20 +728,20 @@ pub fn add_dyn(left: &dyn Array, right: &dyn Array) -> Result<ArrayRef, ArrowErr
             typed_dict_math_op!(left, right, |a, b| a.add_wrapping(b), math_op_dict)
         }
         DataType::Date32 => {
-            let l = as_primitive_array::<Date32Type>(left);
+            let l = left.as_primitive::<Date32Type>();
             match right.data_type() {
                 DataType::Interval(IntervalUnit::YearMonth) => {
-                    let r = as_primitive_array::<IntervalYearMonthType>(right);
+                    let r = right.as_primitive::<IntervalYearMonthType>();
                     let res = math_op(l, r, Date32Type::add_year_months)?;
                     Ok(Arc::new(res))
                 }
                 DataType::Interval(IntervalUnit::DayTime) => {
-                    let r = as_primitive_array::<IntervalDayTimeType>(right);
+                    let r = right.as_primitive::<IntervalDayTimeType>();
                     let res = math_op(l, r, Date32Type::add_day_time)?;
                     Ok(Arc::new(res))
                 }
                 DataType::Interval(IntervalUnit::MonthDayNano) => {
-                    let r = as_primitive_array::<IntervalMonthDayNanoType>(right);
+                    let r = right.as_primitive::<IntervalMonthDayNanoType>();
                     let res = math_op(l, r, Date32Type::add_month_day_nano)?;
                     Ok(Arc::new(res))
                 }
@@ -752,20 +752,20 @@ pub fn add_dyn(left: &dyn Array, right: &dyn Array) -> Result<ArrayRef, ArrowErr
             }
         }
         DataType::Date64 => {
-            let l = as_primitive_array::<Date64Type>(left);
+            let l = left.as_primitive::<Date64Type>();
             match right.data_type() {
                 DataType::Interval(IntervalUnit::YearMonth) => {
-                    let r = as_primitive_array::<IntervalYearMonthType>(right);
+                    let r = right.as_primitive::<IntervalYearMonthType>();
                     let res = math_op(l, r, Date64Type::add_year_months)?;
                     Ok(Arc::new(res))
                 }
                 DataType::Interval(IntervalUnit::DayTime) => {
-                    let r = as_primitive_array::<IntervalDayTimeType>(right);
+                    let r = right.as_primitive::<IntervalDayTimeType>();
                     let res = math_op(l, r, Date64Type::add_day_time)?;
                     Ok(Arc::new(res))
                 }
                 DataType::Interval(IntervalUnit::MonthDayNano) => {
-                    let r = as_primitive_array::<IntervalMonthDayNanoType>(right);
+                    let r = right.as_primitive::<IntervalMonthDayNanoType>();
                     let res = math_op(l, r, Date64Type::add_month_day_nano)?;
                     Ok(Arc::new(res))
                 }
@@ -808,20 +808,20 @@ pub fn add_dyn_checked(
             )
         }
         DataType::Date32 => {
-            let l = as_primitive_array::<Date32Type>(left);
+            let l = left.as_primitive::<Date32Type>();
             match right.data_type() {
                 DataType::Interval(IntervalUnit::YearMonth) => {
-                    let r = as_primitive_array::<IntervalYearMonthType>(right);
+                    let r = right.as_primitive::<IntervalYearMonthType>();
                     let res = math_op(l, r, Date32Type::add_year_months)?;
                     Ok(Arc::new(res))
                 }
                 DataType::Interval(IntervalUnit::DayTime) => {
-                    let r = as_primitive_array::<IntervalDayTimeType>(right);
+                    let r = right.as_primitive::<IntervalDayTimeType>();
                     let res = math_op(l, r, Date32Type::add_day_time)?;
                     Ok(Arc::new(res))
                 }
                 DataType::Interval(IntervalUnit::MonthDayNano) => {
-                    let r = as_primitive_array::<IntervalMonthDayNanoType>(right);
+                    let r = right.as_primitive::<IntervalMonthDayNanoType>();
                     let res = math_op(l, r, Date32Type::add_month_day_nano)?;
                     Ok(Arc::new(res))
                 }
@@ -832,20 +832,20 @@ pub fn add_dyn_checked(
             }
         }
         DataType::Date64 => {
-            let l = as_primitive_array::<Date64Type>(left);
+            let l = left.as_primitive::<Date64Type>();
             match right.data_type() {
                 DataType::Interval(IntervalUnit::YearMonth) => {
-                    let r = as_primitive_array::<IntervalYearMonthType>(right);
+                    let r = right.as_primitive::<IntervalYearMonthType>();
                     let res = math_op(l, r, Date64Type::add_year_months)?;
                     Ok(Arc::new(res))
                 }
                 DataType::Interval(IntervalUnit::DayTime) => {
-                    let r = as_primitive_array::<IntervalDayTimeType>(right);
+                    let r = right.as_primitive::<IntervalDayTimeType>();
                     let res = math_op(l, r, Date64Type::add_day_time)?;
                     Ok(Arc::new(res))
                 }
                 DataType::Interval(IntervalUnit::MonthDayNano) => {
-                    let r = as_primitive_array::<IntervalMonthDayNanoType>(right);
+                    let r = right.as_primitive::<IntervalMonthDayNanoType>();
                     let res = math_op(l, r, Date64Type::add_month_day_nano)?;
                     Ok(Arc::new(res))
                 }
@@ -2079,8 +2079,7 @@ mod tests {
     fn test_primitive_array_add_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let a = as_primitive_array(&a);
-        let actual = add_scalar(a, 3).unwrap();
+        let actual = add_scalar(a.as_primitive(), 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(12), Some(11), None]);
         assert_eq!(actual, expected);
     }
@@ -2110,8 +2109,7 @@ mod tests {
     fn test_primitive_array_subtract_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let a = as_primitive_array(&a);
-        let actual = subtract_scalar(a, 3).unwrap();
+        let actual = subtract_scalar(a.as_primitive(), 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(6), Some(5), None]);
         assert_eq!(actual, expected);
     }
@@ -2141,8 +2139,7 @@ mod tests {
     fn test_primitive_array_multiply_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let a = as_primitive_array(&a);
-        let actual = multiply_scalar(a, 3).unwrap();
+        let actual = multiply_scalar(a.as_primitive(), 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(27), Some(24), None]);
         assert_eq!(actual, expected);
     }
@@ -2171,7 +2168,7 @@ mod tests {
         assert_eq!(0, c.value(4));
 
         let c = modulus_dyn(&a, &b).unwrap();
-        let c = as_primitive_array::<Int32Type>(&c);
+        let c = c.as_primitive::<Int32Type>();
         assert_eq!(0, c.value(0));
         assert_eq!(3, c.value(1));
         assert_eq!(0, c.value(2));
@@ -2262,8 +2259,7 @@ mod tests {
     fn test_primitive_array_divide_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let a = as_primitive_array(&a);
-        let actual = divide_scalar(a, 3).unwrap();
+        let actual = divide_scalar(a.as_primitive(), 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(3), Some(2), None]);
         assert_eq!(actual, expected);
     }
@@ -2277,7 +2273,7 @@ mod tests {
         assert_eq!(c, expected);
 
         let c = modulus_scalar_dyn::<Int32Type>(&a, b).unwrap();
-        let c = as_primitive_array::<Int32Type>(&c);
+        let c = c.as_primitive::<Int32Type>();
         let expected = Int32Array::from(vec![0, 2, 0, 2, 1]);
         assert_eq!(c, &expected);
     }
@@ -2286,13 +2282,13 @@ mod tests {
     fn test_int_array_modulus_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let a = as_primitive_array(&a);
+        let a = a.as_primitive();
         let actual = modulus_scalar(a, 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(0), Some(2), None]);
         assert_eq!(actual, expected);
 
         let actual = modulus_scalar_dyn::<Int32Type>(a, 3).unwrap();
-        let actual = as_primitive_array::<Int32Type>(&actual);
+        let actual = actual.as_primitive::<Int32Type>();
         let expected = Int32Array::from(vec![None, Some(0), Some(2), None]);
         assert_eq!(actual, &expected);
     }
@@ -2313,7 +2309,7 @@ mod tests {
         assert_eq!(0, result.value(0));
 
         let result = modulus_scalar_dyn::<Int32Type>(&a, -1).unwrap();
-        let result = as_primitive_array::<Int32Type>(&result);
+        let result = result.as_primitive::<Int32Type>();
         assert_eq!(0, result.value(0));
     }
 
@@ -3295,7 +3291,8 @@ mod tests {
             .unwrap();
 
         let result = add_scalar_dyn::<Decimal128Type>(&a, 1).unwrap();
-        let result = as_primitive_array::<Decimal128Type>(&result)
+        let result = result
+            .as_primitive::<Decimal128Type>()
             .clone()
             .with_precision_and_scale(38, 2)
             .unwrap();

arrow-arith/src/arity.rs

@@ -523,7 +523,7 @@ mod tests {
         let input =
             Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
         let input_slice = input.slice(1, 4);
-        let input_slice: &Float64Array = as_primitive_array(&input_slice);
+        let input_slice: &Float64Array = input_slice.as_primitive();
         let result = unary(input_slice, |n| n.round());
         assert_eq!(
             result,

arrow-array/src/array/dictionary_array.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::builder::{PrimitiveDictionaryBuilder, StringDictionaryBuilder};
-use crate::cast::as_primitive_array;
+use crate::cast::AsArray;
 use crate::iterator::ArrayIter;
 use crate::types::*;
 use crate::{
@@ -410,8 +410,8 @@ impl<K: ArrowDictionaryKeyType> DictionaryArray<K> {
             return Err(self);
         }
 
-        let key_array = as_primitive_array::<K>(self.keys()).clone();
-        let value_array = as_primitive_array::<V>(self.values()).clone();
+        let key_array = self.keys().clone();
+        let value_array = self.values().as_primitive::<V>().clone();
 
         drop(self.data);
         drop(self.keys);

arrow-array/src/array/map_array.rs

@@ -253,7 +253,7 @@ impl std::fmt::Debug for MapArray {
 
 #[cfg(test)]
 mod tests {
-    use crate::cast::as_primitive_array;
+    use crate::cast::AsArray;
     use crate::types::UInt32Type;
     use crate::{Int32Array, UInt32Array};
     use std::sync::Arc;
@@ -522,7 +522,7 @@ mod tests {
 
         assert_eq!(
             &values_data,
-            as_primitive_array::<UInt32Type>(map_array.values())
+            map_array.values().as_primitive::<UInt32Type>()
         );
         assert_eq!(&DataType::UInt32, map_array.value_type());
         assert_eq!(3, map_array.len());

arrow-array/src/array/run_array.rs

@@ -555,7 +555,7 @@ mod tests {
 
     use super::*;
     use crate::builder::PrimitiveRunBuilder;
-    use crate::cast::as_primitive_array;
+    use crate::cast::AsArray;
     use crate::types::{Int16Type, Int32Type, Int8Type, UInt32Type};
     use crate::{Array, Int32Array, StringArray};
 
@@ -877,8 +877,7 @@ mod tests {
             builder.extend(input_array.clone().into_iter());
 
             let run_array = builder.finish();
-            let physical_values_array =
-                as_primitive_array::<Int32Type>(run_array.values());
+            let physical_values_array = run_array.values().as_primitive::<Int32Type>();
 
             // create an array consisting of all the indices repeated twice and shuffled.
             let mut logical_indices: Vec<u32> = (0_u32..(logical_len as u32)).collect();
@@ -913,7 +912,7 @@ mod tests {
             PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
         builder.extend(input_array.iter().copied());
         let run_array = builder.finish();
-        let physical_values_array = as_primitive_array::<Int32Type>(run_array.values());
+        let physical_values_array = run_array.values().as_primitive::<Int32Type>();
 
         // test for all slice lengths.
         for slice_len in 1..=total_len {

arrow-array/src/array/union_array.rs

@@ -398,7 +398,7 @@ mod tests {
     use super::*;
 
     use crate::builder::UnionBuilder;
-    use crate::cast::{as_primitive_array, as_string_array};
+    use crate::cast::AsArray;
     use crate::types::{Float32Type, Float64Type, Int32Type, Int64Type};
     use crate::RecordBatch;
     use crate::{Float64Array, Int32Array, Int64Array, StringArray};
@@ -1078,36 +1078,36 @@ mod tests {
         let v = array.value(0);
         assert_eq!(v.data_type(), &DataType::Int32);
         assert_eq!(v.len(), 1);
-        assert_eq!(as_primitive_array::<Int32Type>(v.as_ref()).value(0), 5);
+        assert_eq!(v.as_primitive::<Int32Type>().value(0), 5);
 
         let v = array.value(1);
         assert_eq!(v.data_type(), &DataType::Utf8);
         assert_eq!(v.len(), 1);
-        assert_eq!(as_string_array(v.as_ref()).value(0), "foo");
+        assert_eq!(v.as_string::<i32>().value(0), "foo");
 
         let v = array.value(2);
         assert_eq!(v.data_type(), &DataType::Int32);
         assert_eq!(v.len(), 1);
-        assert_eq!(as_primitive_array::<Int32Type>(v.as_ref()).value(0), 6);
+        assert_eq!(v.as_primitive::<Int32Type>().value(0), 6);
 
         let v = array.value(3);
         assert_eq!(v.data_type(), &DataType::Utf8);
         assert_eq!(v.len(), 1);
-        assert_eq!(as_string_array(v.as_ref()).value(0), "bar");
+        assert_eq!(v.as_string::<i32>().value(0), "bar");
 
         let v = array.value(4);
         assert_eq!(v.data_type(), &DataType::Float64);
         assert_eq!(v.len(), 1);
-        assert_eq!(as_primitive_array::<Float64Type>(v.as_ref()).value(0), 10.0);
+        assert_eq!(v.as_primitive::<Float64Type>().value(0), 10.0);
 
         let v = array.value(5);
         assert_eq!(v.data_type(), &DataType::Int32);
         assert_eq!(v.len(), 1);
-        assert_eq!(as_primitive_array::<Int32Type>(v.as_ref()).value(0), 4);
+        assert_eq!(v.as_primitive::<Int32Type>().value(0), 4);
 
         let v = array.value(6);
         assert_eq!(v.data_type(), &DataType::Utf8);
         assert_eq!(v.len(), 1);
-        assert_eq!(as_string_array(v.as_ref()).value(0), "baz");
+        assert_eq!(v.as_string::<i32>().value(0), "baz");
     }
 }