chore: prefer downcast_array_ref (#1786)

danking · web-flow · commit 4eda8900d25d · 2025-01-03T23:22:09.000Z
diff --git a/vortex-array/src/compute/binary_numeric.rs b/vortex-array/src/compute/binary_numeric.rs
@@ -7,7 +7,7 @@ use vortex_scalar::{BinaryNumericOperator, Scalar};
 
 use crate::array::ConstantArray;
 use crate::arrow::{Datum, FromArrowArray};
-use crate::encoding::{downcast_array_ref, Encoding};
+use crate::encoding::Encoding;
 use crate::{ArrayDType, ArrayData, IntoArrayData as _};
 
 pub trait BinaryNumericFn<Array> {
@@ -30,7 +30,7 @@ where
         rhs: &ArrayData,
         op: BinaryNumericOperator,
     ) -> VortexResult<Option<ArrayData>> {
-        let (array_ref, encoding) = downcast_array_ref::<E>(lhs)?;
+        let (array_ref, encoding) = lhs.downcast_array_ref::<E>()?;
         BinaryNumericFn::binary_numeric(encoding, array_ref, rhs, op)
     }
 }
diff --git a/vortex-array/src/compute/boolean.rs b/vortex-array/src/compute/boolean.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use arrow_array::cast::AsArray;
 use arrow_array::ArrayRef;
 use vortex_dtype::DType;
-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
+use vortex_error::{vortex_bail, VortexError, VortexResult};
 
 use crate::arrow::FromArrowArray;
 use crate::encoding::Encoding;
@@ -40,12 +40,7 @@ where
         rhs: &ArrayData,
         op: BinaryOperator,
     ) -> VortexResult<Option<ArrayData>> {
-        let array_ref = <&E::Array>::try_from(lhs)?;
-        let encoding = lhs
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = lhs.downcast_array_ref::<E>()?;
         BinaryBooleanFn::binary_boolean(encoding, array_ref, rhs, op)
     }
 }
diff --git a/vortex-array/src/compute/cast.rs b/vortex-array/src/compute/cast.rs
@@ -1,5 +1,5 @@
 use vortex_dtype::DType;
-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
+use vortex_error::{vortex_bail, VortexError, VortexResult};
 
 use crate::encoding::Encoding;
 use crate::{ArrayDType, ArrayData, IntoArrayData, IntoCanonical};
@@ -14,12 +14,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn cast(&self, array: &ArrayData, dtype: &DType) -> VortexResult<ArrayData> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         CastFn::cast(encoding, array_ref, dtype)
     }
 }
diff --git a/vortex-array/src/compute/compare.rs b/vortex-array/src/compute/compare.rs
@@ -3,7 +3,7 @@ use std::fmt::{Display, Formatter};
 
 use arrow_ord::cmp;
 use vortex_dtype::{DType, Nullability};
-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
+use vortex_error::{vortex_bail, VortexError, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::arrow::{Datum, FromArrowArray};
@@ -92,12 +92,7 @@ where
         rhs: &ArrayData,
         operator: Operator,
     ) -> VortexResult<Option<ArrayData>> {
-        let lhs_ref = <&E::Array>::try_from(lhs)?;
-        let encoding = lhs
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (lhs_ref, encoding) = lhs.downcast_array_ref::<E>()?;
         CompareFn::compare(encoding, lhs_ref, rhs, operator)
     }
 }
diff --git a/vortex-array/src/compute/fill_forward.rs b/vortex-array/src/compute/fill_forward.rs
@@ -18,12 +18,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn fill_forward(&self, array: &ArrayData) -> VortexResult<ArrayData> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         FillForwardFn::fill_forward(encoding, array_ref)
     }
 }
diff --git a/vortex-array/src/compute/fill_null.rs b/vortex-array/src/compute/fill_null.rs
@@ -1,4 +1,4 @@
-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
+use vortex_error::{vortex_bail, VortexError, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::encoding::Encoding;
@@ -17,12 +17,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn fill_null(&self, array: &ArrayData, fill_value: Scalar) -> VortexResult<ArrayData> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         FillNullFn::fill_null(encoding, array_ref, fill_value)
     }
 }
diff --git a/vortex-array/src/compute/filter.rs b/vortex-array/src/compute/filter.rs
@@ -31,12 +31,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn filter(&self, array: &ArrayData, mask: FilterMask) -> VortexResult<ArrayData> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         FilterFn::filter(encoding, array_ref, mask)
     }
 }
diff --git a/vortex-array/src/compute/invert.rs b/vortex-array/src/compute/invert.rs
@@ -1,5 +1,5 @@
 use vortex_dtype::DType;
-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
+use vortex_error::{vortex_bail, VortexError, VortexResult};
 
 use crate::encoding::Encoding;
 use crate::{ArrayDType, ArrayData, IntoArrayData, IntoArrayVariant};
@@ -15,12 +15,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn invert(&self, array: &ArrayData) -> VortexResult<ArrayData> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         InvertFn::invert(encoding, array_ref)
     }
 }
diff --git a/vortex-array/src/compute/like.rs b/vortex-array/src/compute/like.rs
@@ -1,5 +1,5 @@
 use vortex_dtype::DType;
-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
+use vortex_error::{vortex_bail, VortexError, VortexResult};
 
 use crate::arrow::{Datum, FromArrowArray};
 use crate::encoding::Encoding;
@@ -25,12 +25,7 @@ where
         pattern: &ArrayData,
         options: LikeOptions,
     ) -> VortexResult<ArrayData> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         LikeFn::like(encoding, array_ref, pattern, options)
     }
 }
diff --git a/vortex-array/src/compute/scalar_at.rs b/vortex-array/src/compute/scalar_at.rs
@@ -18,12 +18,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn scalar_at(&self, array: &ArrayData, index: usize) -> VortexResult<Scalar> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         ScalarAtFn::scalar_at(encoding, array_ref, index)
     }
 }
diff --git a/vortex-array/src/compute/search_sorted.rs b/vortex-array/src/compute/search_sorted.rs
@@ -4,7 +4,7 @@ use std::fmt::{Debug, Display, Formatter};
 use std::hint;
 
 use itertools::Itertools;
-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
+use vortex_error::{vortex_bail, VortexError, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::compute::scalar_at;
@@ -164,12 +164,7 @@ where
         value: &Scalar,
         side: SearchSortedSide,
     ) -> VortexResult<SearchResult> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         SearchSortedFn::search_sorted(encoding, array_ref, value, side)
     }
 
@@ -179,12 +174,7 @@ where
         values: &[Scalar],
         side: SearchSortedSide,
     ) -> VortexResult<Vec<SearchResult>> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         SearchSortedFn::search_sorted_many(encoding, array_ref, values, side)
     }
 }
@@ -200,12 +190,7 @@ where
         value: usize,
         side: SearchSortedSide,
     ) -> VortexResult<SearchResult> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         SearchSortedUsizeFn::search_sorted_usize(encoding, array_ref, value, side)
     }
 
@@ -215,12 +200,7 @@ where
         values: &[usize],
         side: SearchSortedSide,
     ) -> VortexResult<Vec<SearchResult>> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         SearchSortedUsizeFn::search_sorted_usize_many(encoding, array_ref, values, side)
     }
 }
diff --git a/vortex-array/src/compute/slice.rs b/vortex-array/src/compute/slice.rs
@@ -17,12 +17,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn slice(&self, array: &ArrayData, start: usize, stop: usize) -> VortexResult<ArrayData> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         SliceFn::slice(encoding, array_ref, start, stop)
     }
 }
diff --git a/vortex-array/src/compute/take.rs b/vortex-array/src/compute/take.rs
@@ -32,12 +32,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn take(&self, array: &ArrayData, indices: &ArrayData) -> VortexResult<ArrayData> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         TakeFn::take(encoding, array_ref, indices)
     }
 }
diff --git a/vortex-array/src/data/mod.rs b/vortex-array/src/data/mod.rs
@@ -8,15 +8,15 @@ use owned::OwnedArrayData;
 use viewed::ViewedArrayData;
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
-use vortex_error::{vortex_err, VortexExpect, VortexResult};
+use vortex_error::{vortex_err, VortexError, VortexExpect, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::array::{
     BoolEncoding, ExtensionEncoding, NullEncoding, PrimitiveEncoding, StructEncoding,
     VarBinEncoding, VarBinViewEncoding,
 };
 use crate::compute::scalar_at;
-use crate::encoding::{EncodingId, EncodingRef, EncodingVTable};
+use crate::encoding::{Encoding, EncodingId, EncodingRef, EncodingVTable};
 use crate::iter::{ArrayIterator, ArrayIteratorAdapter};
 use crate::stats::{ArrayStatistics, Stat, Statistics, StatsSet};
 use crate::stream::{ArrayStream, ArrayStreamAdapter};
@@ -371,6 +371,19 @@ impl ArrayData {
             log::warn!("{:?}", bt);
         }
     }
+
+    pub fn downcast_array_ref<E: Encoding>(self: &ArrayData) -> VortexResult<(&E::Array, &E)>
+    where
+        for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
+    {
+        let array_ref = <&E::Array>::try_from(self)?;
+        let encoding = self
+            .encoding()
+            .as_any()
+            .downcast_ref::<E>()
+            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        Ok((array_ref, encoding))
+    }
 }
 
 impl Display for ArrayData {
diff --git a/vortex-array/src/encoding/mod.rs b/vortex-array/src/encoding/mod.rs
@@ -4,8 +4,6 @@ use std::any::Any;
 use std::fmt::{Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
 
-use vortex_error::{vortex_err, VortexError, VortexResult};
-
 use crate::compute::ComputeVTable;
 use crate::stats::StatisticsVTable;
 use crate::validity::ValidityVTable;
@@ -69,19 +67,6 @@ pub trait Encoding: 'static {
     type Metadata: ArrayMetadata;
 }
 
-pub fn downcast_array_ref<E: Encoding>(array: &ArrayData) -> VortexResult<(&E::Array, &E)>
-where
-    for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
-{
-    let array_ref = <&E::Array>::try_from(array)?;
-    let encoding = array
-        .encoding()
-        .as_any()
-        .downcast_ref::<E>()
-        .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
-    Ok((array_ref, encoding))
-}
-
 pub type EncodingRef = &'static dyn EncodingVTable;
 
 /// Object-safe encoding trait for an array.
diff --git a/vortex-array/src/stats/mod.rs b/vortex-array/src/stats/mod.rs
@@ -12,7 +12,7 @@ use num_enum::{IntoPrimitive, TryFromPrimitive};
 pub use statsset::*;
 use vortex_dtype::Nullability::NonNullable;
 use vortex_dtype::{DType, NativePType, PType};
-use vortex_error::{vortex_err, vortex_panic, VortexError, VortexExpect, VortexResult};
+use vortex_error::{vortex_panic, VortexError, VortexExpect, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::encoding::Encoding;
@@ -205,12 +205,7 @@ where
     for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,
 {
     fn compute_statistics(&self, array: &ArrayData, stat: Stat) -> VortexResult<StatsSet> {
-        let array_ref = <&E::Array>::try_from(array)?;
-        let encoding = array
-            .encoding()
-            .as_any()
-            .downcast_ref::<E>()
-            .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
+        let (array_ref, encoding) = array.downcast_array_ref::<E>()?;
         StatisticsVTable::compute_statistics(encoding, array_ref, stat)
     }
 }
diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs
diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs
diff --git a/vortex-array/src/visitor.rs b/vortex-array/src/visitor.rs

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`use vortex_dtype::DType;`
`2`		`-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};`
	`2`	`+use vortex_error::{vortex_bail, VortexError, VortexResult};`
`3`	`3`
`4`	`4`	`use crate::encoding::Encoding;`
`5`	`5`	`use crate::{ArrayDType, ArrayData, IntoArrayData, IntoCanonical};`
`@@ -14,12 +14,7 @@ where`
`14`	`14`	`for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,`
`15`	`15`	`{`
`16`	`16`	`fn cast(&self, array: &ArrayData, dtype: &DType) -> VortexResult<ArrayData> {`
`17`		`- let array_ref = <&E::Array>::try_from(array)?;`
`18`		`- let encoding = array`
`19`		`- .encoding()`
`20`		`- .as_any()`
`21`		`- .downcast_ref::<E>()`
`22`		`- .ok_or_else(\|\| vortex_err!("Mismatched encoding"))?;`
	`17`	`+ let (array_ref, encoding) = array.downcast_array_ref::<E>()?;`
`23`	`18`	`CastFn::cast(encoding, array_ref, dtype)`
`24`	`19`	`}`
`25`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,12 +18,7 @@ where`
`18`	`18`	`for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,`
`19`	`19`	`{`
`20`	`20`	`fn fill_forward(&self, array: &ArrayData) -> VortexResult<ArrayData> {`
`21`		`- let array_ref = <&E::Array>::try_from(array)?;`
`22`		`- let encoding = array`
`23`		`- .encoding()`
`24`		`- .as_any()`
`25`		`- .downcast_ref::<E>()`
`26`		`- .ok_or_else(\|\| vortex_err!("Mismatched encoding"))?;`
	`21`	`+ let (array_ref, encoding) = array.downcast_array_ref::<E>()?;`
`27`	`22`	`FillForwardFn::fill_forward(encoding, array_ref)`
`28`	`23`	`}`
`29`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};`
	`1`	`+use vortex_error::{vortex_bail, VortexError, VortexResult};`
`2`	`2`	`use vortex_scalar::Scalar;`
`3`	`3`
`4`	`4`	`use crate::encoding::Encoding;`
`@@ -17,12 +17,7 @@ where`
`17`	`17`	`for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,`
`18`	`18`	`{`
`19`	`19`	`fn fill_null(&self, array: &ArrayData, fill_value: Scalar) -> VortexResult<ArrayData> {`
`20`		`- let array_ref = <&E::Array>::try_from(array)?;`
`21`		`- let encoding = array`
`22`		`- .encoding()`
`23`		`- .as_any()`
`24`		`- .downcast_ref::<E>()`
`25`		`- .ok_or_else(\|\| vortex_err!("Mismatched encoding"))?;`
	`20`	`+ let (array_ref, encoding) = array.downcast_array_ref::<E>()?;`
`26`	`21`	`FillNullFn::fill_null(encoding, array_ref, fill_value)`
`27`	`22`	`}`
`28`	`23`	`}`
Original file line number	Diff line number	Diff line change
`@@ -31,12 +31,7 @@ where`
`31`	`31`	`for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>,`
`32`	`32`	`{`
`33`	`33`	`fn filter(&self, array: &ArrayData, mask: FilterMask) -> VortexResult<ArrayData> {`
`34`		`- let array_ref = <&E::Array>::try_from(array)?;`
`35`		`- let encoding = array`
`36`		`- .encoding()`
`37`		`- .as_any()`
`38`		`- .downcast_ref::<E>()`
`39`		`- .ok_or_else(\|\| vortex_err!("Mismatched encoding"))?;`
	`34`	`+ let (array_ref, encoding) = array.downcast_array_ref::<E>()?;`
`40`	`35`	`FilterFn::filter(encoding, array_ref, mask)`
`41`	`36`	`}`
`42`	`37`	`}`