Skip to content

Commit d42e7c9

Browse files
joseph-isaacsjoseph-isaacs
authored
feat(compressor): measure inputs and blocks by the bytes their rows reference (#10445)
1 parent 84e0d73 commit d42e7c9

7 files changed

Lines changed: 202 additions & 19 deletions

File tree

‎vortex-compressor/src/compressor/cascade.rs‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ use vortex_array::Canonical;
88
use vortex_array::CanonicalValidity;
99
use vortex_array::ExecutionCtx;
1010
use vortex_array::IntoArray;
11+
use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
12+
use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts;
13+
use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes;
1114
use vortex_array::arrays::Constant;
1215
use vortex_array::arrays::ConstantArray;
1316
use vortex_array::arrays::ExtensionArray;
@@ -29,6 +32,7 @@ use vortex_array::arrays::union::UnionArraySlotsExt;
2932
use vortex_array::arrays::variant::VariantArraySlotsExt;
3033
use vortex_array::scalar::Scalar;
3134
use vortex_error::VortexResult;
35+
use vortex_error::vortex_err;
3236

3337
use super::CascadingCompressor;
3438
use super::constant;
@@ -53,7 +57,7 @@ impl CascadingCompressor {
5357
array: &ArrayRef,
5458
exec_ctx: &mut ExecutionCtx,
5559
) -> VortexResult<ArrayRef> {
56-
let before_nbytes = array.nbytes();
60+
let before_nbytes = exact_uncompressed_size(array, exec_ctx)?;
5761
let span = trace::compress_span(array.len(), array.dtype(), before_nbytes);
5862
let _enter = span.enter();
5963

@@ -226,8 +230,9 @@ impl CascadingCompressor {
226230
/// into a single [`GenerateStatsOptions`], and picks the winner by estimated compression
227231
/// ratio.
228232
///
229-
/// If a winner is found and its compressed output is actually smaller, that output is
230-
/// returned. Otherwise, the original array is returned unchanged.
233+
/// If a winner is found and its compressed output is actually smaller than the exact
234+
/// uncompressed size, that output is returned. Otherwise, the original array is returned
235+
/// unchanged.
231236
///
232237
/// Empty, all-null, and constant arrays are handled by the compressor itself before any
233238
/// scheme evaluation (constant detection is skipped while compressing samples).
@@ -268,7 +273,8 @@ impl CascadingCompressor {
268273
return constant::compress_as_constant(array, exec_ctx);
269274
}
270275

271-
let before_nbytes = array.nbytes();
276+
// A scheme competes against the bytes the rows reference, not dead bytes in the buffers.
277+
let before_nbytes = exact_uncompressed_size(&array, exec_ctx)?;
272278

273279
let merged_opts = eligible_schemes
274280
.iter()
@@ -321,3 +327,17 @@ impl CascadingCompressor {
321327
}
322328
}
323329
}
330+
331+
/// The bytes `array`'s rows use once canonical, read from the exact statistic the writer caches
332+
/// on each chunk and computed otherwise. A dtype that cannot be measured, such as variant, is
333+
/// sized by the bytes it holds.
334+
pub(super) fn exact_uncompressed_size(
335+
array: &ArrayRef,
336+
exec_ctx: &mut ExecutionCtx,
337+
) -> VortexResult<u64> {
338+
if !UncompressedSizeInBytes::supports_dtype(array.dtype()) {
339+
return Ok(array.nbytes());
340+
}
341+
let size = uncompressed_size_in_bytes(array, UncompressedSizeOpts::exact(), exec_ctx)?;
342+
u64::try_from(size).map_err(|e| vortex_err!("uncompressed size does not fit in u64: {e}"))
343+
}

‎vortex-compressor/src/compressor/sample.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use vortex_array::arrays::ChunkedArray;
1414
use vortex_error::VortexExpect;
1515
use vortex_error::VortexResult;
1616

17+
use super::cascade::exact_uncompressed_size;
1718
use crate::CascadingCompressor;
1819
use crate::scheme::CompressorContext;
1920
use crate::scheme::EstimateScore;
@@ -178,7 +179,8 @@ pub(super) fn estimate_compression_ratio_with_sampling<S: Scheme + ?Sized>(
178179
};
179180

180181
let after = compressed.nbytes();
181-
let before = sample_data.array().nbytes();
182+
// Sampled slices still hold the full array's buffers, so measure what their rows reference.
183+
let before = exact_uncompressed_size(sample_data.array(), exec_ctx)?;
182184

183185
let score = EstimateScore::from_sample_sizes(before, after);
184186

‎vortex-compressor/src/compressor/tests.rs‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4+
use std::sync::Arc;
45
use std::sync::LazyLock;
56

67
use parking_lot::Mutex;
@@ -10,11 +11,16 @@ use vortex_array::Canonical;
1011
use vortex_array::ExecutionCtx;
1112
use vortex_array::IntoArray;
1213
use vortex_array::VortexSessionExecute;
14+
use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts;
15+
use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes;
1316
use vortex_array::arrays::BoolArray;
1417
use vortex_array::arrays::Constant;
1518
use vortex_array::arrays::Map;
1619
use vortex_array::arrays::NullArray;
1720
use vortex_array::arrays::PrimitiveArray;
21+
use vortex_array::arrays::VarBinView;
22+
use vortex_array::arrays::VarBinViewArray;
23+
use vortex_array::arrays::varbinview::BinaryView;
1824
use vortex_array::assert_arrays_eq;
1925
use vortex_array::builders::MapBuilder;
2026
use vortex_array::dtype::DType;
@@ -23,6 +29,8 @@ use vortex_array::dtype::Nullability;
2329
use vortex_array::dtype::PType;
2430
use vortex_array::scalar::Scalar;
2531
use vortex_array::validity::Validity;
32+
use vortex_buffer::Buffer;
33+
use vortex_buffer::ByteBuffer;
2634
use vortex_buffer::buffer;
2735
use vortex_error::VortexResult;
2836
use vortex_session::VortexSession;
@@ -845,3 +853,48 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> {
845853
assert_arrays_eq!(&compressed, &array, &mut exec_ctx);
846854
Ok(())
847855
}
856+
857+
/// Distinct 27-byte strings whose views reference a small prefix of one 2 MiB buffer, the way
858+
/// Arrow's Parquet reader shares a decompressed page between batches.
859+
fn views_into_shared_page(rows: u32) -> ArrayRef {
860+
const VALUE_LEN: u32 = 27;
861+
let mut page = vec![0u8; 2 << 20];
862+
for row in 0..rows {
863+
let value = format!("value-{row:021}");
864+
let start = (row * VALUE_LEN) as usize;
865+
page[start..start + VALUE_LEN as usize].copy_from_slice(value.as_bytes());
866+
}
867+
let page = ByteBuffer::from(page);
868+
let views = (0..rows)
869+
.map(|row| {
870+
let start = (row * VALUE_LEN) as usize;
871+
BinaryView::make_view(&page[start..start + VALUE_LEN as usize], 0, row * VALUE_LEN)
872+
})
873+
.collect::<Buffer<BinaryView>>();
874+
// SAFETY: every view points at `VALUE_LEN` bytes inside the page.
875+
unsafe {
876+
VarBinViewArray::new_unchecked(
877+
views,
878+
Arc::from([page]),
879+
DType::Utf8(Nullability::NonNullable),
880+
Validity::NonNullable,
881+
)
882+
}
883+
.into_array()
884+
}
885+
886+
#[test]
887+
fn scheme_must_beat_referenced_bytes_not_held_bytes() -> VortexResult<()> {
888+
let mut exec_ctx = SESSION.create_execution_ctx();
889+
let array = views_into_shared_page(2000);
890+
let exact = uncompressed_size_in_bytes(&array, UncompressedSizeOpts::exact(), &mut exec_ctx)?;
891+
892+
// Every value is distinct, so a dictionary is larger than the referenced bytes but far
893+
// smaller than the 2 MiB the array holds. It must lose.
894+
let compressed = compressor().compress(&array, &mut exec_ctx)?;
895+
896+
assert!(compressed.is::<VarBinView>());
897+
assert_eq!(compressed.nbytes(), exact as u64);
898+
assert_arrays_eq!(&compressed, &array, &mut exec_ctx);
899+
Ok(())
900+
}

‎vortex-file/src/tests.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ async fn test_eof_values() {
149149
#[rstest]
150150
#[case::default(
151151
BtrBlocksCompressorBuilder::from_session(&SESSION),
152-
if cfg!(feature = "zstd") { 70_076 } else { 70_012 }
152+
if cfg!(feature = "zstd") { 70_100 } else { 70_036 }
153153
)]
154154
#[cfg_attr(
155155
feature = "zstd",

‎vortex-layout/src/layouts/compressed.rs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ use futures::StreamExt as _;
88
use vortex_array::ArrayRef;
99
use vortex_array::ExecutionCtx;
1010
use vortex_array::VortexSessionExecute;
11+
use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
12+
use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts;
13+
use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes;
1114
use vortex_array::expr::stats::Stat;
1215
use vortex_btrblocks::BtrBlocksCompressor;
1316
use vortex_error::VortexResult;
@@ -110,8 +113,16 @@ impl LayoutStrategy for CompressingStrategy {
110113
handle.spawn_cpu(move || {
111114
let (sequence_id, chunk) = chunk?;
112115
let mut ctx = session.create_execution_ctx();
113-
// Compute the stats for the chunk prior to compression
116+
// Compute the stats for the chunk prior to compression. The exact
117+
// uncompressed size is measured once here and cached for the compressor.
114118
chunk.statistics().compute_all(&stats, &mut ctx)?;
119+
if UncompressedSizeInBytes::supports_dtype(chunk.dtype()) {
120+
uncompressed_size_in_bytes(
121+
&chunk,
122+
UncompressedSizeOpts::exact(),
123+
&mut ctx,
124+
)?;
125+
}
115126
Ok((sequence_id, compressor.compress_chunk(&chunk, &mut ctx)?))
116127
})
117128
})

0 commit comments

Comments
 (0)