feat(compressor): measure inputs and blocks by the bytes their rows reference

vortex-compressor/src/compressor/cascade.rs

@@ -8,6 +8,9 @@ use vortex_array::Canonical;
 use vortex_array::CanonicalValidity;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
+use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
+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::Constant;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::ExtensionArray;
@@ -29,6 +32,7 @@ use vortex_array::arrays::union::UnionArraySlotsExt;
 use vortex_array::arrays::variant::VariantArraySlotsExt;
 use vortex_array::scalar::Scalar;
 use vortex_error::VortexResult;
+use vortex_error::vortex_err;
 
 use super::CascadingCompressor;
 use super::constant;
@@ -53,7 +57,7 @@ impl CascadingCompressor {
         array: &ArrayRef,
         exec_ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let before_nbytes = array.nbytes();
+        let before_nbytes = exact_uncompressed_size(array, exec_ctx)?;
         let span = trace::compress_span(array.len(), array.dtype(), before_nbytes);
         let _enter = span.enter();
 
@@ -226,8 +230,9 @@ impl CascadingCompressor {
     /// into a single [`GenerateStatsOptions`], and picks the winner by estimated compression
     /// ratio.
     ///
-    /// If a winner is found and its compressed output is actually smaller, that output is
-    /// returned. Otherwise, the original array is returned unchanged.
+    /// 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.
     ///
     /// Empty, all-null, and constant arrays are handled by the compressor itself before any
     /// scheme evaluation (constant detection is skipped while compressing samples).
@@ -268,7 +273,8 @@ impl CascadingCompressor {
             return constant::compress_as_constant(array, exec_ctx);
         }
 
-        let before_nbytes = array.nbytes();
+        // A scheme competes against the bytes the rows reference, not dead bytes in the buffers.
+        let before_nbytes = exact_uncompressed_size(&array, exec_ctx)?;
 
         let merged_opts = eligible_schemes
             .iter()
@@ -321,3 +327,17 @@ impl CascadingCompressor {
         }
     }
 }
+
+/// The bytes `array`'s rows use once canonical, read from the exact statistic the writer caches
+/// on each chunk and computed otherwise. A dtype that cannot be measured, such as variant, is
+/// sized by the bytes it holds.
+pub(super) fn exact_uncompressed_size(
+    array: &ArrayRef,
+    exec_ctx: &mut ExecutionCtx,
+) -> VortexResult<u64> {
+    if !UncompressedSizeInBytes::supports_dtype(array.dtype()) {
+        return Ok(array.nbytes());
+    }
+    let size = uncompressed_size_in_bytes(array, UncompressedSizeOpts::exact(), exec_ctx)?;
+    u64::try_from(size).map_err(|e| vortex_err!("uncompressed size does not fit in u64: {e}"))
+}

vortex-compressor/src/compressor/sample.rs

@@ -14,6 +14,7 @@ use vortex_array::arrays::ChunkedArray;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
+use super::cascade::exact_uncompressed_size;
 use crate::CascadingCompressor;
 use crate::scheme::CompressorContext;
 use crate::scheme::EstimateScore;
@@ -178,7 +179,8 @@ pub(super) fn estimate_compression_ratio_with_sampling<S: Scheme + ?Sized>(
     };
 
     let after = compressed.nbytes();
-    let before = sample_data.array().nbytes();
+    // Sampled slices still hold the full array's buffers, so measure what their rows reference.
+    let before = exact_uncompressed_size(sample_data.array(), exec_ctx)?;
 
     let score = EstimateScore::from_sample_sizes(before, after);
 

vortex-compressor/src/compressor/tests.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
 use std::sync::LazyLock;
 
 use parking_lot::Mutex;
@@ -10,11 +11,16 @@ use vortex_array::Canonical;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 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::Constant;
 use vortex_array::arrays::Map;
 use vortex_array::arrays::NullArray;
 use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::VarBinView;
+use vortex_array::arrays::VarBinViewArray;
+use vortex_array::arrays::varbinview::BinaryView;
 use vortex_array::assert_arrays_eq;
 use vortex_array::builders::MapBuilder;
 use vortex_array::dtype::DType;
@@ -23,6 +29,8 @@ use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 use vortex_array::scalar::Scalar;
 use vortex_array::validity::Validity;
+use vortex_buffer::Buffer;
+use vortex_buffer::ByteBuffer;
 use vortex_buffer::buffer;
 use vortex_error::VortexResult;
 use vortex_session::VortexSession;
@@ -845,3 +853,48 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> {
     assert_arrays_eq!(&compressed, &array, &mut exec_ctx);
     Ok(())
 }
+
+/// Distinct 27-byte strings whose views reference a small prefix of one 2 MiB buffer, the way
+/// Arrow's Parquet reader shares a decompressed page between batches.
+fn views_into_shared_page(rows: u32) -> ArrayRef {
+    const VALUE_LEN: u32 = 27;
+    let mut page = vec![0u8; 2 << 20];
+    for row in 0..rows {
+        let value = format!("value-{row:021}");
+        let start = (row * VALUE_LEN) as usize;
+        page[start..start + VALUE_LEN as usize].copy_from_slice(value.as_bytes());
+    }
+    let page = ByteBuffer::from(page);
+    let views = (0..rows)
+        .map(|row| {
+            let start = (row * VALUE_LEN) as usize;
+            BinaryView::make_view(&page[start..start + VALUE_LEN as usize], 0, row * VALUE_LEN)
+        })
+        .collect::<Buffer<BinaryView>>();
+    // SAFETY: every view points at `VALUE_LEN` bytes inside the page.
+    unsafe {
+        VarBinViewArray::new_unchecked(
+            views,
+            Arc::from([page]),
+            DType::Utf8(Nullability::NonNullable),
+            Validity::NonNullable,
+        )
+    }
+    .into_array()
+}
+
+#[test]
+fn scheme_must_beat_referenced_bytes_not_held_bytes() -> VortexResult<()> {
+    let mut exec_ctx = SESSION.create_execution_ctx();
+    let array = views_into_shared_page(2000);
+    let exact = uncompressed_size_in_bytes(&array, UncompressedSizeOpts::exact(), &mut exec_ctx)?;
+
+    // Every value is distinct, so a dictionary is larger than the referenced bytes but far
+    // smaller than the 2 MiB the array holds. It must lose.
+    let compressed = compressor().compress(&array, &mut exec_ctx)?;
+
+    assert!(compressed.is::<VarBinView>());
+    assert_eq!(compressed.nbytes(), exact as u64);
+    assert_arrays_eq!(&compressed, &array, &mut exec_ctx);
+    Ok(())
+}

vortex-file/src/tests.rs

@@ -149,7 +149,7 @@ async fn test_eof_values() {
 #[rstest]
 #[case::default(
     BtrBlocksCompressorBuilder::from_session(&SESSION),
-    if cfg!(feature = "zstd") { 70_076 } else { 70_012 }
+    if cfg!(feature = "zstd") { 70_100 } else { 70_036 }
 )]
 #[cfg_attr(
     feature = "zstd",

vortex-layout/src/layouts/compressed.rs

@@ -8,6 +8,9 @@ use futures::StreamExt as _;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::VortexSessionExecute;
+use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
+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::expr::stats::Stat;
 use vortex_btrblocks::BtrBlocksCompressor;
 use vortex_error::VortexResult;
@@ -110,8 +113,16 @@ impl LayoutStrategy for CompressingStrategy {
                 handle.spawn_cpu(move || {
                     let (sequence_id, chunk) = chunk?;
                     let mut ctx = session.create_execution_ctx();
-                    // Compute the stats for the chunk prior to compression
+                    // Compute the stats for the chunk prior to compression. The exact
+                    // uncompressed size is measured once here and cached for the compressor.
                     chunk.statistics().compute_all(&stats, &mut ctx)?;
+                    if UncompressedSizeInBytes::supports_dtype(chunk.dtype()) {
+                        uncompressed_size_in_bytes(
+                            &chunk,
+                            UncompressedSizeOpts::exact(),
+                            &mut ctx,
+                        )?;
+                    }
                     Ok((sequence_id, compressor.compress_chunk(&chunk, &mut ctx)?))
                 })
             })