diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 0d1de67c421..a34d945d7e4 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/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 { 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,12 +186,12 @@ 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, @@ -182,13 +199,13 @@ impl CascadingCompressor { // 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::() { - return Ok(scheme_compressed); - } - if let Some(masked) = scheme_compressed.as_opt::() - && masked.child().is::() + if let Selection::Compressed(compressed) = &selection + && (compressed.is::() + || compressed + .as_opt::() + .is_some_and(|masked| masked.child().is::())) { - 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 { + ) -> VortexResult { 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())) } } } diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 1b2005ec00d..8ced45f268a 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/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> { + Ok(vec![*metadata]) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + vortex_ensure!(metadata.len() == 1, "expected one metadata byte"); + Ok(metadata[0]) + } + + fn validate_dtype(dtype: &ExtDType) -> VortexResult<()> { + vortex_ensure!( + matches!(dtype.storage_dtype(), DType::FixedSizeList(..)), + "expected fixed-size list storage" + ); + Ok(()) + } + + fn unpack_native<'a>( + _dtype: &'a ExtDType, + 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 { + 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::::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::()), + "compressed array still contains a lazy slice:\n{}", + compressed.display_tree() + ); + assert_arrays_eq!(&compressed, &array, &mut exec_ctx); + Ok(()) +}