diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f68d7e3..fe7a6c4d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ All significant changes to this project will be documented in this file. * Use serial version 2 for new `CountMinSketch` instances, enabling cross-language interchange with the documented input encodings. New sketches support `1..=127` rows and require a version 2 reader. Existing version 1 data remains readable and updatable, but must be rebuilt from source items to merge with new sketches. * Add `CpcUnion::serialize`, which produces the same bytes as `to_sketch().serialize()` and can be faster. * Add `HllUnion::serialize` to serialize union results directly. +* Add `ThetaUnion::serialize` and `TupleUnion::serialize` to serialize ordered or unordered union results directly in their compact binary formats. Tuple summaries are serialized without cloning, avoiding per-entry allocations for heap-backed summaries. * Add `XorFilter` and `XorFilterBuilder` behind the `xor` feature for immutable membership queries on values or precomputed hashes. Filters support 8-, 16-, and 32-bit fingerprints and can be serialized for later queries. ### Improvements diff --git a/benchmarks/theta/mod.rs b/benchmarks/theta/mod.rs index 4839176b..32c8b7f2 100644 --- a/benchmarks/theta/mod.rs +++ b/benchmarks/theta/mod.rs @@ -15,8 +15,10 @@ // specific language governing permissions and limitations // under the License. +mod serde; mod update; pub fn register(runner: &mut benchmarks::Runner) { + serde::register(runner); update::register(runner); } diff --git a/benchmarks/theta/serde.rs b/benchmarks/theta/serde.rs new file mode 100644 index 00000000..89d2e59c --- /dev/null +++ b/benchmarks/theta/serde.rs @@ -0,0 +1,64 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; + +use benchmarks::Bencher; +use benchmarks::Runner; +use datasketches::theta::ThetaSketchBuilder; +use datasketches::theta::ThetaUnion; +use datasketches::theta::ThetaUnionBuilder; + +fn serialize_union(bencher: Bencher, rows: u64, ordered: bool, direct: bool) { + let union = union(rows); + if direct { + bencher.bench_local(|| black_box(&union).serialize(ordered)); + } else { + bencher.bench_local(|| black_box(&union).to_sketch(ordered).serialize()); + } +} + +fn union(rows: u64) -> ThetaUnion { + let mut sketch = ThetaSketchBuilder::default().lg_k(14).build().unwrap(); + for value in 0..rows { + sketch.update(value); + } + let mut union = ThetaUnionBuilder::default().lg_k(12).build().unwrap(); + union.update(&sketch).unwrap(); + union +} + +pub fn register(runner: &mut Runner) { + // Reuse each union for repeated serialization, including fixture construction in timing. + for (rows, rounds) in [(40, 1_000_000), (8_000, 10_000)] { + for ordered in [false, true] { + let order = if ordered { "ordered" } else { "unordered" }; + runner.compare( + &format!("theta::serde::{order}/{rows}"), + rounds, + &[ + ("via_compact", &|b| { + serialize_union(b, rows, ordered, false); + }), + ("direct", &|b| { + serialize_union(b, rows, ordered, true); + }), + ], + ); + } + } +} diff --git a/benchmarks/tuple/mod.rs b/benchmarks/tuple/mod.rs index 4839176b..32c8b7f2 100644 --- a/benchmarks/tuple/mod.rs +++ b/benchmarks/tuple/mod.rs @@ -15,8 +15,10 @@ // specific language governing permissions and limitations // under the License. +mod serde; mod update; pub fn register(runner: &mut benchmarks::Runner) { + serde::register(runner); update::register(runner); } diff --git a/benchmarks/tuple/serde.rs b/benchmarks/tuple/serde.rs new file mode 100644 index 00000000..700675bf --- /dev/null +++ b/benchmarks/tuple/serde.rs @@ -0,0 +1,188 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; + +use benchmarks::Bencher; +use benchmarks::Runner; +use datasketches::codec::SketchBytes; +use datasketches::codec::SketchSlice; +use datasketches::error::Error; +use datasketches::tuple::DefaultUnionPolicy; +use datasketches::tuple::DefaultUpdatePolicy; +use datasketches::tuple::SummaryCombinePolicy; +use datasketches::tuple::SummaryPolicy; +use datasketches::tuple::SummaryUpdatePolicy; +use datasketches::tuple::TupleSketchBuilder; +use datasketches::tuple::TupleSummaryValue; +use datasketches::tuple::TupleUnion; +use datasketches::tuple::TupleUnionBuilder; + +fn u64_union(rows: u64, lg_k: u8) -> TupleUnion> { + let mut sketch = TupleSketchBuilder::new(DefaultUpdatePolicy::::default()) + .lg_k(lg_k) + .build() + .unwrap(); + for value in 0..rows { + sketch.update(value, 1); + } + let mut union = TupleUnionBuilder::new(DefaultUnionPolicy::::default()) + .lg_k(lg_k) + .build() + .unwrap(); + union.update(&sketch).unwrap(); + union +} + +fn serialize_union(bencher: Bencher, rows: u64, lg_k: u8) { + let union = u64_union(rows, lg_k); + bencher.bench_local(|| black_box(&union).serialize(ORDERED)); +} + +fn serialize_union_via_sketch(bencher: Bencher, rows: u64, lg_k: u8) { + let union = u64_union(rows, lg_k); + bencher.bench_local(|| black_box(&union).to_sketch(ORDERED).serialize()); +} + +/// Heap-backed summary holding a few short labels per key. +#[derive(Clone, Default)] +struct Labels(Vec); + +impl TupleSummaryValue for Labels { + fn serialize_size(&self) -> usize { + 4 + self.0.iter().map(|label| 4 + label.len()).sum::() + } + + fn serialize_value(&self, bytes: &mut SketchBytes) { + bytes.write_u32_le(self.0.len() as u32); + for label in &self.0 { + bytes.write_u32_le(label.len() as u32); + bytes.write(label.as_bytes()); + } + } + + fn deserialize_value(_cursor: &mut SketchSlice<'_>) -> Result { + unimplemented!("only serialization is benchmarked") + } +} + +#[derive(Clone, Copy)] +struct LabelsPolicy; + +impl SummaryPolicy for LabelsPolicy { + type Summary = Labels; + + fn create(&self) -> Self::Summary { + Labels::default() + } +} + +impl SummaryUpdatePolicy for LabelsPolicy { + fn update(&self, summary: &mut Self::Summary, value: u64) { + for i in 0..4 { + summary.0.push(format!("label-{value:08}-{i}")); + } + } +} + +impl SummaryCombinePolicy for LabelsPolicy { + fn combine(&self, summary: &mut Self::Summary, other: &Self::Summary) { + summary.0.extend(other.0.iter().cloned()); + } +} + +fn labels_union(rows: u64, lg_k: u8) -> TupleUnion { + let mut sketch = TupleSketchBuilder::new(LabelsPolicy) + .lg_k(lg_k) + .build() + .unwrap(); + for value in 0..rows { + sketch.update(value, value); + } + let mut union = TupleUnionBuilder::new(LabelsPolicy) + .lg_k(lg_k) + .build() + .unwrap(); + union.update(&sketch).unwrap(); + union +} + +fn serialize_labels_union(bencher: Bencher, rows: u64, lg_k: u8) { + let union = labels_union(rows, lg_k); + bencher.bench_local(|| black_box(&union).serialize(ORDERED)); +} + +fn serialize_labels_union_via_sketch(bencher: Bencher, rows: u64, lg_k: u8) { + let union = labels_union(rows, lg_k); + bencher.bench_local(|| black_box(&union).to_sketch(ORDERED).serialize()); +} + +pub fn register(runner: &mut Runner) { + // Reuse each union for repeated serialization, including fixture construction in timing. + // Cover exact results and estimation results with different retained sizes. + for (rows, lg_k, u64_rounds, labels_rounds) in [ + (0, 12, 10_000_000, 10_000_000), + (4, 12, 2_000_000, 500_000), + (100, 12, 300_000, 20_000), + (1_000, 12, 30_000, 2_000), + (80_000, 8, 100_000, 6_000), + (80_000, 12, 5_000, 500), + (80_000, 16, 250, 20), + ] { + register_order::(runner, rows, lg_k, u64_rounds, labels_rounds); + register_order::(runner, rows, lg_k, u64_rounds, labels_rounds); + } +} + +fn register_order( + runner: &mut Runner, + rows: u64, + lg_k: u8, + u64_rounds: usize, + labels_rounds: usize, +) { + let size = if lg_k == 12 { + rows.to_string() + } else { + format!("{rows}/lg_k={lg_k}") + }; + let order = if ORDERED { "ordered" } else { "unordered" }; + runner.compare( + &format!("tuple::serde::u64::{order}/{size}"), + u64_rounds, + &[ + ("via_compact_sketch", &|bencher| { + serialize_union_via_sketch::(bencher, rows, lg_k) + }), + ("direct", &|bencher| { + serialize_union::(bencher, rows, lg_k) + }), + ], + ); + runner.compare( + &format!("tuple::serde::labels::{order}/{size}"), + labels_rounds, + &[ + ("via_compact_sketch", &|bencher| { + serialize_labels_union_via_sketch::(bencher, rows, lg_k) + }), + ("direct", &|bencher| { + serialize_labels_union::(bencher, rows, lg_k) + }), + ], + ); +} diff --git a/datasketches/src/thetafamily/common/mod.rs b/datasketches/src/thetafamily/common/mod.rs index dbc273ee..e1a9ae89 100644 --- a/datasketches/src/thetafamily/common/mod.rs +++ b/datasketches/src/thetafamily/common/mod.rs @@ -33,6 +33,18 @@ pub(super) trait SketchEntry { fn hash(&self) -> u64; } +impl SketchEntry for u64 { + fn hash(&self) -> u64 { + *self + } +} + +impl SketchEntry for &E { + fn hash(&self) -> u64 { + E::hash(self) + } +} + pub(super) trait KeySketch: Copy { fn metadata(self) -> sketch_state::ThetaFamilySketchMetadata; diff --git a/datasketches/src/thetafamily/common/sketch_state.rs b/datasketches/src/thetafamily/common/sketch_state.rs index d59f74f6..44928c0b 100644 --- a/datasketches/src/thetafamily/common/sketch_state.rs +++ b/datasketches/src/thetafamily/common/sketch_state.rs @@ -82,6 +82,19 @@ impl CompactSketchState { matches!(self, Self::NonEmpty { theta, .. } if *theta < MAX_THETA) } + /// Returns the preamble size in 8-byte words for uncompressed compact Theta and Tuple images. + pub fn uncompressed_preamble_longs(&self) -> u8 { + if self.is_empty() { + 1 + } else if self.is_estimation_mode() { + 3 + } else if self.retained_entries().len() == 1 { + 1 + } else { + 2 + } + } + pub fn seed_hash(&self) -> u16 { match self { Self::Empty { seed_hash } | Self::NonEmpty { seed_hash, .. } => *seed_hash, diff --git a/datasketches/src/thetafamily/common/union.rs b/datasketches/src/thetafamily/common/union.rs index 4ad3b466..b9e5d1db 100644 --- a/datasketches/src/thetafamily/common/union.rs +++ b/datasketches/src/thetafamily/common/union.rs @@ -113,21 +113,37 @@ where where E: Clone, { - let Some(result_theta) = self.result_theta else { + let Some(theta) = self.result_theta else { return CompactSketchState::empty(self.table.seed_hash()); }; + let entries = self.retained_entries(theta).cloned().collect(); + self.finish_compaction(entries, theta, ordered) + } - let mut theta = result_theta.min(self.table.retention_theta()); - let mut retained_entries = if result_theta >= self.table.retention_theta() { - self.table.iter_entries().cloned().collect::>() - } else { - self.table - .iter_entries() - .filter(|entry| entry.hash() < theta) - .cloned() - .collect::>() + /// Returns the union as canonical compact-sketch state that borrows the retained entries. + #[cfg(feature = "tuple")] + pub fn to_borrowed_compact_sketch_state(&self, ordered: bool) -> CompactSketchState<&E> { + let Some(theta) = self.result_theta else { + return CompactSketchState::empty(self.table.seed_hash()); }; + let entries = self.retained_entries(theta).collect(); + self.finish_compaction(entries, theta, ordered) + } + fn retained_entries(&self, theta: u64) -> impl Iterator { + let all_retained = theta >= self.table.retention_theta(); + self.table + .iter_entries() + .filter(move |entry| all_retained || entry.hash() < theta) + } + + fn finish_compaction( + &self, + mut retained_entries: Vec, + result_theta: u64, + ordered: bool, + ) -> CompactSketchState { + let mut theta = result_theta.min(self.table.retention_theta()); let nominal_num = 1usize << self.table.lg_nom_size(); if retained_entries.len() > nominal_num { let (_, kth, _) = diff --git a/datasketches/src/thetafamily/theta/serialization.rs b/datasketches/src/thetafamily/theta/serialization.rs index d10917f0..b25edcf6 100644 --- a/datasketches/src/thetafamily/theta/serialization.rs +++ b/datasketches/src/thetafamily/theta/serialization.rs @@ -15,7 +15,15 @@ // specific language governing permissions and limitations // under the License. -//! Binary serialization format constants for Theta sketches. +//! Binary serialization for compact Theta sketches. + +use crate::codec::SketchBytes; +use crate::codec::family::Family; +use crate::thetacommon::SketchEntry; +use crate::thetacommon::constants::FLAGS_IS_COMPACT; +use crate::thetacommon::constants::FLAGS_IS_ORDERED; +use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; +use crate::thetacommon::sketch_state::CompactSketchState; pub const UNCOMPRESSED_SERIAL_VERSION: u8 = 3; pub const COMPRESSED_SERIAL_VERSION: u8 = 4; @@ -26,3 +34,42 @@ pub const EMPTY_SKETCH_BYTES: [u8; 8] = [1, 3, 3, 0, 0, 0x1e, 0, 0]; pub const V2_PREAMBLE_EMPTY: u8 = 1; pub const V2_PREAMBLE_PRECISE: u8 = 2; pub const V2_PREAMBLE_ESTIMATE: u8 = 3; + +pub(super) fn serialize_compact_uncompressed( + state: &CompactSketchState, +) -> Vec { + if state.is_empty() { + return EMPTY_SKETCH_BYTES.to_vec(); + } + + let entries = state.retained_entries(); + let pre_longs = state.uncompressed_preamble_longs(); + let mut bytes = SketchBytes::with_capacity((usize::from(pre_longs) + entries.len()) * 8); + bytes.write_u8(pre_longs); + bytes.write_u8(UNCOMPRESSED_SERIAL_VERSION); + bytes.write_u8(Family::THETA.id); + bytes.write_u16_be(0); // unused for compact + + let mut flags = 0u8; + flags |= FLAGS_IS_READ_ONLY; + flags |= FLAGS_IS_COMPACT; + if state.is_ordered() { + flags |= FLAGS_IS_ORDERED; + } + bytes.write_u8(flags); + + bytes.write_u16_le(state.seed_hash()); + + if pre_longs > 1 { + bytes.write_u32_le(entries.len() as u32); + // Unused by compact sketches; matches datasketches-java and datasketches-cpp. + bytes.write_u32_be(0); + } + if state.is_estimation_mode() { + bytes.write_u64_le(state.theta()); + } + for entry in entries { + bytes.write_u64_le(entry.hash()); + } + bytes.into_bytes() +} diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index e7d0b9b6..71105cf1 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -531,18 +531,6 @@ impl CompactThetaSketch { .expect("compact theta should always be valid") } - fn preamble_longs(&self, compressed: bool) -> u8 { - if compressed { - if self.is_estimation_mode() { 2 } else { 1 } - } else if self.is_estimation_mode() { - 3 - } else if self.num_retained() == 1 { - 1 - } else { - 2 - } - } - /// Serializes this sketch in compressed form if applicable. /// /// The serialized data can be read by other DataSketches implementations. @@ -568,46 +556,12 @@ impl CompactThetaSketch { /// /// Empty sketches serialize with a zero seed hash. pub fn serialize(&self) -> Vec { - if self.is_empty() { - return serialization::EMPTY_SKETCH_BYTES.to_vec(); - } - - let retained_hashes = self.retained_hashes(); - let mut bytes = SketchBytes::with_capacity(64 + retained_hashes.len() * 8); - - let pre_longs = self.preamble_longs(false); - bytes.write_u8(pre_longs); - bytes.write_u8(serialization::UNCOMPRESSED_SERIAL_VERSION); - bytes.write_u8(Family::THETA.id); - bytes.write_u16_be(0); // unused for compact - - let mut flags = 0u8; - flags |= FLAGS_IS_READ_ONLY; - flags |= FLAGS_IS_COMPACT; - if self.is_ordered() { - flags |= FLAGS_IS_ORDERED; - } - bytes.write_u8(flags); - - bytes.write_u16_le(self.seed_hash()); - - if pre_longs > 1 { - bytes.write_u32_le(retained_hashes.len() as u32); - // Unused by compact sketches; matches datasketches-java and datasketches-cpp. - bytes.write_u32_be(0); - } - if self.is_estimation_mode() { - bytes.write_u64_le(self.theta64()); - } - for hash in retained_hashes { - bytes.write_u64_le(*hash); - } - bytes.into_bytes() + serialization::serialize_compact_uncompressed(&self.compact_state) } fn serialize_v4(&self) -> Vec { let retained_hashes = self.retained_hashes(); - let pre_longs = self.preamble_longs(true); + let pre_longs = if self.is_estimation_mode() { 2u8 } else { 1u8 }; let entry_bits = Self::compute_entry_bits(retained_hashes); let num_entries_bytes = Self::num_entries_bytes(retained_hashes.len()); diff --git a/datasketches/src/thetafamily/theta/union.rs b/datasketches/src/thetafamily/theta/union.rs index 1b26d7a2..a8f21faf 100644 --- a/datasketches/src/thetafamily/theta/union.rs +++ b/datasketches/src/thetafamily/theta/union.rs @@ -21,6 +21,7 @@ use crate::hash::DEFAULT_UPDATE_SEED; use crate::theta::CompactThetaSketch; use crate::theta::ThetaSketchView; use crate::theta::hash_table::ThetaEntry; +use crate::theta::serialization; use crate::thetacommon::constants::DEFAULT_LG_K; use crate::thetacommon::union::UnionMergePolicy; use crate::thetacommon::union::UnionState; @@ -59,6 +60,16 @@ impl ThetaUnion { CompactThetaSketch::from_compact_state(compact_state) } + /// Serializes this union result into the uncompressed compact Theta format. + /// + /// Produces the same bytes as `self.to_sketch(ordered).serialize()`. When `ordered` is `true`, + /// retained hashes are sorted. Otherwise they may be unordered; empty results and exact + /// singletons always use their canonical encoding. + pub fn serialize(&self, ordered: bool) -> Vec { + let state = self.state.to_compact_sketch_state(ordered); + serialization::serialize_compact_uncompressed(&state) + } + /// Resets the union to its empty state. pub fn reset(&mut self) { self.state.reset(); diff --git a/datasketches/src/thetafamily/tuple/serialization.rs b/datasketches/src/thetafamily/tuple/serialization.rs index cc208e0a..1e7b9164 100644 --- a/datasketches/src/thetafamily/tuple/serialization.rs +++ b/datasketches/src/thetafamily/tuple/serialization.rs @@ -25,10 +25,18 @@ //! type is opaque to the sketch, (de)serialization of summaries is delegated to the summary type //! itself via the [`TupleSummaryValue`] trait, mirroring `FrequentItemValue`. +use std::borrow::Borrow; + use crate::codec::SketchBytes; use crate::codec::SketchSlice; use crate::codec::assert::insufficient_data; +use crate::codec::family::Family; use crate::error::Error; +use crate::thetacommon::constants::FLAGS_IS_COMPACT; +use crate::thetacommon::constants::FLAGS_IS_ORDERED; +use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; +use crate::thetacommon::sketch_state::CompactSketchState; +use crate::tuple::hash_table::TupleEntry; /// Current serial version written by this implementation. pub const SERIAL_VERSION: u8 = 3; @@ -43,6 +51,54 @@ pub const SKETCH_TYPE_LEGACY: u8 = 5; /// Canonical empty compact image, independent of the configured seed and summary type. pub const EMPTY_SKETCH_BYTES: [u8; 8] = [1, 3, 9, 1, 0, 0x1e, 0, 0]; +/// Serializes compact Tuple state whose retained entries are owned or borrowed. +pub(super) fn serialize_compact_uncompressed(compact_state: &CompactSketchState) -> Vec +where + S: TupleSummaryValue, + E: Borrow>, +{ + if compact_state.is_empty() { + return EMPTY_SKETCH_BYTES.to_vec(); + } + + let retained_entries = compact_state.retained_entries(); + let is_estimation_mode = compact_state.is_estimation_mode(); + let pre_longs = compact_state.uncompressed_preamble_longs(); + let entries_size: usize = retained_entries + .iter() + .map(|entry| 8 + entry.borrow().summary().serialize_size()) + .sum(); + let mut bytes = SketchBytes::with_capacity(8 * pre_longs as usize + entries_size); + + bytes.write_u8(pre_longs); + bytes.write_u8(SERIAL_VERSION); + bytes.write_u8(Family::TUPLE.id); + bytes.write_u8(SKETCH_TYPE); + bytes.write_u8(0); // unused + + let mut flags = FLAGS_IS_READ_ONLY | FLAGS_IS_COMPACT; + if compact_state.is_ordered() { + flags |= FLAGS_IS_ORDERED; + } + bytes.write_u8(flags); + bytes.write_u16_le(compact_state.seed_hash()); + + if pre_longs > 1 { + bytes.write_u32_le(retained_entries.len() as u32); + bytes.write_u32_le(0); // unused + } + if is_estimation_mode { + bytes.write_u64_le(compact_state.theta()); + } + + for entry in retained_entries { + let entry = entry.borrow(); + bytes.write_u64_le(entry.hash()); + entry.summary().serialize_value(&mut bytes); + } + bytes.into_bytes() +} + /// Trait for values that can be stored as Tuple sketch summaries. /// /// Implement this trait for a summary type to make it (de)serializable by diff --git a/datasketches/src/thetafamily/tuple/sketch.rs b/datasketches/src/thetafamily/tuple/sketch.rs index 37bd946d..0cd368b8 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -25,7 +25,6 @@ use std::hash::Hash; use std::slice; -use crate::codec::SketchBytes; use crate::codec::SketchSlice; use crate::codec::assert::ensure_preamble_longs_in_range; use crate::codec::assert::insufficient_data; @@ -41,10 +40,8 @@ use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::DEFAULT_LG_K; -use crate::thetacommon::constants::FLAGS_IS_COMPACT; use crate::thetacommon::constants::FLAGS_IS_EMPTY; use crate::thetacommon::constants::FLAGS_IS_ORDERED; -use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTableIter; use crate::thetacommon::sketch_state::CompactSketchState; @@ -53,12 +50,12 @@ use crate::tuple::hash_table::TupleEntry; use crate::tuple::hash_table::TupleHashTable; use crate::tuple::policy::SummaryPolicy; use crate::tuple::policy::SummaryUpdatePolicy; -use crate::tuple::serialization::EMPTY_SKETCH_BYTES; use crate::tuple::serialization::SERIAL_VERSION; use crate::tuple::serialization::SERIAL_VERSION_LEGACY; use crate::tuple::serialization::SKETCH_TYPE; use crate::tuple::serialization::SKETCH_TYPE_LEGACY; use crate::tuple::serialization::TupleSummaryValue; +use crate::tuple::serialization::serialize_compact_uncompressed; /// Read-only view of a mutable or compact Tuple sketch. /// @@ -548,16 +545,6 @@ impl CompactTupleSketch { self.compact_state.retained_entries_capacity() * size_of::>() } - fn preamble_longs(&self) -> u8 { - if self.is_estimation_mode() { - 3 - } else if self.num_retained() == 1 { - 1 - } else { - 2 - } - } - /// Serializes this sketch into the compact Tuple binary format. /// /// The serialized data can be read by other DataSketches implementations when both @@ -581,44 +568,7 @@ impl CompactTupleSketch { where S: TupleSummaryValue, { - if self.is_empty() { - return EMPTY_SKETCH_BYTES.to_vec(); - } - - let retained_entries = self.retained_entries(); - let pre_longs = self.preamble_longs(); - let entries_size: usize = retained_entries - .iter() - .map(|entry| 8 + entry.summary().serialize_size()) - .sum(); - let mut bytes = SketchBytes::with_capacity(8 * pre_longs as usize + entries_size); - - bytes.write_u8(pre_longs); - bytes.write_u8(SERIAL_VERSION); - bytes.write_u8(Family::TUPLE.id); - bytes.write_u8(SKETCH_TYPE); - bytes.write_u8(0); // unused - - let mut flags = FLAGS_IS_READ_ONLY | FLAGS_IS_COMPACT; - if self.is_ordered() { - flags |= FLAGS_IS_ORDERED; - } - bytes.write_u8(flags); - bytes.write_u16_le(self.seed_hash()); - - if pre_longs > 1 { - bytes.write_u32_le(retained_entries.len() as u32); - bytes.write_u32_le(0); // unused - } - if self.is_estimation_mode() { - bytes.write_u64_le(self.theta64()); - } - - for entry in retained_entries { - bytes.write_u64_le(entry.hash()); - entry.summary().serialize_value(&mut bytes); - } - bytes.into_bytes() + serialize_compact_uncompressed(&self.compact_state) } /// Deserializes a compact Tuple sketch using the default seed. diff --git a/datasketches/src/thetafamily/tuple/union.rs b/datasketches/src/thetafamily/tuple/union.rs index 029aeba7..f2d419a2 100644 --- a/datasketches/src/thetafamily/tuple/union.rs +++ b/datasketches/src/thetafamily/tuple/union.rs @@ -29,6 +29,8 @@ use crate::thetacommon::constants::DEFAULT_LG_K; use crate::thetacommon::union::UnionState; use crate::tuple::hash_table::TupleEntry; use crate::tuple::policy::SummaryCombinePolicy; +use crate::tuple::serialization::TupleSummaryValue; +use crate::tuple::serialization::serialize_compact_uncompressed; use crate::tuple::sketch::CompactTupleSketch; use crate::tuple::sketch::TupleSketchView; @@ -107,6 +109,44 @@ where CompactTupleSketch::from_compact_state(self.state.to_compact_sketch_state(ordered)) } + /// Serializes the union result into the compact Tuple binary format. + /// + /// The result is the same as `self.to_sketch(ordered).serialize()`: the same retained entries, + /// theta, and seed hash, and identical bytes when `ordered` is `true`. Unordered results may + /// list the retained entries in a different order. Summaries are encoded in + /// place instead of being cloned into an intermediate sketch, so `P::Summary` does not need to + /// implement [`Clone`]. + /// + /// # Examples + /// + /// ``` + /// use datasketches::tuple::CompactTupleSketch; + /// use datasketches::tuple::DefaultUnionPolicy; + /// use datasketches::tuple::DefaultUpdatePolicy; + /// use datasketches::tuple::TupleSketchBuilder; + /// use datasketches::tuple::TupleUnionBuilder; + /// + /// let mut sketch = TupleSketchBuilder::new(DefaultUpdatePolicy::::default()) + /// .build() + /// .unwrap(); + /// sketch.update("apple", 1); + /// + /// let mut union = TupleUnionBuilder::new(DefaultUnionPolicy::::default()) + /// .build() + /// .unwrap(); + /// union.update(&sketch).unwrap(); + /// + /// let bytes = union.serialize(true); + /// let result = CompactTupleSketch::::deserialize(&bytes).unwrap(); + /// assert_eq!(result.num_retained(), 1); + /// ``` + pub fn serialize(&self, ordered: bool) -> Vec + where + P::Summary: TupleSummaryValue, + { + serialize_compact_uncompressed(&self.state.to_borrowed_compact_sketch_state(ordered)) + } + /// Resets the union to its initial empty state. pub fn reset(&mut self) { self.state.reset(); diff --git a/tests-integration/tests/serde_tests/theta.rs b/tests-integration/tests/serde_tests/theta.rs index 4f6dd0f9..04857ca6 100644 --- a/tests-integration/tests/serde_tests/theta.rs +++ b/tests-integration/tests/serde_tests/theta.rs @@ -23,6 +23,7 @@ use datasketches::common::NumStdDev; use datasketches::error::ErrorKind; use datasketches::theta::CompactThetaSketch; use datasketches::theta::ThetaSketchBuilder; +use datasketches::theta::ThetaUnionBuilder; use googletest::assert_that; use googletest::prelude::near; use tests_integration::ZERO_HASH_SEED; @@ -318,3 +319,68 @@ fn test_v2_java_empty_layouts() { } } } + +#[test] +fn union_serialization_matches_compact_result() { + for (rows, probability, mixed_precision) in [ + (0, 1.0, false), + (1, 1.0, false), + (12, 1.0, false), + (40, 1.0, false), + (64, 0.5, false), + (64, 1e-12, false), + (64, 1.0, true), + ] { + let seed = 123; + let mut input = ThetaSketchBuilder::default() + .lg_k(8) + .seed(seed) + .sampling_probability(probability) + .build() + .unwrap(); + for value in 0..rows { + input.update(value); + } + let mut union = ThetaUnionBuilder::default() + .lg_k(5) + .seed(seed) + .build() + .unwrap(); + for ordered in [false, true] { + assert_eq!( + union.serialize(ordered), + union.to_sketch(ordered).serialize() + ); + } + union.update(&input).unwrap(); + if mixed_precision { + let mut other = ThetaSketchBuilder::default() + .lg_k(5) + .seed(seed) + .build() + .unwrap(); + for value in 32..96 { + other.update(value); + } + union.update(&other.compact(true)).unwrap(); + } + + for ordered in [false, true] { + let bytes = union.serialize(ordered); + assert_eq!( + bytes, + union.to_sketch(ordered).serialize(), + "rows={rows}, probability={probability}, mixed={mixed_precision}, ordered={ordered}", + ); + let restored = CompactThetaSketch::deserialize_with_seed(&bytes, seed).unwrap(); + assert_eq!(restored.serialize(), bytes); + } + union.reset(); + for ordered in [false, true] { + assert_eq!( + union.serialize(ordered), + union.to_sketch(ordered).serialize() + ); + } + } +} diff --git a/tests-integration/tests/serde_tests/tuple.rs b/tests-integration/tests/serde_tests/tuple.rs index 3a71e25a..679923ce 100644 --- a/tests-integration/tests/serde_tests/tuple.rs +++ b/tests-integration/tests/serde_tests/tuple.rs @@ -15,13 +15,25 @@ // specific language governing permissions and limitations // under the License. +use std::cell::Cell; +use std::fmt::Debug; use std::fs; +use std::io; use std::path::PathBuf; +use datasketches::codec::SketchBytes; +use datasketches::codec::SketchSlice; +use datasketches::error::Error; use datasketches::error::ErrorKind; use datasketches::tuple::CompactTupleSketch; +use datasketches::tuple::DefaultUnionPolicy; use datasketches::tuple::DefaultUpdatePolicy; +use datasketches::tuple::SummaryCombinePolicy; +use datasketches::tuple::SummaryPolicy; +use datasketches::tuple::SummaryUpdatePolicy; use datasketches::tuple::TupleSketchBuilder; +use datasketches::tuple::TupleSummaryValue; +use datasketches::tuple::TupleUnionBuilder; use googletest::assert_that; use googletest::prelude::each; use googletest::prelude::eq; @@ -214,3 +226,196 @@ fn declared_entry_payload_is_checked_before_allocating() { assert!(CompactTupleSketch::::deserialize(&bytes).is_err()); } + +/// Asserts that two serialized compact Tuple sketches hold the same result. Unordered images may +/// list the same entries in a different order. +fn assert_same_compact_result(actual: &[u8], expected: &[u8], seed: u64) +where + S: TupleSummaryValue + PartialEq + Debug, +{ + let actual = CompactTupleSketch::::deserialize_with_seed(actual, seed).unwrap(); + let expected = CompactTupleSketch::::deserialize_with_seed(expected, seed).unwrap(); + assert_eq!(actual.is_empty(), expected.is_empty()); + assert_eq!(actual.is_ordered(), expected.is_ordered()); + assert_eq!(actual.theta64(), expected.theta64()); + assert_eq!(actual.seed_hash(), expected.seed_hash()); + + fn sorted_entries(sketch: &CompactTupleSketch) -> Vec<(u64, &S)> { + let mut entries: Vec<_> = sketch + .iter() + .map(|entry| (entry.hash(), entry.summary())) + .collect(); + entries.sort_unstable_by_key(|(hash, _)| *hash); + entries + } + assert_eq!(sorted_entries(&actual), sorted_entries(&expected)); +} + +#[test] +fn union_serialization_matches_compact_result() { + let union_builder = || TupleUnionBuilder::new(DefaultUnionPolicy::::default()); + let empty = union_builder().build().unwrap(); + for ordered in [true, false] { + assert_eq!( + empty.serialize(ordered), + empty.to_sketch(ordered).serialize() + ); + } + + // Covers exact mode with zero, one and many entries, trimming to nominal capacity, a sampled + // union, and a non-empty result whose entries were all screened by theta. + let seed = 123; + for (lg_k, sampling_probability, input_probability, count) in [ + (5, 1.0, 1.0, 0), + (5, 1.0, 1.0, 1), + (5, 1.0, 1.0, 20), + (5, 1.0, 1.0, 10_000), + (12, 1.0, 1.0, 10_000), + (8, 0.5, 1.0, 10_000), + (5, 1e-12, 1.0, 10), + (5, 1.0, 0.1, 100), + ] { + let mut a = TupleSketchBuilder::new(DefaultUpdatePolicy::::default()) + .seed(seed) + .build() + .unwrap(); + let mut b = TupleSketchBuilder::new(DefaultUpdatePolicy::::default()) + .seed(seed) + .sampling_probability(input_probability) + .build() + .unwrap(); + for value in 0..count { + a.update(value, 1u64); + b.update(value + count / 2, 2u64); + } + let mut union = union_builder() + .seed(seed) + .lg_k(lg_k) + .sampling_probability(sampling_probability) + .build() + .unwrap(); + union.update(&a).unwrap(); + union.update(&b).unwrap(); + + for ordered in [true, false] { + let bytes = union.serialize(ordered); + let expected = union.to_sketch(ordered).serialize(); + if ordered { + assert_eq!(bytes, expected); + } + assert_same_compact_result::(&bytes, &expected, seed); + } + } +} + +/// Variable-size, heap-backed summary: the tags recorded for a key. +#[derive(Debug, Default, PartialEq)] +struct Tags(Vec); + +thread_local! { + static SUMMARY_CLONES: Cell = const { Cell::new(0) }; +} + +impl Clone for Tags { + fn clone(&self) -> Self { + SUMMARY_CLONES.with(|count| count.set(count.get() + 1)); + Self(self.0.clone()) + } +} + +impl TupleSummaryValue for Tags { + fn serialize_size(&self) -> usize { + 4 + self.0.iter().map(|tag| 4 + tag.len()).sum::() + } + + fn serialize_value(&self, bytes: &mut SketchBytes) { + bytes.write_u32_le(self.0.len() as u32); + for tag in &self.0 { + bytes.write_u32_le(tag.len() as u32); + bytes.write(tag.as_bytes()); + } + } + + fn deserialize_value(cursor: &mut SketchSlice<'_>) -> Result { + let malformed = |_: io::Error| Error::new(ErrorKind::InvalidData, "malformed tags"); + let count = cursor.read_u32_le().map_err(malformed)?; + let mut tags = Vec::new(); + for _ in 0..count { + let mut tag = vec![0; cursor.read_u32_le().map_err(malformed)? as usize]; + cursor.read_exact(&mut tag).map_err(malformed)?; + tags.push( + String::from_utf8(tag) + .map_err(|_| Error::new(ErrorKind::InvalidData, "tag is not UTF-8"))?, + ); + } + Ok(Tags(tags)) + } +} + +#[derive(Debug, Clone, Copy)] +struct TagsPolicy; + +impl SummaryPolicy for TagsPolicy { + type Summary = Tags; + + fn create(&self) -> Self::Summary { + Tags::default() + } +} + +impl SummaryUpdatePolicy<&str> for TagsPolicy { + fn update(&self, summary: &mut Self::Summary, value: &str) { + summary.0.push(value.to_string()); + } +} + +impl SummaryCombinePolicy for TagsPolicy { + fn combine(&self, summary: &mut Self::Summary, other: &Self::Summary) { + summary.0.extend(other.0.iter().cloned()); + } +} + +#[test] +fn union_serialization_with_heap_backed_summaries_supports_later_merges() { + for lg_k in [5, 9] { + let seed = 123; + let sketch = |values: std::ops::Range, tag: &str| { + let mut sketch = TupleSketchBuilder::new(TagsPolicy) + .lg_k(lg_k) + .seed(seed) + .build() + .unwrap(); + for value in values { + sketch.update(value, tag); + } + sketch + }; + let union_builder = || TupleUnionBuilder::new(TagsPolicy).lg_k(lg_k).seed(seed); + + let mut union = union_builder().build().unwrap(); + union.update(&sketch(0..1_000, "first")).unwrap(); + union + .update(&sketch(500..1_500, "a longer second tag")) + .unwrap(); + let later = sketch(1_000..2_000, "third"); + + for ordered in [true, false] { + let clones = SUMMARY_CLONES.get(); + let bytes = union.serialize(ordered); + assert_eq!(SUMMARY_CLONES.get(), clones); + let expected = union.to_sketch(ordered).serialize(); + assert_same_compact_result::(&bytes, &expected, seed); + + let restored = CompactTupleSketch::::deserialize_with_seed(&bytes, seed).unwrap(); + let mut from_bytes = union_builder().build().unwrap(); + from_bytes.update(&restored).unwrap(); + from_bytes.update(&later).unwrap(); + + let mut from_sketch = union_builder().build().unwrap(); + from_sketch.update(&union.to_sketch(ordered)).unwrap(); + from_sketch.update(&later).unwrap(); + + assert_eq!(from_bytes.serialize(true), from_sketch.serialize(true)); + } + } +}