From 1f6df340306e1458ef711f7d65a93718dd0fcbf5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jul 2024 20:20:59 +0100 Subject: [PATCH] Extension DType canonicalization behaves as other array types --- vortex-array/src/array/chunked/canonical.rs | 10 +++++++++- vortex-array/src/array/extension/mod.rs | 7 ++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/array/chunked/canonical.rs b/vortex-array/src/array/chunked/canonical.rs index 778a33f61e..05a53046ff 100644 --- a/vortex-array/src/array/chunked/canonical.rs +++ b/vortex-array/src/array/chunked/canonical.rs @@ -47,9 +47,17 @@ pub(crate) fn try_canonicalize_chunks( // Extension arrays wrap an internal storage array, which can hold a ChunkedArray until // it is safe to unpack them. DType::Extension(ext_dtype, _) => { + if chunks.is_empty() { + panic!("Can't canonicalize empty extension array") + } + let ext_chunk = ExtensionArray::try_from(chunks.first().unwrap()).unwrap(); + let storage_chunks = chunks + .into_iter() + .map(|c| ExtensionArray::try_from(c).unwrap().storage()) + .collect::>(); let ext_array = ExtensionArray::new( ext_dtype.clone(), - ChunkedArray::try_new(chunks, dtype.clone())?.into_array(), + try_canonicalize_chunks(storage_chunks, ext_chunk.storage_dtype())?.into_array(), ); Ok(Canonical::Extension(ext_array)) diff --git a/vortex-array/src/array/extension/mod.rs b/vortex-array/src/array/extension/mod.rs index 8f3f11e0db..859d68c504 100644 --- a/vortex-array/src/array/extension/mod.rs +++ b/vortex-array/src/array/extension/mod.rs @@ -33,7 +33,7 @@ impl ExtensionArray { pub fn storage(&self) -> Array { self.array() - .child(0, &self.metadata().storage_dtype, self.len()) + .child(0, self.storage_dtype(), self.len()) .expect("Missing storage array") } @@ -50,6 +50,11 @@ impl ExtensionArray { }; ext } + + #[inline] + pub(crate) fn storage_dtype(&self) -> &DType { + &self.metadata().storage_dtype + } } impl ArrayTrait for ExtensionArray {}