From 6c2a06b049b5efd9ee6decf3911b5a2cb241230e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 16:02:32 +0800 Subject: [PATCH 1/2] refactor(encoding): pass mini-block compression context --- rust/lance-encoding/src/compression.rs | 39 ++++++++++----- .../src/encodings/logical/primitive.rs | 8 +++- .../encodings/logical/primitive/miniblock.rs | 29 +++++++++++- .../logical/primitive/sparse/writer.rs | 6 ++- .../src/encodings/physical/binary.rs | 9 +++- .../src/encodings/physical/bitpacking.rs | 8 +++- .../encodings/physical/byte_stream_split.rs | 20 ++++++-- .../src/encodings/physical/fsst.rs | 10 ++-- .../src/encodings/physical/general.rs | 34 ++++++++++---- .../src/encodings/physical/packed.rs | 10 ++-- .../src/encodings/physical/rle.rs | 47 +++++++++++-------- .../src/encodings/physical/value.rs | 22 ++++++--- 12 files changed, 176 insertions(+), 66 deletions(-) diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 5c15d48c25a..2ac2bedfc75 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -1372,6 +1372,13 @@ mod tests { use arrow_schema::{DataType, Field as ArrowField}; use std::collections::HashMap; + fn miniblock_context() + -> crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext { + crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext::new( + 0, true, true, + ) + } + fn create_test_field(name: &str, data_type: DataType) -> Field { let arrow_field = ArrowField::new(name, data_type, true); let mut field = Field::try_from(&arrow_field).unwrap(); @@ -1759,12 +1766,16 @@ mod tests { let compressor = strategy .create_miniblock_compressor(&field, &fixed_data) .unwrap(); - let (_block, encoding) = compressor.compress(fixed_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(fixed_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, false); let compressor = strategy .create_miniblock_compressor(&field, &variable_data) .unwrap(); - let (_block, encoding) = compressor.compress(variable_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(variable_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, true); // Test pervalue @@ -1794,13 +1805,17 @@ mod tests { let compressor = strategy .create_miniblock_compressor(&field, &fixed_data) .unwrap(); - let (_block, encoding) = compressor.compress(fixed_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(fixed_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, false); let compressor = strategy .create_miniblock_compressor(&field, &variable_data) .unwrap(); - let (_block, encoding) = compressor.compress(variable_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(variable_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, true); // Test pervalue @@ -2097,7 +2112,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2122,7 +2137,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(version); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}"); } } @@ -2150,7 +2165,7 @@ mod tests { let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2173,7 +2188,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2196,7 +2211,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8); } @@ -2233,7 +2248,7 @@ mod tests { let data = DataBlock::FixedWidth(data); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let rle = expect_rle_encoding(&encoding); assert!( @@ -2281,7 +2296,7 @@ mod tests { let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { panic!("expected RLE encoding"); }; @@ -2331,7 +2346,7 @@ mod tests { "expected RLE to beat inline bitpacking after child selection, got: {debug_str}" ); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let rle = expect_rle_encoding(&encoding); assert!(matches!( rle.values.as_ref().unwrap().compression.as_ref().unwrap(), diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 1b16ab73b80..cdcb0e0c482 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -55,7 +55,7 @@ use crate::{ encodings::logical::primitive::fullzip::PerValueDataBlock, }; use crate::{ - encodings::logical::primitive::miniblock::MiniBlockCompressed, + encodings::logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressionContext}, statistics::{ComputeStat, GetStat, Stat}, }; use crate::{ @@ -5282,7 +5282,11 @@ impl PrimitiveStructuralEncoder { let num_items = data.num_values(); let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; - let (compressed_data, value_encoding) = compressor.compress(data)?; + let common_chunk_buffers = + u64::from(repdef.rep_slicer().is_some()) + u64::from(repdef.def_slicer().is_some()); + let compression_context = + MiniBlockCompressionContext::new(common_chunk_buffers, support_large_chunk, true); + let (compressed_data, value_encoding) = compressor.compress(data, compression_context)?; let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16; diff --git a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs index edfba526670..9461229310b 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs @@ -54,6 +54,29 @@ pub struct MiniBlockCompressed { pub num_values: u64, } +/// Per-page framing details that can affect a mini-block compressor's choice. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MiniBlockCompressionContext { + common_chunk_buffers: u64, + support_large_chunk: bool, + allow_generic_offsets: bool, +} + +impl MiniBlockCompressionContext { + /// Creates the framing context supplied by the owning mini-block page. + pub fn new( + common_chunk_buffers: u64, + support_large_chunk: bool, + allow_generic_offsets: bool, + ) -> Self { + Self { + common_chunk_buffers, + support_large_chunk, + allow_generic_offsets, + } + } +} + /// Describes the size of a mini-block chunk of data /// /// Mini-block chunks are designed to be small (just a few disk sectors) @@ -113,7 +136,11 @@ pub trait MiniBlockCompressor: std::fmt::Debug + Send + Sync { /// /// This method also returns a description of the encoding applied that will be /// used at decode time to read the data. - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)>; + fn compress( + &self, + page: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)>; } #[cfg(test)] diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs index 3d82b4a5ef6..809536a2c1b 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -24,7 +24,8 @@ use super::{ SparseValidityMeaning, SparseValiditySet, }; use crate::encodings::logical::primitive::{ - FILL_BYTE, MINIBLOCK_ALIGNMENT, miniblock::MiniBlockCompressed, + FILL_BYTE, MINIBLOCK_ALIGNMENT, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext}, }; #[derive(Clone, Copy, Default)] @@ -563,7 +564,8 @@ pub fn prepare_values( let num_values = data.num_values(); let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; - let (compressed, value_compression) = compressor.compress(data)?; + let compression_context = MiniBlockCompressionContext::new(0, support_large_chunk, false); + let (compressed, value_compression) = compressor.compress(data, compression_context)?; let values = serialize_value_chunks(with_explicit_value_counts(compressed)?, support_large_chunk)?; Ok(PreparedSparseValues { diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index d02cf2da693..7c84c2774b7 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -21,7 +21,8 @@ use crate::buffer::LanceBuffer; use crate::data::{BlockInfo, DataBlock, VariableWidthBlock}; use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}; use crate::encodings::logical::primitive::miniblock::{ - MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, + MiniBlockCompressor, }; use crate::format::pb21::CompressiveEncoding; use crate::format::pb21::compressive_encoding::Compression; @@ -245,7 +246,11 @@ impl BinaryMiniBlockEncoder { } impl MiniBlockCompressor for BinaryMiniBlockEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::VariableWidth(variable_width) => Ok(self.chunk_data(variable_width)), _ => Err(Error::invalid_input_source( diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index be0b747e7dc..ad1b567edd7 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -28,7 +28,7 @@ use crate::compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompress use crate::data::BlockInfo; use crate::data::{DataBlock, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ - MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::pb21::CompressiveEncoding; use crate::format::{ProtobufUtils21, pb21}; @@ -215,7 +215,11 @@ impl InlineBitpacking { } impl MiniBlockCompressor for InlineBitpacking { - fn compress(&self, chunk: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + chunk: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match chunk { DataBlock::FixedWidth(fixed_width) => Ok(self.chunk_data(fixed_width)), _ => Err(Error::invalid_input_source( diff --git a/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs b/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs index c2b7aac9b9c..b861098d3e4 100644 --- a/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs +++ b/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs @@ -62,7 +62,7 @@ use crate::compression::MiniBlockDecompressor; use crate::compression_config::BssMode; use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ - MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::ProtobufUtils21; use crate::format::pb21::CompressiveEncoding; @@ -107,7 +107,11 @@ impl ByteStreamSplitEncoder { } impl MiniBlockCompressor for ByteStreamSplitEncoder { - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + page: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match page { DataBlock::FixedWidth(data) => { let num_values = data.num_values; @@ -345,7 +349,9 @@ mod tests { }); // Compress - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress let decompressed = decompressor @@ -391,7 +397,9 @@ mod tests { }); // Compress - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress let decompressed = decompressor @@ -424,7 +432,9 @@ mod tests { }); // Compress empty data - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress empty data let decompressed = decompressor.decompress(compressed.data, 0).unwrap(); diff --git a/rust/lance-encoding/src/encodings/physical/fsst.rs b/rust/lance-encoding/src/encodings/physical/fsst.rs index 8c1fe4141df..295ffcd88b1 100644 --- a/rust/lance-encoding/src/encodings/physical/fsst.rs +++ b/rust/lance-encoding/src/encodings/physical/fsst.rs @@ -23,7 +23,7 @@ use crate::{ data::{BlockInfo, DataBlock, VariableWidthBlock}, encodings::logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor}, }, format::{ ProtobufUtils21, @@ -138,7 +138,11 @@ impl FsstMiniBlockEncoder { } impl MiniBlockCompressor for FsstMiniBlockEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { let compressed = FsstCompressed::fsst_compress(data)?; let data_block = DataBlock::VariableWidth(compressed.data); @@ -148,7 +152,7 @@ impl MiniBlockCompressor for FsstMiniBlockEncoder { as Box; let (binary_miniblock_compressed, binary_array_encoding) = - binary_compressor.compress(data_block)?; + binary_compressor.compress(data_block, context)?; Ok(( binary_miniblock_compressed, diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index 53c61928870..769f308ccec 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -9,7 +9,9 @@ use crate::{ compression::MiniBlockDecompressor, data::DataBlock, encodings::{ - logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + logical::primitive::miniblock::{ + MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, + }, physical::block::{CompressionConfig, GeneralBufferCompressor}, }, format::{ProtobufUtils21, pb21::CompressiveEncoding}, @@ -35,9 +37,13 @@ const MIN_BUFFER_SIZE_FOR_COMPRESSION: usize = 4 * 1024; use super::super::logical::primitive::miniblock::MiniBlockChunk; impl MiniBlockCompressor for GeneralMiniBlockCompressor { - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + page: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { // First, compress with the inner compressor - let (inner_compressed, inner_encoding) = self.inner.compress(page)?; + let (inner_compressed, inner_encoding) = self.inner.compress(page, context)?; // Return the original encoding without compression if there's no data or // the first buffer is not large enough @@ -146,6 +152,10 @@ mod tests { use crate::format::pb21::compressive_encoding::Compression; use arrow_array::{Float64Array, Int32Array}; + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + #[derive(Debug)] struct TestCase { name: &'static str, @@ -249,7 +259,9 @@ mod tests { GeneralMiniBlockCompressor::new(test_case.inner_encoder, test_case.compression); // Compress the data - let (compressed, encoding) = compressor.compress(test_case.data).unwrap(); + let (compressed, encoding) = compressor + .compress(test_case.data, miniblock_context()) + .unwrap(); // Check if compression was applied as expected match &encoding.compression { @@ -461,7 +473,7 @@ mod tests { let compressor = GeneralMiniBlockCompressor::new(inner, compression); // Compress the data - let (compressed, encoding) = compressor.compress(block).unwrap(); + let (compressed, encoding) = compressor.compress(block, miniblock_context()).unwrap(); // Should get GeneralMiniBlock encoding since buffer is 4KB match &encoding.compression { @@ -503,7 +515,7 @@ mod tests { }, ); - let (compressed, _) = compressor.compress(data).unwrap(); + let (compressed, _) = compressor.compress(data, miniblock_context()).unwrap(); // RLE produces 2 buffers, but only the first one is compressed assert_eq!(compressed.data.len(), 2); } @@ -539,7 +551,9 @@ mod tests { }, ); - let (_compressed, encoding) = compressor.compress(test_32.data).unwrap(); + let (_compressed, encoding) = compressor + .compress(test_32.data, miniblock_context()) + .unwrap(); // Verify the encoding structure match &encoding.compression { @@ -596,7 +610,9 @@ mod tests { }, ); - let (_compressed_64, encoding_64) = compressor_64.compress(block_64).unwrap(); + let (_compressed_64, encoding_64) = compressor_64 + .compress(block_64, miniblock_context()) + .unwrap(); // Verify the encoding structure for 64-bit match &encoding_64.compression { @@ -650,7 +666,7 @@ mod tests { }, ); - let result = compressor.compress(empty_block); + let result = compressor.compress(empty_block, miniblock_context()); match result { Ok((compressed, _)) => { assert_eq!(compressed.num_values, 0); diff --git a/rust/lance-encoding/src/encodings/physical/packed.rs b/rust/lance-encoding/src/encodings/physical/packed.rs index ad2221dffed..96aa61e69d6 100644 --- a/rust/lance-encoding/src/encodings/physical/packed.rs +++ b/rust/lance-encoding/src/encodings/physical/packed.rs @@ -27,7 +27,7 @@ use crate::{ }, encodings::logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor}, }, format::{ ProtobufUtils21, @@ -73,7 +73,11 @@ fn struct_data_block_to_fixed_width_data_block( pub struct PackedStructFixedWidthMiniBlockEncoder {} impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::Struct(struct_data_block) => { let bits_per_values = struct_data_block.children.iter().map(|data_block| data_block.as_fixed_width_ref().unwrap().bits_per_value).collect::>(); @@ -84,7 +88,7 @@ impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { // store and transformed fixed-width data block. let value_miniblock_compressor = Box::new(ValueEncoder::default()) as Box; let (value_miniblock_compressed, value_array_encoding) = - value_miniblock_compressor.compress(data_block)?; + value_miniblock_compressor.compress(data_block, context)?; Ok(( value_miniblock_compressed, diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index b04716c2b44..0127817d08e 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -63,7 +63,7 @@ use crate::data::DataBlock; use crate::data::{BlockInfo, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, - MiniBlockCompressor, + MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::encodings::physical::block::{CompressionConfig, GeneralBufferCompressor}; use crate::format::ProtobufUtils21; @@ -1036,7 +1036,11 @@ impl RleChildCandidate { } impl MiniBlockCompressor for RleEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::FixedWidth(fixed_width) => { let num_values = fixed_width.num_values; @@ -1853,6 +1857,13 @@ mod tests { use arrow_array::Int32Array; use rstest::rstest; + fn compress_miniblock( + compressor: &dyn MiniBlockCompressor, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + compressor.compress(data, MiniBlockCompressionContext::new(0, true, true)) + } + fn expand_u16_runs(runs: &RleRuns) -> Vec { let mut expanded = Vec::with_capacity(runs.num_values()); for (value, length) in runs.iter() { @@ -2035,7 +2046,7 @@ mod tests { let array = Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 3, 3]); let data_block = DataBlock::from_array(array); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, data_block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, data_block).unwrap(); assert_eq!(compressed.num_values, 9); assert_eq!(compressed.chunks.len(), 1); @@ -2056,8 +2067,7 @@ mod tests { data.extend(&[100i32; 300]); // Will be split into 255+45 let array = Int32Array::from(data); - let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Should have 6 runs total (4 for first value, 2 for second) let lengths_buffer = &compressed.data[1]; @@ -2071,7 +2081,7 @@ mod tests { let data = vec![42i32; 1000]; let array = Int32Array::from(data); let (compressed, encoding) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); assert_eq!(compressed.data[0].len(), 4); assert_eq!(compressed.data[1].len(), 2); @@ -2112,7 +2122,7 @@ mod tests { RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false); let array = Int32Array::from(repeating_runs(1024, 4)); let (compressed, encoding) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); let rle = expect_rle(&encoding); assert!(matches!( @@ -2145,7 +2155,7 @@ mod tests { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false); let expected = repeating_runs(1024, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2181,7 +2191,7 @@ mod tests { use crate::encodings::physical::bitpacking::OutOfLineBitpacking; let expected = repeating_runs(1024, 4); - let (compressed, _) = MiniBlockCompressor::compress( + let (compressed, _) = compress_miniblock( &RleEncoder::new(), DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2275,7 +2285,7 @@ mod tests { false, ); let expected = repeating_runs(8192, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2306,7 +2316,7 @@ mod tests { fn test_rle_miniblock_bitpacks_values_child_when_smaller() { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); let expected = monotonic_runs(2048, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2336,7 +2346,7 @@ mod tests { fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); let expected = high_entropy_runs(2048, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2464,7 +2474,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, block).unwrap(); let decompressor = RleDecompressor::new(bits_per_value); let decompressed = MiniBlockDecompressor::decompress( &decompressor, @@ -2498,7 +2508,7 @@ mod tests { let array = Int32Array::from(data); let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Verify all non-last chunks have power-of-2 values for (i, chunk) in compressed.chunks.iter().enumerate() { @@ -2737,7 +2747,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, empty_block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, empty_block).unwrap(); assert_eq!(compressed.num_values, 0); assert!(compressed.data.is_empty()); @@ -2771,8 +2781,7 @@ mod tests { data.extend(vec![777i32; 2000]); let array = Int32Array::from(data.clone()); - let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Manually decompress all chunks let mut reconstructed = Vec::new(); @@ -2885,7 +2894,7 @@ mod tests { // Compress the data let array = Int32Array::from(data.clone()); let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Decompress and verify match MiniBlockDecompressor::decompress( @@ -2955,7 +2964,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, block).unwrap(); // Debug first few chunks for (i, chunk) in compressed.chunks.iter().take(5).enumerate() { diff --git a/rust/lance-encoding/src/encodings/physical/value.rs b/rust/lance-encoding/src/encodings/physical/value.rs index 606f49b699a..8b7c385f601 100644 --- a/rust/lance-encoding/src/encodings/physical/value.rs +++ b/rust/lance-encoding/src/encodings/physical/value.rs @@ -13,7 +13,7 @@ use crate::data::{ use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}; use crate::encodings::logical::primitive::miniblock::{ MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, - MiniBlockCompressor, + MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::ProtobufUtils21; use crate::format::pb21::compressive_encoding::Compression; @@ -471,7 +471,11 @@ impl BlockCompressor for ValueEncoder { } impl MiniBlockCompressor for ValueEncoder { - fn compress(&self, chunk: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + chunk: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match chunk { DataBlock::FixedWidth(fixed_width) => { let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); @@ -775,7 +779,7 @@ mod tests { encodings::{ logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::MiniBlockCompressor, + miniblock::{MiniBlockCompressionContext, MiniBlockCompressor}, }, physical::value::ValueDecompressor, }, @@ -789,6 +793,10 @@ mod tests { use super::ValueEncoder; + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + const PRIMITIVE_TYPES: &[DataType] = &[ DataType::Null, DataType::FixedSizeBinary(2), @@ -969,7 +977,8 @@ mod tests { let starting_data = DataBlock::from_array(sample_list.clone()); let encoder = ValueEncoder::default(); - let (data, compression) = MiniBlockCompressor::compress(&encoder, starting_data).unwrap(); + let (data, compression) = + MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()).unwrap(); assert_eq!(data.num_values, 3); assert_eq!(data.data.len(), 3); @@ -1030,7 +1039,7 @@ mod tests { let starting_data = DataBlock::from_array(array); let encoder = ValueEncoder::default(); - let result = MiniBlockCompressor::compress(&encoder, starting_data); + let result = MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()); let err = result.expect_err("wide values should not be encodable as miniblock"); assert!( @@ -1142,7 +1151,8 @@ mod tests { ); let encoder = ValueEncoder::default(); - let (data, compression) = MiniBlockCompressor::compress(&encoder, starting_data).unwrap(); + let (data, compression) = + MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()).unwrap(); let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { panic!() From 92f9370741f36b209acd114875fb6e8c856f87c3 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 30 Jul 2026 16:44:32 +0800 Subject: [PATCH 2/2] refactor(encoding): pass mini-block context first --- rust/lance-encoding/src/compression.rs | 32 +++++++++---------- .../src/encodings/logical/primitive.rs | 2 +- .../encodings/logical/primitive/miniblock.rs | 2 +- .../logical/primitive/sparse/writer.rs | 2 +- .../src/encodings/physical/binary.rs | 2 +- .../src/encodings/physical/bitpacking.rs | 2 +- .../encodings/physical/byte_stream_split.rs | 8 ++--- .../src/encodings/physical/fsst.rs | 4 +-- .../src/encodings/physical/general.rs | 16 +++++----- .../src/encodings/physical/packed.rs | 4 +-- .../src/encodings/physical/rle.rs | 4 +-- .../src/encodings/physical/value.rs | 8 ++--- 12 files changed, 42 insertions(+), 44 deletions(-) diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 2ac2bedfc75..7d147b2b4e3 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -1367,16 +1367,14 @@ mod tests { use super::*; use crate::buffer::LanceBuffer; use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; + use crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext; use crate::statistics::ComputeStat; use crate::testing::extract_array_encoding_chain; use arrow_schema::{DataType, Field as ArrowField}; use std::collections::HashMap; - fn miniblock_context() - -> crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext { - crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext::new( - 0, true, true, - ) + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) } fn create_test_field(name: &str, data_type: DataType) -> Field { @@ -1767,14 +1765,14 @@ mod tests { .create_miniblock_compressor(&field, &fixed_data) .unwrap(); let (_block, encoding) = compressor - .compress(fixed_data.clone(), miniblock_context()) + .compress(miniblock_context(), fixed_data.clone()) .unwrap(); check_uncompressed_encoding(&encoding, false); let compressor = strategy .create_miniblock_compressor(&field, &variable_data) .unwrap(); let (_block, encoding) = compressor - .compress(variable_data.clone(), miniblock_context()) + .compress(miniblock_context(), variable_data.clone()) .unwrap(); check_uncompressed_encoding(&encoding, true); @@ -1806,7 +1804,7 @@ mod tests { .create_miniblock_compressor(&field, &fixed_data) .unwrap(); let (_block, encoding) = compressor - .compress(fixed_data.clone(), miniblock_context()) + .compress(miniblock_context(), fixed_data.clone()) .unwrap(); check_uncompressed_encoding(&encoding, false); @@ -1814,7 +1812,7 @@ mod tests { .create_miniblock_compressor(&field, &variable_data) .unwrap(); let (_block, encoding) = compressor - .compress(variable_data.clone(), miniblock_context()) + .compress(miniblock_context(), variable_data.clone()) .unwrap(); check_uncompressed_encoding(&encoding, true); @@ -2112,7 +2110,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2137,7 +2135,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(version); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}"); } } @@ -2165,7 +2163,7 @@ mod tests { let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); - let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2188,7 +2186,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2211,7 +2209,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8); } @@ -2248,7 +2246,7 @@ mod tests { let data = DataBlock::FixedWidth(data); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); let rle = expect_rle_encoding(&encoding); assert!( @@ -2296,7 +2294,7 @@ mod tests { let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); - let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { panic!("expected RLE encoding"); }; @@ -2346,7 +2344,7 @@ mod tests { "expected RLE to beat inline bitpacking after child selection, got: {debug_str}" ); - let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap(); let rle = expect_rle_encoding(&encoding); assert!(matches!( rle.values.as_ref().unwrap().compression.as_ref().unwrap(), diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index cdcb0e0c482..3267366b288 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -5286,7 +5286,7 @@ impl PrimitiveStructuralEncoder { u64::from(repdef.rep_slicer().is_some()) + u64::from(repdef.def_slicer().is_some()); let compression_context = MiniBlockCompressionContext::new(common_chunk_buffers, support_large_chunk, true); - let (compressed_data, value_encoding) = compressor.compress(data, compression_context)?; + let (compressed_data, value_encoding) = compressor.compress(compression_context, data)?; let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16; diff --git a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs index 9461229310b..0ee408a8510 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs @@ -138,8 +138,8 @@ pub trait MiniBlockCompressor: std::fmt::Debug + Send + Sync { /// used at decode time to read the data. fn compress( &self, - page: DataBlock, context: MiniBlockCompressionContext, + page: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)>; } diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs index 809536a2c1b..e8365451094 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -565,7 +565,7 @@ pub fn prepare_values( let num_values = data.num_values(); let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; let compression_context = MiniBlockCompressionContext::new(0, support_large_chunk, false); - let (compressed, value_compression) = compressor.compress(data, compression_context)?; + let (compressed, value_compression) = compressor.compress(compression_context, data)?; let values = serialize_value_chunks(with_explicit_value_counts(compressed)?, support_large_chunk)?; Ok(PreparedSparseValues { diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index 7c84c2774b7..a4cda602aa1 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -248,8 +248,8 @@ impl BinaryMiniBlockEncoder { impl MiniBlockCompressor for BinaryMiniBlockEncoder { fn compress( &self, - data: DataBlock, _context: MiniBlockCompressionContext, + data: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::VariableWidth(variable_width) => Ok(self.chunk_data(variable_width)), diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index ad1b567edd7..5a749cee6e7 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -217,8 +217,8 @@ impl InlineBitpacking { impl MiniBlockCompressor for InlineBitpacking { fn compress( &self, - chunk: DataBlock, _context: MiniBlockCompressionContext, + chunk: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match chunk { DataBlock::FixedWidth(fixed_width) => Ok(self.chunk_data(fixed_width)), diff --git a/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs b/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs index b861098d3e4..683e4ad8236 100644 --- a/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs +++ b/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs @@ -109,8 +109,8 @@ impl ByteStreamSplitEncoder { impl MiniBlockCompressor for ByteStreamSplitEncoder { fn compress( &self, - page: DataBlock, _context: MiniBlockCompressionContext, + page: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match page { DataBlock::FixedWidth(data) => { @@ -350,7 +350,7 @@ mod tests { // Compress let (compressed, _encoding) = encoder - .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .compress(MiniBlockCompressionContext::new(0, true, true), data_block) .unwrap(); // Decompress @@ -398,7 +398,7 @@ mod tests { // Compress let (compressed, _encoding) = encoder - .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .compress(MiniBlockCompressionContext::new(0, true, true), data_block) .unwrap(); // Decompress @@ -433,7 +433,7 @@ mod tests { // Compress empty data let (compressed, _encoding) = encoder - .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .compress(MiniBlockCompressionContext::new(0, true, true), data_block) .unwrap(); // Decompress empty data diff --git a/rust/lance-encoding/src/encodings/physical/fsst.rs b/rust/lance-encoding/src/encodings/physical/fsst.rs index 295ffcd88b1..92189dbf40a 100644 --- a/rust/lance-encoding/src/encodings/physical/fsst.rs +++ b/rust/lance-encoding/src/encodings/physical/fsst.rs @@ -140,8 +140,8 @@ impl FsstMiniBlockEncoder { impl MiniBlockCompressor for FsstMiniBlockEncoder { fn compress( &self, - data: DataBlock, context: MiniBlockCompressionContext, + data: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { let compressed = FsstCompressed::fsst_compress(data)?; @@ -152,7 +152,7 @@ impl MiniBlockCompressor for FsstMiniBlockEncoder { as Box; let (binary_miniblock_compressed, binary_array_encoding) = - binary_compressor.compress(data_block, context)?; + binary_compressor.compress(context, data_block)?; Ok(( binary_miniblock_compressed, diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index 769f308ccec..479a47c36c4 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -39,11 +39,11 @@ use super::super::logical::primitive::miniblock::MiniBlockChunk; impl MiniBlockCompressor for GeneralMiniBlockCompressor { fn compress( &self, - page: DataBlock, context: MiniBlockCompressionContext, + page: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { // First, compress with the inner compressor - let (inner_compressed, inner_encoding) = self.inner.compress(page, context)?; + let (inner_compressed, inner_encoding) = self.inner.compress(context, page)?; // Return the original encoding without compression if there's no data or // the first buffer is not large enough @@ -260,7 +260,7 @@ mod tests { // Compress the data let (compressed, encoding) = compressor - .compress(test_case.data, miniblock_context()) + .compress(miniblock_context(), test_case.data) .unwrap(); // Check if compression was applied as expected @@ -473,7 +473,7 @@ mod tests { let compressor = GeneralMiniBlockCompressor::new(inner, compression); // Compress the data - let (compressed, encoding) = compressor.compress(block, miniblock_context()).unwrap(); + let (compressed, encoding) = compressor.compress(miniblock_context(), block).unwrap(); // Should get GeneralMiniBlock encoding since buffer is 4KB match &encoding.compression { @@ -515,7 +515,7 @@ mod tests { }, ); - let (compressed, _) = compressor.compress(data, miniblock_context()).unwrap(); + let (compressed, _) = compressor.compress(miniblock_context(), data).unwrap(); // RLE produces 2 buffers, but only the first one is compressed assert_eq!(compressed.data.len(), 2); } @@ -552,7 +552,7 @@ mod tests { ); let (_compressed, encoding) = compressor - .compress(test_32.data, miniblock_context()) + .compress(miniblock_context(), test_32.data) .unwrap(); // Verify the encoding structure @@ -611,7 +611,7 @@ mod tests { ); let (_compressed_64, encoding_64) = compressor_64 - .compress(block_64, miniblock_context()) + .compress(miniblock_context(), block_64) .unwrap(); // Verify the encoding structure for 64-bit @@ -666,7 +666,7 @@ mod tests { }, ); - let result = compressor.compress(empty_block, miniblock_context()); + let result = compressor.compress(miniblock_context(), empty_block); match result { Ok((compressed, _)) => { assert_eq!(compressed.num_values, 0); diff --git a/rust/lance-encoding/src/encodings/physical/packed.rs b/rust/lance-encoding/src/encodings/physical/packed.rs index 96aa61e69d6..6c97f368178 100644 --- a/rust/lance-encoding/src/encodings/physical/packed.rs +++ b/rust/lance-encoding/src/encodings/physical/packed.rs @@ -75,8 +75,8 @@ pub struct PackedStructFixedWidthMiniBlockEncoder {} impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { fn compress( &self, - data: DataBlock, context: MiniBlockCompressionContext, + data: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::Struct(struct_data_block) => { @@ -88,7 +88,7 @@ impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { // store and transformed fixed-width data block. let value_miniblock_compressor = Box::new(ValueEncoder::default()) as Box; let (value_miniblock_compressed, value_array_encoding) = - value_miniblock_compressor.compress(data_block, context)?; + value_miniblock_compressor.compress(context, data_block)?; Ok(( value_miniblock_compressed, diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index 0127817d08e..cac33d7ecf6 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -1038,8 +1038,8 @@ impl RleChildCandidate { impl MiniBlockCompressor for RleEncoder { fn compress( &self, - data: DataBlock, _context: MiniBlockCompressionContext, + data: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::FixedWidth(fixed_width) => { @@ -1861,7 +1861,7 @@ mod tests { compressor: &dyn MiniBlockCompressor, data: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { - compressor.compress(data, MiniBlockCompressionContext::new(0, true, true)) + compressor.compress(MiniBlockCompressionContext::new(0, true, true), data) } fn expand_u16_runs(runs: &RleRuns) -> Vec { diff --git a/rust/lance-encoding/src/encodings/physical/value.rs b/rust/lance-encoding/src/encodings/physical/value.rs index 8b7c385f601..4aa701d6d5f 100644 --- a/rust/lance-encoding/src/encodings/physical/value.rs +++ b/rust/lance-encoding/src/encodings/physical/value.rs @@ -473,8 +473,8 @@ impl BlockCompressor for ValueEncoder { impl MiniBlockCompressor for ValueEncoder { fn compress( &self, - chunk: DataBlock, _context: MiniBlockCompressionContext, + chunk: DataBlock, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match chunk { DataBlock::FixedWidth(fixed_width) => { @@ -978,7 +978,7 @@ mod tests { let encoder = ValueEncoder::default(); let (data, compression) = - MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()).unwrap(); + MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data).unwrap(); assert_eq!(data.num_values, 3); assert_eq!(data.data.len(), 3); @@ -1039,7 +1039,7 @@ mod tests { let starting_data = DataBlock::from_array(array); let encoder = ValueEncoder::default(); - let result = MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()); + let result = MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data); let err = result.expect_err("wide values should not be encodable as miniblock"); assert!( @@ -1152,7 +1152,7 @@ mod tests { let encoder = ValueEncoder::default(); let (data, compression) = - MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()).unwrap(); + MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data).unwrap(); let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { panic!()