This is an automated email from the ASF dual-hosted git repository. tisonkun pushed a commit to branch codex/direct-cpc-serde in repository https://gitbox.apache.org/repos/asf/datasketches-rust.git
commit 34565d0354470e7711e7d3c39a16165bdfcd623e Author: tison <[email protected]> AuthorDate: Wed Aug 26 17:26:27 2026 +0800 perf(cpc): encode and decode payloads directly --- CHANGELOG.md | 1 + datasketches/src/codec/encode.rs | 12 + datasketches/src/cpc/compression.rs | 824 ++++++++++++------------------------ datasketches/src/cpc/pair_table.rs | 4 + datasketches/src/cpc/sketch.rs | 120 +++--- 5 files changed, 358 insertions(+), 603 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 269fa4d..eba1dae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ All significant changes to this project will be documented in this file. ### Performance improvements * Reduce T-Digest allocation overhead and retained memory across updates, compression, merges, serialization, deserialization, and freezing while preserving the serialized format. +* Reduce CPC serialization and deserialization allocations by encoding directly into the output buffer and decoding directly from the input payload. ### Bug fixes diff --git a/datasketches/src/codec/encode.rs b/datasketches/src/codec/encode.rs index 20ac0ea..fa3cd75 100644 --- a/datasketches/src/codec/encode.rs +++ b/datasketches/src/codec/encode.rs @@ -33,6 +33,18 @@ impl SketchBytes { self.bytes } + /// Returns the number of bytes written so far. + #[cfg(feature = "cpc")] + pub(crate) fn len(&self) -> usize { + self.bytes.len() + } + + /// Overwrites a previously written 32-bit unsigned integer in little-endian byte order. + #[cfg(feature = "cpc")] + pub(crate) fn overwrite_u32_le(&mut self, offset: usize, n: u32) { + self.bytes[offset..offset + size_of::<u32>()].copy_from_slice(&n.to_le_bytes()); + } + /// Writes the given byte slice to the `SketchBytes`. pub fn write(&mut self, buf: &[u8]) { self.bytes.extend_from_slice(buf); diff --git a/datasketches/src/cpc/compression.rs b/datasketches/src/cpc/compression.rs index db07f8e..281a7cf 100644 --- a/datasketches/src/cpc/compression.rs +++ b/datasketches/src/cpc/compression.rs @@ -17,6 +17,7 @@ use std::cmp::Ordering; +use crate::codec::SketchBytes; use crate::cpc::CpcSketch; use crate::cpc::Flavor; use crate::cpc::compression_data::COLUMN_PERMUTATIONS_FOR_DECODING; @@ -30,521 +31,329 @@ use crate::cpc::determine_flavor; use crate::cpc::pair_table::PairTable; use crate::error::Error; -#[derive(Default)] -pub(super) struct CompressedState { - pub(super) table_data: Vec<u32>, - pub(super) table_data_words: usize, - // can be different from the number of entries in the sketch in hybrid mode - pub(super) table_num_entries: u32, - pub(super) window_data: Vec<u32>, - pub(super) window_data_words: usize, -} - -impl CompressedState { - pub fn compress(&mut self, source: &CpcSketch) { - match source.flavor() { - Flavor::Empty => { - // do nothing - } - Flavor::Sparse => { - self.compress_sparse_flavor(source); - debug_assert!(self.window_data.is_empty(), "window is not expected"); - debug_assert!(!self.table_data.is_empty(), "table is expected"); - } - Flavor::Hybrid => { - self.compress_hybrid_flavor(source); - debug_assert!(self.window_data.is_empty(), "window is not expected"); - debug_assert!(!self.table_data.is_empty(), "table is expected"); - } - Flavor::Pinned => { - self.compress_pinned_flavor(source); - debug_assert!(!self.window_data.is_empty(), "window is expected"); - } - Flavor::Sliding => { - self.compress_sliding_flavor(source); - debug_assert!(!self.window_data.is_empty(), "window is expected"); - } +/// Appends the compressed payload directly to the serialized sketch. +/// +/// The return value is `(table_words, window_words)`. The payload itself is written in wire order: +/// window first, then table. +pub(super) fn encode_payload(source: &CpcSketch, output: &mut SketchBytes) -> (usize, usize) { + match source.flavor() { + Flavor::Empty => (0, 0), + Flavor::Sparse => { + debug_assert!(source.sliding_window.is_empty()); + let mut pairs = source.surprising_value_table().unwrapping_get_items(); + pairs.sort_unstable(); + (compress_surprising_values(&pairs, source.lg_k(), output), 0) } - } + Flavor::Hybrid => { + debug_assert!(!source.sliding_window.is_empty()); + debug_assert_eq!(source.window_offset, 0); - fn compress_sparse_flavor(&mut self, source: &CpcSketch) { - debug_assert!(source.sliding_window.is_empty()); - let mut pairs = source.surprising_value_table().unwrapping_get_items(); - pairs.sort_unstable(); - self.compress_surprising_values(&pairs, source.lg_k()); - } + let k = 1 << source.lg_k(); + let mut pairs = source.surprising_value_table().unwrapping_get_items(); + pairs.sort_unstable(); + let num_pairs_from_table = pairs.len(); + let mut all_pairs = vec![0; source.num_coupons() as usize]; - fn compress_hybrid_flavor(&mut self, source: &CpcSketch) { - debug_assert!(!source.sliding_window.is_empty()); - debug_assert_eq!(source.window_offset, 0); - - let k = 1 << source.lg_k(); - let mut pairs = source.surprising_value_table().unwrapping_get_items(); - pairs.sort_unstable(); - let num_pairs_from_table = pairs.len(); - let num_pairs_from_window = (source.num_coupons() as usize) - num_pairs_from_table; - - let all_pairs_len = num_pairs_from_table + num_pairs_from_window; - let mut all_pairs = vec![0; all_pairs_len]; - // tricky read pairs from sliding_window - { - // The empty space that this leaves at the beginning of the output array will be filled - // later. - let mut idx = num_pairs_from_table; + let mut index = num_pairs_from_table; for row_index in 0..k { let mut byte = source.sliding_window[row_index]; while byte != 0 { let col_index = byte.trailing_zeros(); - byte ^= 1 << col_index; // erase the 1 - all_pairs[idx] = ((row_index << 6) as u32) | col_index; - idx += 1; + byte ^= 1 << col_index; + all_pairs[index] = ((row_index << 6) as u32) | col_index; + index += 1; } } - assert_eq!(idx, all_pairs_len); - } - // two-way merge of pairs_from_table and pairs_from_window into all_pairs - { - let mut final_idx = 0; - let mut table_idx = 0; - let mut window_idx = num_pairs_from_table; - - while final_idx < all_pairs_len { - if table_idx < num_pairs_from_table - && (window_idx >= all_pairs_len || pairs[table_idx] <= all_pairs[window_idx]) + assert_eq!(index, all_pairs.len()); + + let mut table_index = 0; + let mut window_index = num_pairs_from_table; + for final_index in 0..all_pairs.len() { + if table_index < num_pairs_from_table + && (window_index >= all_pairs.len() + || pairs[table_index] <= all_pairs[window_index]) { - all_pairs[final_idx] = pairs[table_idx]; - table_idx += 1; + all_pairs[final_index] = pairs[table_index]; + table_index += 1; } else { - all_pairs[final_idx] = all_pairs[window_idx]; - window_idx += 1; + all_pairs[final_index] = all_pairs[window_index]; + window_index += 1; } - final_idx += 1; } - } - - self.compress_surprising_values(&all_pairs, source.lg_k()); - } - - fn compress_pinned_flavor(&mut self, source: &CpcSketch) { - self.compress_sliding_window(&source.sliding_window, source.lg_k(), source.num_coupons()); - let mut pairs = source.surprising_value_table().unwrapping_get_items(); - if !pairs.is_empty() { - // Here we subtract 8 from the column indices. Because they are stored in the low 6 bits - // of each row_col pair, and because no column index is less than 8 for a "Pinned" - // sketch, we can simply subtract 8 from the pairs themselves. - // shift the columns over by 8 positions before compressing (because of the window) + ( + compress_surprising_values(&all_pairs, source.lg_k(), output), + 0, + ) + } + Flavor::Pinned => { + let window_words = compress_sliding_window(source, output); + let mut pairs = source.surprising_value_table().unwrapping_get_items(); for pair in &mut pairs { assert!(*pair & 63 >= 8, "pair column index is less than 8: {pair}"); *pair -= 8; } - pairs.sort_unstable(); - self.compress_surprising_values(&pairs, source.lg_k()); + let table_words = if pairs.is_empty() { + 0 + } else { + compress_surprising_values(&pairs, source.lg_k(), output) + }; + (table_words, window_words) } - } - - // Complicated by the existence of both a left fringe and a right fringe. - fn compress_sliding_flavor(&mut self, source: &CpcSketch) { - self.compress_sliding_window(&source.sliding_window, source.lg_k(), source.num_coupons()); - let mut pairs = source.surprising_value_table().unwrapping_get_items(); - if !pairs.is_empty() { - // Here we apply a complicated transformation to the column indices, which - // changes the implied ordering of the pairs, so we must do it before sorting. - - let pseudo_phase = determine_pseudo_phase(source.lg_k(), source.num_coupons()); - let permutation = &COLUMN_PERMUTATIONS_FOR_ENCODING[pseudo_phase as usize]; - let offset = source.window_offset; - debug_assert!(offset <= 56); - for pair in &mut pairs { - let row_col = *pair; - let row = row_col >> 6; - let mut col = (row_col & 63) as u8; - // first rotate the columns into a canonical configuration: - // new = ((old - (offset+8)) + 64) mod 64 - col = (col + 56 - offset) & 63; - debug_assert!(col < 56); - // then apply the permutation - col = permutation[col as usize]; - *pair = (row << 6) | (col as u32); + Flavor::Sliding => { + let window_words = compress_sliding_window(source, output); + let mut pairs = source.surprising_value_table().unwrapping_get_items(); + if !pairs.is_empty() { + let pseudo_phase = determine_pseudo_phase(source.lg_k(), source.num_coupons()); + let permutation = &COLUMN_PERMUTATIONS_FOR_ENCODING[pseudo_phase as usize]; + let offset = source.window_offset; + debug_assert!(offset <= 56); + for pair in &mut pairs { + let row = *pair >> 6; + let col = ((*pair & 63) as u8 + 56 - offset) & 63; + debug_assert!(col < 56); + *pair = (row << 6) | u32::from(permutation[col as usize]); + } + pairs.sort_unstable(); } - - pairs.sort_unstable(); - self.compress_surprising_values(&pairs, source.lg_k()); + let table_words = if pairs.is_empty() { + 0 + } else { + compress_surprising_values(&pairs, source.lg_k(), output) + }; + (table_words, window_words) } } +} - fn compress_surprising_values(&mut self, pairs: &[u32], lg_k: u8) { - let k = 1 << lg_k; - let num_pairs = pairs.len() as u32; - let num_base_bits = golomb_choose_number_of_base_bits(k + num_pairs, num_pairs as u64); - let table_len = safe_length_for_compressed_pair_buf(k, num_pairs, num_base_bits); - self.table_data.resize(table_len, 0); - - let compressed_surprising_values = self.low_level_compress_pairs(pairs, num_base_bits); +fn compress_sliding_window(source: &CpcSketch, output: &mut SketchBytes) -> usize { + let pseudo_phase = determine_pseudo_phase(source.lg_k(), source.num_coupons()); + low_level_compress_bytes( + &source.sliding_window, + 1 << source.lg_k(), + &ENCODING_TABLES_FOR_HIGH_ENTROPY_BYTE[pseudo_phase as usize], + output, + ) +} - // At this point we could free the unused portion of the compression output buffer, - // but it is not necessary if it is temporary - // Note: realloc caused strange timing spikes for lgK = 11 and 12. +fn compress_surprising_values(pairs: &[u32], lg_k: u8, output: &mut SketchBytes) -> usize { + let num_pairs = pairs.len() as u32; + let num_base_bits = + golomb_choose_number_of_base_bits((1 << lg_k) + num_pairs, u64::from(num_pairs)); + low_level_compress_pairs(pairs, num_base_bits, output) +} - self.table_data_words = compressed_surprising_values; - self.table_num_entries = num_pairs; +fn low_level_compress_bytes( + bytes: &[u8], + num_bytes_to_encode: u32, + encoding_table: &[u16], + output: &mut SketchBytes, +) -> usize { + let mut bits = BitWriter::new(output); + for &byte in &bytes[..num_bytes_to_encode as usize] { + let code_info = encoding_table[byte as usize]; + bits.write(u64::from(code_info & 0xfff), (code_info >> 12) as u8); } + bits.pad(11); + bits.finish() +} - fn compress_sliding_window(&mut self, window: &[u8], lg_k: u8, num_coupons: u32) { - let k = 1 << lg_k; - let window_buf_len = safe_length_for_compressed_window_buf(k); - self.window_data.resize(window_buf_len, 0); - let pseudo_phase = determine_pseudo_phase(lg_k, num_coupons); - let data_words = self.low_level_compress_bytes( - window, - k, - &ENCODING_TABLES_FOR_HIGH_ENTROPY_BYTE[pseudo_phase as usize], - ); - - // At this point we could free the unused portion of the compression output buffer, - // but it is not necessary if it is temporary - // Note: realloc caused strange timing spikes for lgK = 11 and 12. - - self.window_data_words = data_words; - } +fn low_level_compress_pairs(pairs: &[u32], num_base_bits: u8, output: &mut SketchBytes) -> usize { + let mut bits = BitWriter::new(output); + let golomb_lo_mask = (1 << num_base_bits) - 1; + let mut predicted_row_index = 0; + let mut predicted_col_index = 0; - /// Returns the number of compressed words that were actually used. - /// - /// It is the caller's responsibility to ensure that the window_data is long enough. - fn low_level_compress_bytes( - &mut self, - byte_array: &[u8], - num_bytes_to_encode: u32, - encoding_table: &[u16], - ) -> usize { - // bits are packed into this first, then are flushed to window_data - let mut bitbuf = 0; - // number of bits currently in bitbuf; must be between 0 and 31 - let mut bufbits = 0; - let mut next_word_index = 0; - - for byte_index in 0..num_bytes_to_encode { - let code_info = encoding_table[byte_array[byte_index as usize] as usize]; - let code_val = (code_info & 0xfff) as u64; - let code_len = (code_info >> 12) as u8; - bitbuf |= code_val << bufbits; - bufbits += code_len; - maybe_flush_bitbuf( - &mut bitbuf, - &mut bufbits, - &mut self.window_data, - &mut next_word_index, - ); + for &row_col in pairs { + let row_index = row_col >> 6; + let col_index = row_col & 63; + if row_index != predicted_row_index { + predicted_col_index = 0; } + assert!(row_index >= predicted_row_index); + assert!(col_index >= predicted_col_index); - // Pad the bitstream with 11 zero-bits so that the decompressor's 12-bit peek can't overrun - // its input. - bufbits += 11; - maybe_flush_bitbuf( - &mut bitbuf, - &mut bufbits, - &mut self.window_data, - &mut next_word_index, - ); - - if bufbits > 0 { - // We are done encoding now, so we flush the bit buffer. - debug_assert!(bufbits < 32); - self.window_data[next_word_index] = (bitbuf & 0xffffffff) as u32; - next_word_index += 1; - - // not really necessary unset since no more use - //bitbuf = 0; - //bufbits = 0; - } + let y_delta = row_index - predicted_row_index; + let x_delta = col_index - predicted_col_index; + predicted_row_index = row_index; + predicted_col_index = col_index + 1; - next_word_index + let code_info = LENGTH_LIMITED_UNARY_ENCODING_TABLE65[x_delta as usize]; + bits.write(u64::from(code_info & 0xfff), (code_info >> 12) as u8); + bits.write_unary(u64::from(y_delta >> num_base_bits)); + bits.write(u64::from(y_delta & golomb_lo_mask), num_base_bits); } - /// Returns the number of table_data actually used. - /// - /// Here "pairs" refers to row/column pairs that specify the positions of surprising values in - /// the bit matrix. - fn low_level_compress_pairs(&mut self, pairs: &[u32], num_base_bits: u8) -> usize { - let mut bitbuf = 0; - let mut bufbits = 0; - let mut next_word_index = 0; - let golomb_lo_mask = ((1 << num_base_bits) - 1) as u64; - let mut predicted_row_index = 0; - let mut predicted_col_index = 0; - - for pair_index in 0..pairs.len() { - let row_col = pairs[pair_index]; - let row_index = row_col >> 6; - let col_index = row_col & 63; - - if row_index != predicted_row_index { - predicted_col_index = 0; - } - - assert!(row_index >= predicted_row_index); - assert!(col_index >= predicted_col_index); - - let y_delta = row_index - predicted_row_index; - let x_delta = col_index - predicted_col_index; - - predicted_row_index = row_index; - predicted_col_index = col_index + 1; - - let code_info = LENGTH_LIMITED_UNARY_ENCODING_TABLE65[x_delta as usize]; - let code_val = (code_info & 0xfff) as u64; - let code_len = (code_info >> 12) as u8; - bitbuf |= code_val << bufbits; - bufbits += code_len; - - maybe_flush_bitbuf( - &mut bitbuf, - &mut bufbits, - &mut self.table_data, - &mut next_word_index, - ); - - let golomb_lo = (y_delta as u64) & golomb_lo_mask; - let golomb_hi = (y_delta as u64) >> num_base_bits; - write_unary( - &mut self.table_data, - &mut next_word_index, - &mut bitbuf, - &mut bufbits, - golomb_hi, - ); - - bitbuf |= golomb_lo << bufbits; - bufbits += num_base_bits; - maybe_flush_bitbuf( - &mut bitbuf, - &mut bufbits, - &mut self.table_data, - &mut next_word_index, - ); - } - - // Pad the bitstream so that the decompressor's 12-bit peek can't overrun its input. - let padding = 10u8.saturating_sub(num_base_bits); - bufbits += padding; - maybe_flush_bitbuf( - &mut bitbuf, - &mut bufbits, - &mut self.table_data, - &mut next_word_index, - ); - - if bufbits > 0 { - // We are done encoding now, so we flush the bit buffer - assert!(bufbits < 32); - self.table_data[next_word_index] = (bitbuf & 0xffffffff) as u32; - next_word_index += 1; - - // not really necessary unset since no more use - //bitbuf = 0; - //bufbits = 0; - } - - next_word_index - } + bits.pad(10u8.saturating_sub(num_base_bits)); + bits.finish() } -pub(super) struct UncompressedState { - pub(super) table: PairTable, - pub(super) window: Vec<u8>, +struct BitWriter<'a> { + output: &'a mut SketchBytes, + buffer: u64, + buffered_bits: u8, + words_written: usize, } -impl CompressedState { - pub fn uncompress(&self, lg_k: u8, num_coupons: u32) -> Result<UncompressedState, Error> { - match determine_flavor(lg_k, num_coupons) { - Flavor::Empty => Ok(UncompressedState { - table: PairTable::new(2, lg_k + 6), - window: vec![], - }), - Flavor::Sparse => self.uncompress_sparse_flavor(lg_k), - Flavor::Hybrid => self.uncompress_hybrid_flavor(lg_k), - Flavor::Pinned => self.uncompress_pinned_flavor(lg_k, num_coupons), - Flavor::Sliding => self.uncompress_sliding_flavor(lg_k, num_coupons), +impl<'a> BitWriter<'a> { + fn new(output: &'a mut SketchBytes) -> Self { + Self { + output, + buffer: 0, + buffered_bits: 0, + words_written: 0, } } - fn uncompress_sparse_flavor(&self, lg_k: u8) -> Result<UncompressedState, Error> { - debug_assert!(self.window_data.is_empty(), "window is not expected"); + fn write(&mut self, value: u64, count: u8) { + self.buffer |= value << self.buffered_bits; + self.buffered_bits += count; + self.flush_full_word(); + } - let pairs = uncompress_surprising_values( - &self.table_data, - self.table_data_words, - self.table_num_entries, - lg_k, - )?; + fn write_unary(&mut self, value: u64) { + let mut remaining = value; + while remaining >= 16 { + self.pad(16); + remaining -= 16; + } + self.write(1 << remaining, (remaining + 1) as u8); + } - Ok(UncompressedState { - table: PairTable::from_slots(lg_k, self.table_num_entries, pairs)?, - window: vec![], - }) + fn pad(&mut self, count: u8) { + self.buffered_bits += count; + self.flush_full_word(); } - fn uncompress_hybrid_flavor(&self, lg_k: u8) -> Result<UncompressedState, Error> { - debug_assert!(self.window_data.is_empty(), "window is not expected"); - - let mut pairs = uncompress_surprising_values( - &self.table_data, - self.table_data_words, - self.table_num_entries, - lg_k, - )?; - - // In the hybrid flavor, some of these pairs actually belong in the window, so we will - // separate them out, moving the "true" pairs to the bottom of the array. - let k = 1 << lg_k; - let mut window = vec![0u8; k]; // important: zero the memory - let mut next_true_pair = 0; - for i in 0..self.table_num_entries { - let row_col = pairs[i as usize]; - let col = row_col & 63; - if col < 8 { - let row = row_col >> 6; - window[row as usize] |= 1 << col; // set the window bit - } else { - pairs[next_true_pair as usize] = row_col; - next_true_pair += 1; - } + fn flush_full_word(&mut self) { + if self.buffered_bits >= 32 { + self.output.write_u32_le(self.buffer as u32); + self.words_written += 1; + self.buffer >>= 32; + self.buffered_bits -= 32; } - - Ok(UncompressedState { - table: PairTable::from_slots(lg_k, next_true_pair, pairs)?, - window, - }) } - fn uncompress_pinned_flavor( - &self, - lg_k: u8, - num_coupons: u32, - ) -> Result<UncompressedState, Error> { - let mut window = vec![]; - uncompress_sliding_window( - &self.window_data, - self.window_data_words, - &mut window, - lg_k, - num_coupons, - )?; - let num_pairs = self.table_num_entries; - let table = if num_pairs == 0 { - PairTable::new(2, lg_k + 6) - } else { - let mut pairs = uncompress_surprising_values( - &self.table_data, - self.table_data_words, - num_pairs, - lg_k, - )?; - // undo the compressor's 8-column shift - for i in 0..num_pairs { - let i = i as usize; - if (pairs[i] & 63) >= 56 { - return Err(Error::deserial(format!( - "CPC pinned table pair column index is invalid: {}", - pairs[i] - ))); - } - pairs[i] += 8; - } - PairTable::from_slots(lg_k, num_pairs, pairs)? - }; - Ok(UncompressedState { table, window }) + fn finish(mut self) -> usize { + if self.buffered_bits > 0 { + debug_assert!(self.buffered_bits < 32); + self.output.write_u32_le(self.buffer as u32); + self.words_written += 1; + } + self.words_written } +} - fn uncompress_sliding_flavor( - &self, - lg_k: u8, - num_coupons: u32, - ) -> Result<UncompressedState, Error> { - let mut window = vec![]; - uncompress_sliding_window( - &self.window_data, - self.window_data_words, - &mut window, - lg_k, - num_coupons, - )?; - let num_pairs = self.table_num_entries; - let table = if num_pairs == 0 { - PairTable::new(2, lg_k + 6) - } else { - let mut pairs = uncompress_surprising_values( - &self.table_data, - self.table_data_words, - num_pairs, - lg_k, - )?; - let pseudo_phase = determine_pseudo_phase(lg_k, num_coupons); - let permutation = &COLUMN_PERMUTATIONS_FOR_DECODING[pseudo_phase as usize]; - let offset = determine_correct_offset(lg_k, num_coupons); - if offset > 56 { - return Err(Error::deserial(format!( - "CPC sliding window offset is invalid: {offset}" - ))); +/// Decodes the wire payload directly from the caller's byte slice into the two fields owned by a +/// `CpcSketch`. No compressed payload is copied into an intermediate state object. +pub(super) fn decode_payload( + lg_k: u8, + num_coupons: u32, + table_num_entries: u32, + table_data: &[u8], + window_data: &[u8], +) -> Result<(PairTable, Vec<u8>), Error> { + match determine_flavor(lg_k, num_coupons) { + Flavor::Empty => Ok((PairTable::new(2, lg_k + 6), vec![])), + Flavor::Sparse => { + debug_assert!(window_data.is_empty(), "window is not expected"); + let pairs = uncompress_surprising_values(table_data, table_num_entries, lg_k)?; + Ok(( + PairTable::from_slots(lg_k, table_num_entries, pairs)?, + vec![], + )) + } + Flavor::Hybrid => { + debug_assert!(window_data.is_empty(), "window is not expected"); + let mut pairs = uncompress_surprising_values(table_data, table_num_entries, lg_k)?; + let mut window = vec![0u8; 1 << lg_k]; + let mut next_true_pair = 0; + for index in 0..table_num_entries as usize { + let row_col = pairs[index]; + let col = row_col & 63; + if col < 8 { + window[(row_col >> 6) as usize] |= 1 << col; + } else { + pairs[next_true_pair as usize] = row_col; + next_true_pair += 1; + } } - - for i in 0..num_pairs { - let i = i as usize; - let row_col = pairs[i]; - let row = row_col >> 6; - let col = (row_col & 63) as usize; - // first undo the permutation - if col >= permutation.len() { + Ok((PairTable::from_slots(lg_k, next_true_pair, pairs)?, window)) + } + Flavor::Pinned => { + let window = uncompress_sliding_window(window_data, lg_k, num_coupons)?; + let table = if table_num_entries == 0 { + PairTable::new(2, lg_k + 6) + } else { + let mut pairs = uncompress_surprising_values(table_data, table_num_entries, lg_k)?; + for pair in &mut pairs { + if (*pair & 63) >= 56 { + return Err(Error::deserial(format!( + "CPC pinned table pair column index is invalid: {pair}" + ))); + } + *pair += 8; + } + PairTable::from_slots(lg_k, table_num_entries, pairs)? + }; + Ok((table, window)) + } + Flavor::Sliding => { + let window = uncompress_sliding_window(window_data, lg_k, num_coupons)?; + let table = if table_num_entries == 0 { + PairTable::new(2, lg_k + 6) + } else { + let mut pairs = uncompress_surprising_values(table_data, table_num_entries, lg_k)?; + let pseudo_phase = determine_pseudo_phase(lg_k, num_coupons); + let permutation = &COLUMN_PERMUTATIONS_FOR_DECODING[pseudo_phase as usize]; + let offset = determine_correct_offset(lg_k, num_coupons); + if offset > 56 { return Err(Error::deserial(format!( - "CPC sliding table pair column index is invalid: {}", - pairs[i] + "CPC sliding window offset is invalid: {offset}" ))); } - let mut col = permutation[col]; - // then undo the rotation: old = (new + (offset+8)) mod 64 - col = (col + (offset + 8)) & 63; - pairs[i] = (row << 6) | (col as u32); - } - PairTable::from_slots(lg_k, num_pairs, pairs)? - }; - Ok(UncompressedState { table, window }) + + for pair in &mut pairs { + let row = *pair >> 6; + let col = (*pair & 63) as usize; + if col >= permutation.len() { + return Err(Error::deserial(format!( + "CPC sliding table pair column index is invalid: {pair}" + ))); + } + let col = (permutation[col] + offset + 8) & 63; + *pair = (row << 6) | u32::from(col); + } + PairTable::from_slots(lg_k, table_num_entries, pairs)? + }; + Ok((table, window)) + } } } -fn uncompress_surprising_values( - data: &[u32], - data_words: usize, - num_pairs: u32, - lg_k: u8, -) -> Result<Vec<u32>, Error> { +fn uncompress_surprising_values(data: &[u8], num_pairs: u32, lg_k: u8) -> Result<Vec<u32>, Error> { if num_pairs == 0 { return Ok(vec![]); } let k = 1 << lg_k; let mut pairs = vec![0; num_pairs as usize]; let num_base_bits = golomb_choose_number_of_base_bits(k + num_pairs, num_pairs as u64); - low_level_uncompress_pairs(&mut pairs, num_pairs, k, num_base_bits, data, data_words)?; + low_level_uncompress_pairs(&mut pairs, num_pairs, k, num_base_bits, data)?; Ok(pairs) } -fn uncompress_sliding_window( - data: &[u32], - data_words: usize, - window: &mut Vec<u8>, - lg_k: u8, - num_coupons: u32, -) -> Result<(), Error> { +fn uncompress_sliding_window(data: &[u8], lg_k: u8, num_coupons: u32) -> Result<Vec<u8>, Error> { let k = 1 << lg_k; - window.resize(k, 0); + let mut window = vec![0; k]; let pseudo_phase = determine_pseudo_phase(lg_k, num_coupons); low_level_uncompress_bytes( - window, + &mut window, k as u32, data, - data_words, &DECODING_TABLES_FOR_HIGH_ENTROPY_BYTE[pseudo_phase as usize], - ) + )?; + Ok(window) } fn low_level_uncompress_pairs( @@ -552,10 +361,9 @@ fn low_level_uncompress_pairs( num_pairs_to_decode: u32, k: u32, num_base_bits: u8, - compressed_words: &[u32], - num_compressed_words: usize, + compressed_bytes: &[u8], ) -> Result<(), Error> { - let mut bits = BitReader::new(compressed_words, num_compressed_words)?; + let mut bits = BitReader::new(compressed_bytes); let golomb_lo_mask = (1 << num_base_bits) - 1; let mut predicted_row_index = 0u32; let mut predicted_col_index = 0u32; @@ -608,11 +416,10 @@ fn low_level_uncompress_pairs( fn low_level_uncompress_bytes( byte_array: &mut [u8], num_bytes_to_decode: u32, - compressed_words: &[u32], - num_compressed_words: usize, + compressed_bytes: &[u8], decoding_table: &[u16], ) -> Result<(), Error> { - let mut bits = BitReader::new(compressed_words, num_compressed_words)?; + let mut bits = BitReader::new(compressed_bytes); for byte_index in 0..num_bytes_to_decode { // These 12 bits will include an entire Huffman codeword. @@ -627,33 +434,31 @@ fn low_level_uncompress_bytes( } struct BitReader<'a> { - words: &'a [u32], - next_word: usize, + bytes: &'a [u8], + next_byte: usize, buffer: u64, buffered_bits: u8, } impl<'a> BitReader<'a> { - fn new(words: &'a [u32], num_words: usize) -> Result<Self, Error> { - let words = words - .get(..num_words) - .ok_or_else(|| Error::deserial("CPC compressed word count exceeds payload"))?; - Ok(Self { - words, - next_word: 0, + fn new(bytes: &'a [u8]) -> Self { + Self { + bytes, + next_byte: 0, buffer: 0, buffered_bits: 0, - }) + } } fn fill(&mut self, needed: u8) -> Result<(), Error> { if self.buffered_bits < needed { let word = self - .words - .get(self.next_word) + .bytes + .get(self.next_byte..self.next_byte + size_of::<u32>()) .ok_or_else(|| Error::deserial("CPC compressed stream is truncated"))?; - self.buffer |= u64::from(*word) << self.buffered_bits; - self.next_word += 1; + let word = u32::from_le_bytes(word.try_into().unwrap()); + self.buffer |= u64::from(word) << self.buffered_bits; + self.next_byte += size_of::<u32>(); self.buffered_bits += 32; } Ok(()) @@ -735,82 +540,6 @@ fn determine_pseudo_phase(lg_k: u8, num_coupons: u32) -> u8 { } } -fn write_unary( - compressed_words: &mut [u32], - next_word_index: &mut usize, - bitbuf: &mut u64, - bufbits: &mut u8, - value: u64, -) { - assert!(*bufbits <= 31); - - let mut remaining = value; - while remaining >= 16 { - remaining -= 16; - // Here we output 16 zeros, but we don't need to physically write them into bitbuf - // because it already contains zeros in that region. - *bufbits += 16; // Record the fact that 16 bits of output have occurred. - maybe_flush_bitbuf(bitbuf, bufbits, compressed_words, next_word_index); - } - - let the_unary_code = 1 << remaining; - *bitbuf |= the_unary_code << *bufbits; - *bufbits += (remaining + 1) as u8; - maybe_flush_bitbuf(bitbuf, bufbits, compressed_words, next_word_index); -} - -fn maybe_flush_bitbuf( - bitbuf: &mut u64, - bufbits: &mut u8, - word: &mut [u32], - word_index: &mut usize, -) { - if *bufbits >= 32 { - word[*word_index] = (*bitbuf & 0xffffffff) as u32; - *word_index += 1; - *bitbuf >>= 32; - *bufbits -= 32; - } -} - -// Explanation of padding: we write -// 1) xdelta (huffman, provides at least 1 bit, requires 12-bit lookahead) -// 2) ydeltaGolombHi (unary, provides at least 1 bit, requires 8-bit lookahead) -// 3) ydeltaGolombLo (straight B bits). -// So the 12-bit lookahead is the tight constraint, but there are at least (2 + B) bits emitted, -// so we would be safe with max (0, 10 - B) bits of padding at the end of the bitstream. -fn safe_length_for_compressed_window_buf(k: u32) -> usize { - // 11 bits of padding, due to 12-bit lookahead, with 1 bit certainly present. - let bits = 12 * k + 11; - divide_longs_rounding_up(bits as usize, 32) -} - -fn safe_length_for_compressed_pair_buf(k: u32, num_pairs: u32, num_base_bits: u8) -> usize { - // Long ybits = k + numPairs; // simpler and safer UB - // The following tighter UB on ybits is based on page 198 - // of the textbook "Managing Gigabytes" by Witten, Moffat, and Bell. - // Notice that if numBaseBits == 0 it coincides with (k + numPairs). - - let k = k as usize; - let num_pairs = num_pairs as usize; - let num_base_bits = num_base_bits as usize; - - let ybits = num_pairs * (1 + num_base_bits) + (k >> num_base_bits); - let xbits = 12 * (num_pairs); - let padding = 10usize.saturating_sub(num_base_bits); - divide_longs_rounding_up(xbits + ybits + padding, 32) -} - -fn divide_longs_rounding_up(x: usize, y: usize) -> usize { - debug_assert_ne!(y, 0); - let quotient = x / y; - if quotient * y == x { - quotient - } else { - quotient + 1 - } -} - /// Returns an integer that is between zero and ceil(log_2(k)) - 1, inclusive. fn golomb_choose_number_of_base_bits(k: u32, count: u64) -> u8 { debug_assert!(k > 0); @@ -841,9 +570,10 @@ fn floor_log2_of_long(x: u64) -> u8 { #[cfg(test)] mod tests { - use super::CompressedState; + use super::compress_surprising_values; use super::determine_pseudo_phase; use super::uncompress_surprising_values; + use crate::codec::SketchBytes; #[test] fn pseudo_phase_handles_maximum_lg_k() { @@ -853,15 +583,11 @@ mod tests { #[test] fn pair_decoder_rejects_empty_table_sentinel() { - let mut compressed = CompressedState::default(); - compressed.compress_surprising_values(&[u32::MAX], 26); - - let error = uncompress_surprising_values( - &compressed.table_data, - compressed.table_data_words, - 1, - 26, - ); + let mut compressed = SketchBytes::with_capacity(16); + compress_surprising_values(&[u32::MAX], 26, &mut compressed); + let compressed = compressed.into_bytes(); + + let error = uncompress_surprising_values(&compressed, 1, 26); assert_eq!( error.unwrap_err().message(), "CPC pair uses the reserved empty-table sentinel" diff --git a/datasketches/src/cpc/pair_table.rs b/datasketches/src/cpc/pair_table.rs index 8c0fb62..17811d4 100644 --- a/datasketches/src/cpc/pair_table.rs +++ b/datasketches/src/cpc/pair_table.rs @@ -86,6 +86,10 @@ impl PairTable { &self.slots } + pub fn len(&self) -> u32 { + self.num_items + } + pub fn clear(&mut self) { self.slots.fill(u32::MAX); self.num_items = 0; diff --git a/datasketches/src/cpc/sketch.rs b/datasketches/src/cpc/sketch.rs index 2ae9164..ce81620 100644 --- a/datasketches/src/cpc/sketch.rs +++ b/datasketches/src/cpc/sketch.rs @@ -29,7 +29,8 @@ use crate::cpc::DEFAULT_LG_K; use crate::cpc::Flavor; use crate::cpc::MAX_LG_K; use crate::cpc::MIN_LG_K; -use crate::cpc::compression::CompressedState; +use crate::cpc::compression::decode_payload; +use crate::cpc::compression::encode_payload; use crate::cpc::count_bits_set_in_matrix; use crate::cpc::determine_correct_offset; use crate::cpc::determine_flavor; @@ -470,12 +471,15 @@ impl CpcSketch { pub fn serialize(&self) -> Vec<u8> { let mut bytes = SketchBytes::with_capacity(256); - let mut compressed = CompressedState::default(); - compressed.compress(self); - + let flavor = self.flavor(); let has_hip = !self.merge_flag; - let has_table = !compressed.table_data.is_empty(); - let has_window = !compressed.window_data.is_empty(); + let has_window = matches!(flavor, Flavor::Pinned | Flavor::Sliding); + let table_num_entries = match flavor { + Flavor::Empty => 0, + Flavor::Sparse | Flavor::Hybrid => self.num_coupons, + Flavor::Pinned | Flavor::Sliding => self.surprising_value_table().len(), + }; + let has_table = table_num_entries > 0; let preamble_ints = make_preamble_ints(self.num_coupons, has_hip, has_table, has_window); bytes.write_u8(preamble_ints); bytes.write_u8(SERIAL_VERSION); @@ -493,34 +497,37 @@ impl CpcSketch { bytes.write_u32_le(self.num_coupons); if has_table && has_window { // if there is no window it is the same as number of coupons - bytes.write_u32_le(compressed.table_num_entries); + bytes.write_u32_le(table_num_entries); // HIP values can be in two different places in the sequence of fields // this is the first HIP decision point if has_hip { self.write_hip(&mut bytes); } } - if has_table { - debug_assert!(compressed.table_data_words <= u32::MAX as usize); - bytes.write_u32_le(compressed.table_data_words as u32); - } - if has_window { - debug_assert!(compressed.window_data_words <= u32::MAX as usize); - bytes.write_u32_le(compressed.window_data_words as u32); - } + let table_words_offset = has_table.then(|| { + let offset = bytes.len(); + bytes.write_u32_le(0); + offset + }); + let window_words_offset = has_window.then(|| { + let offset = bytes.len(); + bytes.write_u32_le(0); + offset + }); // this is the second HIP decision point if has_hip && !(has_table && has_window) { self.write_hip(&mut bytes); } - if has_window { - for i in 0..compressed.window_data_words { - bytes.write_u32_le(compressed.window_data[i]); - } + let (table_words, window_words) = encode_payload(self, &mut bytes); + debug_assert_eq!(has_table, table_words > 0); + debug_assert_eq!(has_window, window_words > 0); + debug_assert!(table_words <= u32::MAX as usize); + debug_assert!(window_words <= u32::MAX as usize); + if let Some(offset) = table_words_offset { + bytes.overwrite_u32_le(offset, table_words as u32); } - if has_table { - for i in 0..compressed.table_data_words { - bytes.write_u32_le(compressed.table_data[i]); - } + if let Some(offset) = window_words_offset { + bytes.overwrite_u32_le(offset, window_words as u32); } } bytes.into_bytes() @@ -561,8 +568,10 @@ impl CpcSketch { let has_table = flags & (1 << FLAG_HAS_TABLE) != 0; let has_window = flags & (1 << FLAG_HAS_WINDOW) != 0; - let mut compressed = CompressedState::default(); let mut num_coupons = 0; + let mut table_num_entries = 0; + let mut table_data_words = 0; + let mut window_data_words = 0; let mut kxp = 0.0; let mut hip_est_accum = 0.0; @@ -571,7 +580,7 @@ impl CpcSketch { .read_u32_le() .map_err(insufficient_data("num_coupons"))?; if has_table && has_window { - compressed.table_num_entries = cursor + table_num_entries = cursor .read_u32_le() .map_err(insufficient_data("table_num_entries"))?; if has_hip { @@ -582,13 +591,13 @@ impl CpcSketch { } } if has_table { - compressed.table_data_words = cursor + table_data_words = cursor .read_u32_le() .map_err(insufficient_data("table_data_words"))? as usize; } if has_window { - compressed.window_data_words = cursor + window_data_words = cursor .read_u32_le() .map_err(insufficient_data("window_data_words"))? as usize; @@ -599,24 +608,8 @@ impl CpcSketch { .read_f64_le() .map_err(insufficient_data("hip_est_accum"))?; } - if has_window { - for _ in 0..compressed.window_data_words { - let word = cursor - .read_u32_le() - .map_err(insufficient_data("window_data"))?; - compressed.window_data.push(word); - } - } - if has_table { - for _ in 0..compressed.table_data_words { - let word = cursor - .read_u32_le() - .map_err(insufficient_data("table_data"))?; - compressed.table_data.push(word); - } - } if !has_window { - compressed.table_num_entries = num_coupons; + table_num_entries = num_coupons; } } @@ -668,41 +661,60 @@ impl CpcSketch { } // The number of stored table entries can never exceed the number of coupons. - if compressed.table_num_entries > num_coupons { + if table_num_entries > num_coupons { return Err(Error::deserial(format!( "table_num_entries ({}) exceeds num_coupons ({})", - compressed.table_num_entries, num_coupons + table_num_entries, num_coupons ))); } // A pair requires at least one bit to encode, so the declared number of table entries can // never exceed the number of bits available in the table data. This also bounds the size // of the allocation made while decoding, rejecting corrupt inputs that claim an enormous // entry count backed by only a few data words. - if (compressed.table_num_entries as usize) > compressed.table_data_words.saturating_mul(32) - { + if (table_num_entries as usize) > table_data_words.saturating_mul(32) { return Err(Error::deserial(format!( "table_num_entries ({}) exceeds capacity of table data ({} words)", - compressed.table_num_entries, compressed.table_data_words + table_num_entries, table_data_words ))); } let k = 1usize << lg_k; - if has_window && compressed.window_data_words.saturating_mul(32) < k { + if has_window && window_data_words.saturating_mul(32) < k { return Err(Error::deserial(format!( "window data ({} words) is too short for lg_k = {lg_k}", - compressed.window_data_words + window_data_words ))); } - let uncompressed = compressed.uncompress(lg_k, num_coupons)?; + let window_data_bytes = window_data_words.checked_mul(4).ok_or_else(|| { + Error::deserial("CPC window data word count overflows payload length") + })?; + let table_data_bytes = table_data_words + .checked_mul(4) + .ok_or_else(|| Error::deserial("CPC table data word count overflows payload length"))?; + let payload_bytes = window_data_bytes + .checked_add(table_data_bytes) + .ok_or_else(|| Error::deserial("CPC payload length overflows"))?; + let payload = cursor + .remaining() + .get(..payload_bytes) + .ok_or_else(|| Error::deserial("insufficient data for CPC compressed payload"))?; + let (window_data, table_data) = payload.split_at(window_data_bytes); + let (table, window) = decode_payload( + lg_k, + num_coupons, + table_num_entries, + table_data, + window_data, + )?; Ok(CpcSketch { lg_k, seed, seed_hash, first_interesting_column, num_coupons, - surprising_value_table: Some(uncompressed.table), + surprising_value_table: Some(table), window_offset: determine_correct_offset(lg_k, num_coupons), - sliding_window: uncompressed.window, + sliding_window: window, merge_flag: !has_hip, kxp, hip_est_accum, --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
