Feature: introduces fuzzing with extension arrays

fuzz/src/array/compare.rs

@@ -123,8 +123,16 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper
             )
             .into_array()
         }
-        d @ (DType::Null | DType::Extension(_)) => {
-            unreachable!("DType {d} not supported for fuzzing")
+        DType::Null => {
+            unreachable!("DType null not supported for fuzzing")
+        }
+        DType::Extension(..) => {
+            // Extension arrays delegate comparison to their storage type
+            compare_canonical_array(
+                array.to_extension().storage(),
+                &value.as_extension().storage(),
+                operator,
+            )
         }
     }
 }

fuzz/src/array/fill_null.rs

@@ -1,13 +1,16 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
+
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::IntoArray;
 use vortex_array::ToCanonical;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::compute::fill_null;
@@ -16,13 +19,13 @@ use vortex_array::vtable::ValidityHelper;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_dtype::DType;
+use vortex_dtype::ExtDType;
 use vortex_dtype::Nullability;
 use vortex_dtype::match_each_decimal_value_type;
 use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
-
 /// Apply fill_null on the canonical form of the array to get a consistent baseline.
 /// This implementation manually fills null values for each canonical type
 /// without using the fill_null method, to serve as an independent baseline for testing.
@@ -40,13 +43,33 @@ pub fn fill_null_canonical_array(
         Canonical::VarBinView(array) => {
             fill_varbinview_array(&array, fill_value, result_nullability)
         }
-        Canonical::Struct(_)
-        | Canonical::List(_)
-        | Canonical::FixedSizeList(_)
-        | Canonical::Extension(_) => fill_null(canonical.as_ref(), fill_value)?,
+        Canonical::Extension(array) => fill_extension_array(&array, fill_value),
+        Canonical::Struct(_) | Canonical::List(_) | Canonical::FixedSizeList(_) => {
+            fill_null(canonical.as_ref(), fill_value)?
+        }
     })
 }
 
+fn fill_extension_array(array: &ExtensionArray, fill_value: &Scalar) -> ArrayRef {
+    let filled_storage = fill_null_canonical_array(
+        array.storage().to_canonical(),
+        &fill_value.as_extension().storage(),
+    )
+    .vortex_expect("fill_null should succeed in canonical form");
+
+    if filled_storage.dtype().nullability() == array.ext_dtype().storage_dtype().nullability() {
+        ExtensionArray::new(array.ext_dtype().clone(), filled_storage).into_array()
+    } else {
+        let new_ext_dtype = Arc::new(ExtDType::new(
+            array.ext_dtype().id().clone(),
+            Arc::new(filled_storage.dtype().clone()),
+            array.ext_dtype().metadata().cloned(),
+        ));
+
+        ExtensionArray::new(new_ext_dtype, filled_storage).into_array()
+    }
+}
+
 fn fill_bool_array(
     array: &BoolArray,
     fill_value: &Scalar,

fuzz/src/array/filter.rs

@@ -1,20 +1,24 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
+
 use vortex_array::Array;
 use vortex_array::ArrayRef;
 use vortex_array::IntoArray;
 use vortex_array::ToCanonical;
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
 use vortex_buffer::Buffer;
 use vortex_dtype::DType;
+use vortex_dtype::ExtDType;
 use vortex_dtype::match_each_decimal_value_type;
 use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
@@ -115,8 +119,23 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul
             }
             take_canonical_array_non_nullable_indices(array, indices.as_slice())
         }
-        d @ (DType::Null | DType::Extension(_)) => {
-            unreachable!("DType {d} not supported for fuzzing")
+        DType::Extension(ext_dtype) => {
+            // Extension arrays delegate filter to their storage type
+            let filtered_storage = filter_canonical_array(array.to_extension().storage(), filter)?;
+
+            if filtered_storage.dtype().nullability() == ext_dtype.storage_dtype().nullability() {
+                Ok(ExtensionArray::new(ext_dtype.clone(), filtered_storage).into_array())
+            } else {
+                let new_ext_dtype = Arc::new(ExtDType::new(
+                    ext_dtype.id().clone(),
+                    Arc::new(filtered_storage.dtype().clone()),
+                    ext_dtype.metadata().cloned(),
+                ));
+                Ok(ExtensionArray::new(new_ext_dtype, filtered_storage).into_array())
+            }
+        }
+        DType::Null => {
+            unreachable!("Cannot search sorted on Null array")
         }
     }
 }

fuzz/src/array/mod.rs

@@ -662,6 +662,9 @@ pub fn assert_array_eq(
     rhs: &ArrayRef,
     step: usize,
 ) -> crate::error::VortexFuzzResult<()> {
+    use vortex_array::ToCanonical;
+    use vortex_array::arrays::validator_for_ext_type;
+
     use crate::error::Backtrace;
     use crate::error::VortexFuzzError;
 
@@ -699,7 +702,22 @@ pub fn assert_array_eq(
                 Backtrace::capture(),
             ));
         }
+
+        // Also validate the expected array's domain constraints for extension types
+        if matches!(lhs.dtype(), DType::Extension(..)) {
+            let validator = validator_for_ext_type(lhs.to_extension().ext_dtype());
+            if !validator(&l) {
+                return Err(VortexFuzzError::DomainValidationFailed(
+                    l,
+                    idx,
+                    lhs.clone(),
+                    step,
+                    Backtrace::capture(),
+                ));
+            }
+        }
     }
+
     Ok(())
 }
 

fuzz/src/array/search_sorted.rs

@@ -129,8 +129,16 @@ pub fn search_sorted_canonical_array(
             let scalar_vals = (0..array.len()).map(|i| array.scalar_at(i)).collect_vec();
             Ok(scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side))
         }
-        d @ (DType::Null | DType::Extension(_)) => {
-            unreachable!("DType {d} not supported for fuzzing")
+        DType::Extension(..) => {
+            // Extension arrays delegate search to their storage type
+            search_sorted_canonical_array(
+                array.to_extension().storage(),
+                &scalar.as_extension().storage(),
+                side,
+            )
+        }
+        DType::Null => {
+            unreachable!("Cannot search sorted on Null array")
         }
     }
 }

fuzz/src/array/slice.rs

@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
+use std::sync::Arc;
 
 use vortex_array::Array;
 use vortex_array::ArrayRef;
@@ -8,13 +9,15 @@ use vortex_array::ToCanonical;
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::ListViewArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::validity::Validity;
 use vortex_dtype::DType;
+use vortex_dtype::ExtDType;
 use vortex_dtype::match_each_decimal_value_type;
 use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
@@ -113,8 +116,24 @@ pub fn slice_canonical_array(
                 .to_array(),
             )
         }
-        d @ (DType::Null | DType::Extension(_)) => {
-            unreachable!("DType {d} not supported for fuzzing")
+        DType::Extension(ext_dtype) => {
+            // Extension arrays delegate slicing to their storage type
+            let sliced_storage =
+                slice_canonical_array(array.to_extension().storage(), start, stop)?;
+
+            if sliced_storage.dtype().nullability() == ext_dtype.storage_dtype().nullability() {
+                Ok(ExtensionArray::new(ext_dtype.clone(), sliced_storage).into_array())
+            } else {
+                let new_ext_dtype = Arc::new(ExtDType::new(
+                    ext_dtype.id().clone(),
+                    Arc::new(sliced_storage.dtype().clone()),
+                    ext_dtype.metadata().cloned(),
+                ));
+                Ok(ExtensionArray::new(new_ext_dtype, sliced_storage).into_array())
+            }
+        }
+        DType::Null => {
+            unreachable!("Cannot search sorted on Null array")
         }
     }
 }

fuzz/src/array/sort.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::cmp::Ordering;
+use std::sync::Arc;
 
 use vortex_array::Array;
 use vortex_array::ArrayRef;
@@ -10,9 +11,11 @@ use vortex_array::ToCanonical;
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_dtype::DType;
+use vortex_dtype::ExtDType;
 use vortex_dtype::NativePType;
 use vortex_dtype::match_each_decimal_value_type;
 use vortex_dtype::match_each_native_ptype;
@@ -80,8 +83,24 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult<ArrayRef> {
             });
             take_canonical_array_non_nullable_indices(array, &sort_indices)
         }
-        d @ (DType::Null | DType::Extension(_)) => {
-            unreachable!("DType {d} not supported for fuzzing")
+        DType::Null => {
+            // Null arrays don't need sorting - all elements are null
+            Ok(array.to_array())
+        }
+        DType::Extension(ext_dtype) => {
+            // Extension arrays delegate sorting to their storage type
+            let sorted_storage = sort_canonical_array(array.to_extension().storage())?;
+
+            if sorted_storage.dtype().nullability() == ext_dtype.storage_dtype().nullability() {
+                Ok(ExtensionArray::new(ext_dtype.clone(), sorted_storage).into_array())
+            } else {
+                let new_ext_dtype = Arc::new(ExtDType::new(
+                    ext_dtype.id().clone(),
+                    Arc::new(sorted_storage.dtype().clone()),
+                    ext_dtype.metadata().cloned(),
+                ));
+                Ok(ExtensionArray::new(new_ext_dtype, sorted_storage).into_array())
+            }
         }
     }
 }

fuzz/src/array/take.rs

@@ -1,13 +1,16 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
+
 use vortex_array::Array;
 use vortex_array::ArrayRef;
 use vortex_array::IntoArray;
 use vortex_array::ToCanonical;
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::VarBinViewArray;
@@ -16,14 +19,14 @@ use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_dtype::DType;
 use vortex_dtype::DecimalDType;
+use vortex_dtype::ExtDType;
 use vortex_dtype::NativeDecimalType;
 use vortex_dtype::NativePType;
 use vortex_dtype::Nullability;
 use vortex_dtype::match_each_decimal_value_type;
 use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-
 pub fn take_canonical_array_non_nullable_indices(
     array: &dyn Array,
     indices: &[usize],
@@ -141,8 +144,23 @@ pub fn take_canonical_array(
             }
             Ok(builder.finish())
         }
-        d @ (DType::Null | DType::Extension(_)) => {
-            unreachable!("DType {d} not supported for fuzzing")
+        DType::Extension(ext_dtype) => {
+            // Extension arrays delegate take to their storage type
+            let taken_storage = take_canonical_array(array.to_extension().storage(), indices)?;
+
+            if taken_storage.dtype().nullability() == ext_dtype.storage_dtype().nullability() {
+                Ok(ExtensionArray::new(ext_dtype.clone(), taken_storage).into_array())
+            } else {
+                let new_ext_dtype = Arc::new(ExtDType::new(
+                    ext_dtype.id().clone(),
+                    Arc::new(taken_storage.dtype().clone()),
+                    ext_dtype.metadata().cloned(),
+                ));
+                Ok(ExtensionArray::new(new_ext_dtype, taken_storage).into_array())
+            }
+        }
+        DType::Null => {
+            unreachable!("Null type not supported for fuzzing")
         }
     }
 }

fuzz/src/error.rs

@@ -59,6 +59,8 @@ pub enum VortexFuzzError {
     LengthMismatch(usize, usize, ArrayRef, ArrayRef, usize, Backtrace),
 
     VortexError(VortexError, Backtrace),
+
+    DomainValidationFailed(Scalar, usize, ArrayRef, usize, Backtrace),
 }
 
 impl Debug for VortexFuzzError {
@@ -121,6 +123,13 @@ impl Display for VortexFuzzError {
                     rhs.display_tree(),
                 )
             }
+            VortexFuzzError::DomainValidationFailed(expected, idx, lhs, step, backtrace) => {
+                write!(
+                    f,
+                    "Domain validation failed:\n  Scalar: {expected}\n  Index: {idx}\n  Array: {}\n  Step: {step}\nBacktrace:\n{backtrace}",
+                    lhs.display_tree(),
+                )
+            }
             VortexFuzzError::VortexError(err, backtrace) => {
                 write!(f, "{err}\nBacktrace:\n{backtrace}")
             }
@@ -137,6 +146,7 @@ impl Error for VortexFuzzError {
             | VortexFuzzError::LengthMismatch(..)
             | VortexFuzzError::ScalarMismatch(..)
             | VortexFuzzError::MinMaxMismatch(..)
+            | VortexFuzzError::DomainValidationFailed(..)
             | VortexFuzzError::DTypeMismatch(..) => None,
         }
     }