Don't pass extension arrays through the compressor unchanged

vortex-compressor/src/compressor/cascade.rs

@@ -44,6 +44,23 @@ use crate::stats::ArrayAndStats;
 use crate::stats::GenerateStatsOptions;
 use crate::trace;
 
+/// The outcome of [`CascadingCompressor::choose_and_compress`].
+enum Selection {
+    /// A scheme, or the compressor's own empty, null or constant handling, produced a new array.
+    Compressed(ArrayRef),
+    /// Nothing beat the input, which is handed back as it was.
+    Unchanged(ArrayRef),
+}
+
+impl Selection {
+    /// Returns the array, compressed or not.
+    fn into_array(self) -> ArrayRef {
+        match self {
+            Selection::Compressed(array) | Selection::Unchanged(array) => array,
+        }
+    }
+}
+
 impl CascadingCompressor {
     /// Compresses an array using cascading adaptive compression.
     ///
@@ -114,15 +131,15 @@ impl CascadingCompressor {
     ) -> VortexResult<ArrayRef> {
         match array {
             Canonical::Null(null_array) => Ok(null_array.into_array()),
-            Canonical::Bool(bool_array) => {
-                self.choose_and_compress(Canonical::Bool(bool_array), compress_ctx, exec_ctx)
-            }
-            Canonical::Primitive(primitive) => {
-                self.choose_and_compress(Canonical::Primitive(primitive), compress_ctx, exec_ctx)
-            }
-            Canonical::Decimal(decimal) => {
-                self.choose_and_compress(Canonical::Decimal(decimal), compress_ctx, exec_ctx)
-            }
+            Canonical::Bool(bool_array) => self
+                .choose_and_compress(Canonical::Bool(bool_array), compress_ctx, exec_ctx)
+                .map(Selection::into_array),
+            Canonical::Primitive(primitive) => self
+                .choose_and_compress(Canonical::Primitive(primitive), compress_ctx, exec_ctx)
+                .map(Selection::into_array),
+            Canonical::Decimal(decimal) => self
+                .choose_and_compress(Canonical::Decimal(decimal), compress_ctx, exec_ctx)
+                .map(Selection::into_array),
             Canonical::Struct(struct_array) => {
                 let fields = struct_array
                     .iter_unmasked_fields()
@@ -169,26 +186,26 @@ impl CascadingCompressor {
                 )?
                 .into_array())
             }
-            Canonical::VarBinView(varbinview) => {
-                self.choose_and_compress(Canonical::VarBinView(varbinview), compress_ctx, exec_ctx)
-            }
+            Canonical::VarBinView(varbinview) => self
+                .choose_and_compress(Canonical::VarBinView(varbinview), compress_ctx, exec_ctx)
+                .map(Selection::into_array),
             Canonical::Extension(ext_array) => {
                 // Try scheme-based compression first.
-                let scheme_compressed = self.choose_and_compress(
+                let selection = self.choose_and_compress(
                     Canonical::Extension(ext_array.clone()),
                     compress_ctx,
                     exec_ctx,
                 )?;
 
                 // A constant extension array (that might be masked) is already in its terminal
                 // representation, and compressing the storage separately cannot do better.
-                if scheme_compressed.is::<Constant>() {
-                    return Ok(scheme_compressed);
-                }
-                if let Some(masked) = scheme_compressed.as_opt::<Masked>()
-                    && masked.child().is::<Constant>()
+                if let Selection::Compressed(compressed) = &selection
+                    && (compressed.is::<Constant>()
+                        || compressed
+                            .as_opt::<Masked>()
+                            .is_some_and(|masked| masked.child().is::<Constant>()))
                 {
-                    return Ok(scheme_compressed);
+                    return Ok(selection.into_array());
                 }
 
                 // Also compress the underlying storage array. Some extension schemes can beat the
@@ -198,10 +215,15 @@ impl CascadingCompressor {
                     ExtensionArray::new(ext_array.ext_dtype().clone(), compressed_storage)
                         .into_array();
 
-                if scheme_compressed.nbytes() < storage_compressed.nbytes() {
-                    Ok(scheme_compressed)
-                } else {
-                    Ok(storage_compressed)
+                match selection {
+                    Selection::Compressed(scheme_compressed)
+                        if scheme_compressed.nbytes() < storage_compressed.nbytes() =>
+                    {
+                        Ok(scheme_compressed)
+                    }
+                    // A canonical extension array's storage is only canonical at its top level,
+                    // so an unchanged input can still hold lazy, unserializable arrays.
+                    Selection::Compressed(_) | Selection::Unchanged(_) => Ok(storage_compressed),
                 }
             }
             Canonical::Variant(variant_array) => {
@@ -231,8 +253,10 @@ impl CascadingCompressor {
     /// ratio.
     ///
     /// If a winner is found and its compressed output is actually smaller than the exact
-    /// uncompressed size, that output is returned. Otherwise, the original array is returned
-    /// unchanged.
+    /// uncompressed size, that output is returned as [`Selection::Compressed`]. Otherwise, the
+    /// original array is returned as [`Selection::Unchanged`], and the caller decides whether it
+    /// is a valid result: it is for leaf arrays, but not for arrays like extensions whose children
+    /// have yet to be compressed.
     ///
     /// Empty, all-null, and constant arrays are handled by the compressor itself before any
     /// scheme evaluation (constant detection is skipped while compressing samples).
@@ -244,7 +268,7 @@ impl CascadingCompressor {
         canonical: Canonical,
         compress_ctx: CompressorContext,
         exec_ctx: &mut ExecutionCtx,
-    ) -> VortexResult<ArrayRef> {
+    ) -> VortexResult<Selection> {
         let eligible_schemes: Vec<&'static dyn Scheme> = self
             .schemes
             .iter()
@@ -255,13 +279,13 @@ impl CascadingCompressor {
         let array: ArrayRef = canonical.into();
 
         if array.is_empty() {
-            return Ok(array);
+            return Ok(Selection::Unchanged(array));
         }
 
         if array.all_invalid(exec_ctx)? {
-            return Ok(
+            return Ok(Selection::Compressed(
                 ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()).into_array(),
-            );
+            ));
         }
 
         // Constant detection is built into the compressor: a constant leaf always short-circuits
@@ -270,7 +294,7 @@ impl CascadingCompressor {
         //
         // Compare valid values directly so even nullable constants avoid distinct-value counting.
         if !compress_ctx.is_sample() && constant::is_constant_for_compression(&array, exec_ctx)? {
-            return constant::compress_as_constant(array, exec_ctx);
+            return constant::compress_as_constant(array, exec_ctx).map(Selection::Compressed);
         }
 
         // A scheme competes against the bytes the rows reference, not dead bytes in the buffers.
@@ -286,13 +310,13 @@ impl CascadingCompressor {
         let data = ArrayAndStats::new(array, merged_opts);
 
         if eligible_schemes.is_empty() {
-            return Ok(data.into_array());
+            return Ok(Selection::Unchanged(data.into_array()));
         }
 
         let Some((winner, winner_estimate)) =
             self.choose_best_scheme(&eligible_schemes, &data, compress_ctx.clone(), exec_ctx)?
         else {
-            return Ok(data.into_array());
+            return Ok(Selection::Unchanged(data.into_array()));
         };
 
         // Run the winning scheme's `compress`. On failure, emit an ERROR event carrying the
@@ -321,9 +345,9 @@ impl CascadingCompressor {
         );
 
         if accepted {
-            Ok(compressed)
+            Ok(Selection::Compressed(compressed))
         } else {
-            Ok(data.into_array())
+            Ok(Selection::Unchanged(data.into_array()))
         }
     }
 }

vortex-compressor/src/compressor/tests.rs

@@ -14,10 +14,16 @@ use vortex_array::VortexSessionExecute;
 use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts;
 use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes;
 use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::ChunkedArray;
 use vortex_array::arrays::Constant;
+use vortex_array::arrays::DictArray;
+use vortex_array::arrays::ExtensionArray;
+use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::Map;
 use vortex_array::arrays::NullArray;
 use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::Slice;
+use vortex_array::arrays::SliceArray;
 use vortex_array::arrays::VarBinView;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::arrays::varbinview::BinaryView;
@@ -27,12 +33,17 @@ use vortex_array::dtype::DType;
 use vortex_array::dtype::MapDType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
+use vortex_array::dtype::extension::ExtDType;
+use vortex_array::dtype::extension::ExtId;
+use vortex_array::dtype::extension::ExtVTable;
 use vortex_array::scalar::Scalar;
+use vortex_array::scalar::ScalarValue;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_buffer::ByteBuffer;
 use vortex_buffer::buffer;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
 use vortex_session::VortexSession;
 
 use super::CascadingCompressor;
@@ -881,3 +892,73 @@ fn scheme_must_beat_referenced_bytes_not_held_bytes() -> VortexResult<()> {
     assert_arrays_eq!(&compressed, &array, &mut exec_ctx);
     Ok(())
 }
+
+/// An extension type over fixed-size lists, like UUIDs over 16 bytes.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
+struct BytesExtension;
+
+impl ExtVTable for BytesExtension {
+    type Metadata = u8;
+    type NativeValue<'a> = &'a ScalarValue;
+
+    #[expect(clippy::disallowed_methods, reason = "test-only id")]
+    fn id(&self) -> ExtId {
+        ExtId::new("test.bytes")
+    }
+
+    fn serialize_metadata(&self, metadata: &u8) -> VortexResult<Vec<u8>> {
+        Ok(vec![*metadata])
+    }
+
+    fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<u8> {
+        vortex_ensure!(metadata.len() == 1, "expected one metadata byte");
+        Ok(metadata[0])
+    }
+
+    fn validate_dtype(dtype: &ExtDType<Self>) -> VortexResult<()> {
+        vortex_ensure!(
+            matches!(dtype.storage_dtype(), DType::FixedSizeList(..)),
+            "expected fixed-size list storage"
+        );
+        Ok(())
+    }
+
+    fn unpack_native<'a>(
+        _dtype: &'a ExtDType<Self>,
+        value: &'a ScalarValue,
+    ) -> VortexResult<&'a ScalarValue> {
+        Ok(value)
+    }
+}
+
+/// Canonicalizing an extension array only canonicalizes its storage's top level, so the elements
+/// of fixed-size list storage can stay a lazy slice. When no scheme compresses such an array it
+/// must not be returned as-is, because the slice is not serializable.
+#[test]
+fn extension_with_lazy_storage_is_not_passed_through() -> VortexResult<()> {
+    // Dictionaries of wide values keep the lazy array smaller than its canonical form, and
+    // chunked arrays only slice on the execute path, so the slice stays lazy.
+    let dict = |offset: u64| -> VortexResult<ArrayRef> {
+        let codes = PrimitiveArray::from_iter((0..1024u32).map(|i| (i % 4) as u8)).into_array();
+        let values = PrimitiveArray::from_iter([offset, offset + 1, offset + 2, offset + 3]);
+        Ok(DictArray::try_new(codes, values.into_array())?.into_array())
+    };
+    let chunks = vec![dict(0)?, dict(100)?];
+    let elements_dtype = chunks[0].dtype().clone();
+    let chunked = ChunkedArray::try_new(chunks, elements_dtype)?.into_array();
+    let elements = SliceArray::new(chunked, 16..2048).into_array();
+    let storage = FixedSizeListArray::new(elements, 16, Validity::NonNullable, 127).into_array();
+    let dtype = ExtDType::<BytesExtension>::try_new(0, storage.dtype().clone())?.erased();
+    let array = ExtensionArray::new(dtype, storage).into_array();
+    let mut exec_ctx = SESSION.create_execution_ctx();
+
+    let compressed = CascadingCompressor::new(Vec::new()).compress(&array, &mut exec_ctx)?;
+
+    assert!(
+        !compressed.depth_first_traversal().any(|a| a.is::<Slice>()),
+        "compressed array still contains a lazy slice:\n{}",
+        compressed.display_tree()
+    );
+    assert_arrays_eq!(&compressed, &array, &mut exec_ctx);
+    Ok(())
+}