Fix incorrect downcast in estimated_size_bytes (jorgecarleitao#1351)

ritchie46 · Mar 29, 2023 · 8741bad · 8741bad
1 parent a529955
commit 8741bad
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/src/compute/aggregate/memory.rs b/src/compute/aggregate/memory.rs
@@ -67,7 +67,7 @@ pub fn estimated_bytes_size(array: &dyn Array) -> usize {
                 + validity_size(array.validity())
         }
         FixedSizeList => {
-            let array = array.as_any().downcast_ref::<ListArray<i64>>().unwrap();
+            let array = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
             estimated_bytes_size(array.values().as_ref()) + validity_size(array.validity())
         }
         LargeList => {

diff --git a/tests/it/compute/aggregate/memory.rs b/tests/it/compute/aggregate/memory.rs
@@ -1,4 +1,8 @@
-use arrow2::{array::*, compute::aggregate::estimated_bytes_size};
+use arrow2::{
+    array::*,
+    compute::aggregate::estimated_bytes_size,
+    datatypes::{DataType, Field},
+};
 
 #[test]
 fn primitive() {
@@ -17,3 +21,12 @@ fn utf8() {
     let a = Utf8Array::<i32>::from_slice(["aaa"]);
     assert_eq!(3 + 2 * std::mem::size_of::<i32>(), estimated_bytes_size(&a));
 }
+
+#[test]
+fn fixed_size_list() {
+    let data_type =
+        DataType::FixedSizeList(Box::new(Field::new("elem", DataType::Float32, false)), 3);
+    let values = Box::new(Float32Array::from_slice([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]));
+    let a = FixedSizeListArray::new(data_type, values, None);
+    assert_eq!(6 * std::mem::size_of::<f32>(), estimated_bytes_size(&a));
+}