Canonicalize constant array with extension dtype (#1322)

also `storage_array` -> `storage_data` seems appropriate.

---------

Co-authored-by: Robert Kruszewski <[email protected]>

Author: AdamGS

encodings/datetime-parts/src/array.rs

@@ -122,7 +122,7 @@ impl ArrayVariants for DateTimePartsArray {
 }
 
 impl ExtensionArrayTrait for DateTimePartsArray {
-    fn storage_array(&self) -> ArrayData {
+    fn storage_data(&self) -> ArrayData {
         // FIXME(ngates): this needs to be a tuple array so we can implement Compare
         // we don't want to write validity twice, so we pull it up to the top
         let days = try_cast(self.days(), &self.days().dtype().as_nonnullable()).vortex_unwrap();

vortex-array/src/array/chunked/variants.rs

@@ -122,10 +122,10 @@ impl StructArrayTrait for ChunkedArray {
 impl ListArrayTrait for ChunkedArray {}
 
 impl ExtensionArrayTrait for ChunkedArray {
-    fn storage_array(&self) -> ArrayData {
+    fn storage_data(&self) -> ArrayData {
         ChunkedArray::from_iter(
             self.chunks()
-                .map(|chunk| chunk.with_dyn(|a| a.as_extension_array_unchecked().storage_array())),
+                .map(|chunk| chunk.with_dyn(|a| a.as_extension_array_unchecked().storage_data())),
         )
         .into_array()
     }

vortex-array/src/array/constant/canonical.rs

@@ -2,12 +2,14 @@ use arrow_array::builder::make_view;
 use arrow_buffer::{BooleanBuffer, BufferBuilder};
 use vortex_buffer::Buffer;
 use vortex_dtype::{match_each_native_ptype, DType, Nullability, PType};
-use vortex_error::{vortex_bail, VortexResult};
-use vortex_scalar::{BinaryScalar, BoolScalar, Utf8Scalar};
+use vortex_error::{vortex_bail, vortex_panic, VortexResult};
+use vortex_scalar::{BinaryScalar, BoolScalar, ExtScalar, Scalar, Utf8Scalar};
 
 use crate::array::constant::ConstantArray;
 use crate::array::primitive::PrimitiveArray;
-use crate::array::{BinaryView, BoolArray, NullArray, VarBinViewArray, VIEW_SIZE_BYTES};
+use crate::array::{
+    BinaryView, BoolArray, ExtensionArray, NullArray, VarBinViewArray, VIEW_SIZE_BYTES,
+};
 use crate::validity::Validity;
 use crate::{ArrayDType, Canonical, IntoArrayData, IntoCanonical};
 
@@ -63,6 +65,17 @@ impl IntoCanonical for ConstantArray {
             return Ok(Canonical::Null(NullArray::new(self.len())));
         }
 
+        if let Ok(s) = ExtScalar::try_from(scalar) {
+            let DType::Extension(ext_dtype) = s.dtype() else {
+                vortex_panic!("ExtScalar has a non-ext dtype {}", s.dtype());
+            };
+
+            let storage_dtype = ext_dtype.storage_dtype();
+            let storage_scalar = Scalar::new(storage_dtype.clone(), s.value().clone());
+            let storage_array = ConstantArray::new(storage_scalar, self.len()).into_array();
+            return ExtensionArray::new(ext_dtype.clone(), storage_array).into_canonical();
+        }
+
         vortex_bail!("Unsupported scalar type {}", self.dtype())
     }
 }

vortex-array/src/array/constant/variants.rs

@@ -210,7 +210,7 @@ impl StructArrayTrait for ConstantArray {
 impl ListArrayTrait for ConstantArray {}
 
 impl ExtensionArrayTrait for ConstantArray {
-    fn storage_array(&self) -> ArrayData {
+    fn storage_data(&self) -> ArrayData {
         let storage_dtype = self.ext_dtype().storage_dtype().clone();
         ConstantArray::new(
             Scalar::new(storage_dtype, self.scalar_value().clone()),

vortex-array/src/array/extension/mod.rs

@@ -66,7 +66,7 @@ impl ArrayVariants for ExtensionArray {
 }
 
 impl ExtensionArrayTrait for ExtensionArray {
-    fn storage_array(&self) -> ArrayData {
+    fn storage_data(&self) -> ArrayData {
         self.storage()
     }
 }

vortex-array/src/array/sparse/variants.rs

@@ -123,11 +123,11 @@ impl StructArrayTrait for SparseArray {
 impl ListArrayTrait for SparseArray {}
 
 impl ExtensionArrayTrait for SparseArray {
-    fn storage_array(&self) -> ArrayData {
+    fn storage_data(&self) -> ArrayData {
         SparseArray::try_new_with_offset(
             self.indices().clone(),
             self.values()
-                .with_dyn(|a| a.as_extension_array_unchecked().storage_array()),
+                .with_dyn(|a| a.as_extension_array_unchecked().storage_data()),
             self.len(),
             self.indices_offset(),
             self.fill_value().clone(),

vortex-array/src/variants.rs

@@ -317,5 +317,5 @@ pub trait ExtensionArrayTrait: ArrayTrait {
     }
 
     /// Returns the underlying [`ArrayData`], without the [`ExtDType`].
-    fn storage_array(&self) -> ArrayData;
+    fn storage_data(&self) -> ArrayData;
 }