throw arrow error instead of panic

apache · Sep 25, 2024 · 72759ab · 72759ab
1 parent e67f17e
commit 72759ab
Showing 1 changed file with 44 additions and 12 deletions.
diff --git a/arrow-cast/src/cast/dictionary.rs b/arrow-cast/src/cast/dictionary.rs
@@ -91,12 +91,17 @@ pub(crate) fn dictionary_cast<K: ArrowDictionaryKeyType>(
             let dict_array = array
                 .as_dictionary::<K>()
                 .downcast_dict::<StringArray>()
-                .unwrap();
+                .ok_or_else(|| {
+                    ArrowError::ComputeError(
+                        "Internal Error: Cannot cast Utf8View to StringArray of expected type"
+                            .to_string(),
+                    )
+                })?;
 
             let string_view = view_from_dict_values::<K, StringViewType, GenericStringType<i32>>(
                 dict_array.values(),
                 dict_array.keys(),
-            );
+            )?;
             Ok(Arc::new(string_view))
         }
         BinaryView => {
@@ -105,12 +110,17 @@ pub(crate) fn dictionary_cast<K: ArrowDictionaryKeyType>(
             let dict_array = array
                 .as_dictionary::<K>()
                 .downcast_dict::<BinaryArray>()
-                .unwrap();
+                .ok_or_else(|| {
+                    ArrowError::ComputeError(
+                        "Internal Error: Cannot cast BinaryView to BinaryArray of expected type"
+                            .to_string(),
+                    )
+                })?;
 
             let binary_view = view_from_dict_values::<K, BinaryViewType, BinaryType>(
                 dict_array.values(),
                 dict_array.keys(),
-            );
+            )?;
             Ok(Arc::new(binary_view))
         }
         _ => unpack_dictionary::<K>(array, to_type, cast_options),
@@ -120,15 +130,17 @@ pub(crate) fn dictionary_cast<K: ArrowDictionaryKeyType>(
 fn view_from_dict_values<K: ArrowDictionaryKeyType, T: ByteViewType, V: ByteArrayType>(
     array: &GenericByteArray<V>,
     keys: &PrimitiveArray<K>,
-) -> GenericByteViewArray<T> {
+) -> Result<GenericByteViewArray<T>, ArrowError> {
     let value_buffer = array.values();
     let value_offsets = array.value_offsets();
     let mut builder = GenericByteViewBuilder::<T>::with_capacity(keys.len());
     builder.append_block(value_buffer.clone());
     for i in keys.iter() {
         match i {
             Some(v) => {
-                let idx = v.to_usize().unwrap();
+                let idx = v.to_usize().ok_or_else(|| {
+                    ArrowError::ComputeError("Invalid dictionary index".to_string())
+                })?;
 
                 // Safety
                 // (1) The index is within bounds as they are offsets
@@ -145,7 +157,7 @@ fn view_from_dict_values<K: ArrowDictionaryKeyType, T: ByteViewType, V: ByteArra
             }
         }
     }
-    builder.finish()
+    Ok(builder.finish())
 }
 
 // Unpack a dictionary where the keys are of type <K> into a flattened array of type to_type
@@ -211,7 +223,11 @@ pub(crate) fn cast_to_dictionary<K: ArrowDictionaryKeyType>(
             let dict = dict
                 .as_dictionary::<K>()
                 .downcast_dict::<Decimal128Array>()
-                .unwrap();
+                .ok_or_else(|| {
+                    ArrowError::ComputeError(
+                        "Internal Error: Cannot cast dict to Decimal128Array".to_string(),
+                    )
+                })?;
             let value = dict.values().clone();
             // Set correct precision/scale
             let value = value.with_precision_and_scale(p, s)?;
@@ -229,7 +245,11 @@ pub(crate) fn cast_to_dictionary<K: ArrowDictionaryKeyType>(
             let dict = dict
                 .as_dictionary::<K>()
                 .downcast_dict::<Decimal256Array>()
-                .unwrap();
+                .ok_or_else(|| {
+                    ArrowError::ComputeError(
+                        "Internal Error: Cannot cast dict to Decimal256Array".to_string(),
+                    )
+                })?;
             let value = dict.values().clone();
             // Set correct precision/scale
             let value = value.with_precision_and_scale(p, s)?;
@@ -350,7 +370,12 @@ where
         1024,
         1024,
     );
-    let string_view = array.as_any().downcast_ref::<StringViewArray>().unwrap();
+    let string_view = array
+        .as_any()
+        .downcast_ref::<StringViewArray>()
+        .ok_or_else(|| {
+            ArrowError::ComputeError("Internal Error: Cannot cast to StringViewArray".to_string())
+        })?;
     for v in string_view.iter() {
         match v {
             Some(v) => {
@@ -376,7 +401,12 @@ where
         1024,
         1024,
     );
-    let binary_view = array.as_any().downcast_ref::<BinaryViewArray>().unwrap();
+    let binary_view = array
+        .as_any()
+        .downcast_ref::<BinaryViewArray>()
+        .ok_or_else(|| {
+            ArrowError::ComputeError("Internal Error: Cannot cast to BinaryViewArray".to_string())
+        })?;
     for v in binary_view.iter() {
         match v {
             Some(v) => {
@@ -405,7 +435,9 @@ where
     let values = cast_values
         .as_any()
         .downcast_ref::<GenericByteArray<T>>()
-        .unwrap();
+        .ok_or_else(|| {
+            ArrowError::ComputeError("Internal Error: Cannot cast to GenericByteArray".to_string())
+        })?;
     let mut b = GenericByteDictionaryBuilder::<K, T>::with_capacity(values.len(), 1024, 1024);
 
     // copy each element one at a time