canonicalize before ExtensionArray::try_from

spiraldb · Jul 22, 2024 · 7e08d95 · 7e08d95
1 parent 6c5e1f2
commit 7e08d95
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/vortex-array/src/array/chunked/canonical.rs b/vortex-array/src/array/chunked/canonical.rs
@@ -51,7 +51,11 @@ pub(crate) fn try_canonicalize_chunks(
         DType::Extension(ext_dtype, _) => {
             let storage_chunks: Vec<Array> = chunks
                 .iter()
-                .map(|chunk| ExtensionArray::try_from(chunk).unwrap().storage())
+                .map(|chunk| {
+                    ExtensionArray::try_from(chunk.clone().into_canonical())
+                        .unwrap()
+                        .storage()
+                })
                 .collect();
             let storage_dtype = storage_chunks.first().unwrap().dtype().clone();
             let ext_array = ExtensionArray::new(

diff --git a/vortex-array/src/array/chunked/compute/take.rs b/vortex-array/src/array/chunked/compute/take.rs
@@ -8,10 +8,12 @@ use crate::array::primitive::PrimitiveArray;
 use crate::compute::unary::{scalar_at, subtract_scalar, try_cast};
 use crate::compute::{search_sorted, slice, take, SearchSortedSide, TakeFn};
 use crate::stats::ArrayStatistics;
-use crate::{Array, ArrayDType, IntoArray, ToArray};
+use crate::{Array, ArrayDType, IntoArray, IntoCanonical, ToArray};
 
 impl TakeFn for ChunkedArray {
     fn take(&self, indices: &Array) -> VortexResult<Array> {
+        let indices = indices.clone().into_canonical()?.into_array();
+
         // Fast path for strict sorted indices.
         if indices
             .statistics()
@@ -22,11 +24,10 @@ impl TakeFn for ChunkedArray {
                 return Ok(self.to_array());
             }
 
-            return take_strict_sorted(self, indices);
+            return take_strict_sorted(self, &indices);
         }
 
-        // FIXME(ngates): this is wrong, need to canonicalise
-        let indices = PrimitiveArray::try_from(try_cast(indices, PType::U64.into())?)?;
+        let indices = PrimitiveArray::try_from(try_cast(&indices, PType::U64.into())?)?;
 
         // While the chunk idx remains the same, accumulate a list of chunk indices.
         let mut chunks = Vec::new();