TryFromArrayRef support handling cast failures (#2740)

Author: robert3005

vortex-array/src/array/convert.rs

@@ -1,4 +1,3 @@
-use arrow_array::Array;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
@@ -13,12 +12,7 @@ pub trait TryIntoArray {
     fn try_into_array(self) -> VortexResult<ArrayRef>;
 }
 
-/// Trait for converting a type from a Vortex [`ArrayRef`], returning an error if the conversion fails.
-pub trait TryFromArray: Sized {
-    fn try_from_array(array: &dyn Array) -> VortexResult<Self>;
-}
-
 /// Trait for converting a type from a Vortex [`ArrayRef`], returning an error if the conversion fails.
 pub trait TryFromArrayRef: Sized {
-    fn try_from_array(array: ArrayRef) -> VortexResult<Self>;
+    fn try_from_array(array: ArrayRef) -> Result<Self, ArrayRef>;
 }

vortex-array/src/array/mod.rs

@@ -7,7 +7,7 @@ mod validity;
 mod variants;
 mod visitor;
 
-use std::any::{Any, type_name};
+use std::any::Any;
 use std::fmt::{Debug, Display, Formatter};
 use std::sync::Arc;
 
@@ -20,7 +20,7 @@ pub use validity::*;
 pub use variants::*;
 pub use visitor::*;
 use vortex_dtype::DType;
-use vortex_error::{VortexExpect, VortexResult, vortex_err};
+use vortex_error::{VortexExpect, VortexResult};
 use vortex_mask::Mask;
 
 use crate::arrays::{
@@ -228,22 +228,22 @@ impl ToOwned for dyn Array {
 }
 
 impl<A: Array + Clone + 'static> TryFromArrayRef for A {
-    fn try_from_array(array: ArrayRef) -> VortexResult<Self> {
-        Ok(Arc::unwrap_or_clone(
-            array
-                .as_any_arc()
-                .downcast::<A>()
-                .map_err(|_| vortex_err!("Cannot downcast to {}", type_name::<A>()))?,
-        ))
+    fn try_from_array(array: ArrayRef) -> Result<Self, ArrayRef> {
+        let fallback = array.clone();
+        if let Ok(array) = array.as_any_arc().downcast::<A>() {
+            // manually drop the fallback value so `Arc::unwrap_or_clone` doesn't always have to clone
+            drop(fallback);
+            Ok(Arc::unwrap_or_clone(array))
+        } else {
+            Err(fallback)
+        }
     }
 }
 
 impl<A: Array + Clone + 'static> TryFromArrayRef for Arc<A> {
-    fn try_from_array(array: ArrayRef) -> VortexResult<Self> {
-        array
-            .as_any_arc()
-            .downcast::<A>()
-            .map_err(|_| vortex_err!("Cannot downcast to {}", type_name::<A>()))
+    fn try_from_array(array: ArrayRef) -> Result<Self, ArrayRef> {
+        let fallback = array.clone();
+        array.as_any_arc().downcast::<A>().map_err(|_| fallback)
     }
 }
 

vortex-array/src/arrays/varbin/canonical.rs

@@ -1,6 +1,6 @@
 use arrow_schema::DataType;
 use vortex_dtype::DType;
-use vortex_error::VortexResult;
+use vortex_error::{VortexResult, vortex_err};
 
 use crate::arrays::VarBinViewArray;
 use crate::arrays::varbin::VarBinArray;
@@ -21,6 +21,7 @@ impl ArrayCanonicalImpl for VarBinArray {
         };
         VarBinViewArray::try_from_array(ArrayRef::from_arrow(array, nullable))
             .map(Canonical::VarBinView)
+            .map_err(|_| vortex_err!("Array wasn't a VarBinViewArray"))
     }
 }
 

vortex-array/src/arrays/varbinview/mod.rs

@@ -11,7 +11,9 @@ use arrow_buffer::ScalarBuffer;
 use static_assertions::{assert_eq_align, assert_eq_size};
 use vortex_buffer::{Alignment, Buffer, ByteBuffer};
 use vortex_dtype::DType;
-use vortex_error::{VortexExpect, VortexResult, VortexUnwrap, vortex_bail, vortex_panic};
+use vortex_error::{
+    VortexExpect, VortexResult, VortexUnwrap, vortex_bail, vortex_err, vortex_panic,
+};
 use vortex_mask::Mask;
 
 use crate::array::{ArrayCanonicalImpl, ArrayValidityImpl};
@@ -397,6 +399,7 @@ impl VarBinViewArray {
                     &string_view_array,
                     nullability.into(),
                 ))
+                .map_err(|_| vortex_err!("Array was not a VarBinViewArray"))
                 .vortex_expect("StringViewArray to VarBinViewArray downcast")
             }
             DType::Binary(nullability) => {
@@ -408,6 +411,7 @@ impl VarBinViewArray {
                     &binary_view_array,
                     nullability.into(),
                 ))
+                .map_err(|_| vortex_err!("Array was not a VarBinViewArray"))
                 .vortex_expect("BinaryViewArray to VarBinViewArray downcast")
             }
             other => vortex_panic!("VarBinViewArray must be Utf8 or Binary, was {other}"),
@@ -422,6 +426,7 @@ impl VarBinViewArray {
         }
         let array = ArrayRef::from_arrow(&builder.finish(), false);
         VarBinViewArray::try_from_array(array)
+            .map_err(|_| vortex_err!("Array was not a VarBinViewArray"))
             .vortex_expect("VarBinViewArray from StringViewBuilder")
     }
 
@@ -434,6 +439,7 @@ impl VarBinViewArray {
 
         let array = ArrayRef::from_arrow(&builder.finish(), true);
         VarBinViewArray::try_from_array(array)
+            .map_err(|_| vortex_err!("Array was not a VarBinViewArray"))
             .vortex_expect("VarBinViewArray from StringViewBuilder")
     }
 
@@ -445,6 +451,7 @@ impl VarBinViewArray {
         }
         let array = ArrayRef::from_arrow(&builder.finish(), false);
         VarBinViewArray::try_from_array(array)
+            .map_err(|_| vortex_err!("Array was not a VarBinViewArray"))
             .vortex_expect("VarBinViewArray from StringViewBuilder")
     }
 
@@ -456,6 +463,7 @@ impl VarBinViewArray {
         builder.extend(iter);
         let array = ArrayRef::from_arrow(&builder.finish(), true);
         VarBinViewArray::try_from_array(array)
+            .map_err(|_| vortex_err!("Array was not a VarBinViewArray"))
             .vortex_expect("VarBinViewArray from StringViewBuilder")
     }
 }