From 0a287e2d2d296cdfb4a45d9e08430942452e154d Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 7 Nov 2026 14:39:22 +0800 Subject: [PATCH 1/7] feat(theta,tuple): serialize union results directly Add compact serialization to ThetaUnion and TupleUnion with shared result filtering, trimming, and ordering. Reuse each family's compact encoder and borrow Tuple entries to avoid cloning summaries during serialization. Cover canonical results, custom seeds, summary clone avoidance, and subsequent merges. Compare direct serialization with compact-sketch serialization for scalar and heap-backed summaries. Co-authored-by: Nash Aschenbath Co-authored-by: Claude Opus 5.5 --- CHANGELOG.md | 1 + benchmarks/theta/mod.rs | 2 + benchmarks/theta/serde.rs | 64 ++++++ benchmarks/tuple/mod.rs | 2 + benchmarks/tuple/serde.rs | 170 +++++++++++++++ datasketches/src/thetafamily/common/mod.rs | 6 + datasketches/src/thetafamily/common/union.rs | 36 +++- .../src/thetafamily/theta/serialization.rs | 55 ++++- datasketches/src/thetafamily/theta/sketch.rs | 50 +---- datasketches/src/thetafamily/theta/union.rs | 11 + .../src/thetafamily/tuple/serialization.rs | 62 ++++++ datasketches/src/thetafamily/tuple/sketch.rs | 54 +---- datasketches/src/thetafamily/tuple/union.rs | 40 ++++ tests-integration/tests/serde_tests/theta.rs | 66 ++++++ tests-integration/tests/serde_tests/tuple.rs | 202 ++++++++++++++++++ 15 files changed, 710 insertions(+), 111 deletions(-) create mode 100644 benchmarks/theta/serde.rs create mode 100644 benchmarks/tuple/serde.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d7158fb..66b25615 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ All significant changes to this project will be documented in this file. * 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..096c7d03 --- /dev/null +++ b/benchmarks/tuple/serde.rs @@ -0,0 +1,170 @@ +// 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) -> TupleUnion> { + let mut sketch = TupleSketchBuilder::new(DefaultUpdatePolicy::::default()) + .build() + .unwrap(); + for value in 0..rows { + sketch.update(value, 1); + } + let mut union = TupleUnionBuilder::new(DefaultUnionPolicy::::default()) + .build() + .unwrap(); + union.update(&sketch).unwrap(); + union +} + +fn serialize_union(bencher: Bencher, rows: u64) { + let union = u64_union(rows); + bencher.bench_local(|| black_box(&union).serialize(ORDERED)); +} + +fn serialize_union_via_sketch(bencher: Bencher, rows: u64) { + let union = u64_union(rows); + 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) -> TupleUnion { + let mut sketch = TupleSketchBuilder::new(LabelsPolicy).build().unwrap(); + for value in 0..rows { + sketch.update(value, value); + } + let mut union = TupleUnionBuilder::new(LabelsPolicy).build().unwrap(); + union.update(&sketch).unwrap(); + union +} + +fn serialize_labels_union(bencher: Bencher, rows: u64) { + let union = labels_union(rows); + bencher.bench_local(|| black_box(&union).serialize(ORDERED)); +} + +fn serialize_labels_union_via_sketch(bencher: Bencher, rows: u64) { + let union = labels_union(rows); + 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. + for (rows, u64_rounds, labels_rounds) in [ + (0, 10_000_000, 10_000_000), + (4, 2_000_000, 500_000), + (100, 300_000, 20_000), + (80_000, 5_000, 500), + ] { + register_order::(runner, rows, u64_rounds, labels_rounds); + register_order::(runner, rows, u64_rounds, labels_rounds); + } +} + +fn register_order( + runner: &mut Runner, + rows: u64, + u64_rounds: usize, + labels_rounds: usize, +) { + let order = if ORDERED { "ordered" } else { "unordered" }; + runner.compare( + &format!("tuple::serde::u64::{order}/{rows}"), + u64_rounds, + &[ + ("via_compact_sketch", &|bencher| { + serialize_union_via_sketch::(bencher, rows) + }), + ("direct", &|bencher| { + serialize_union::(bencher, rows) + }), + ], + ); + runner.compare( + &format!("tuple::serde::labels::{order}/{rows}"), + labels_rounds, + &[ + ("via_compact_sketch", &|bencher| { + serialize_labels_union_via_sketch::(bencher, rows) + }), + ("direct", &|bencher| { + serialize_labels_union::(bencher, rows) + }), + ], + ); +} diff --git a/datasketches/src/thetafamily/common/mod.rs b/datasketches/src/thetafamily/common/mod.rs index dbc273ee..358fb756 100644 --- a/datasketches/src/thetafamily/common/mod.rs +++ b/datasketches/src/thetafamily/common/mod.rs @@ -33,6 +33,12 @@ pub(super) trait SketchEntry { fn hash(&self) -> u64; } +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/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..003e0a05 100644 --- a/datasketches/src/thetafamily/theta/serialization.rs +++ b/datasketches/src/thetafamily/theta/serialization.rs @@ -15,7 +15,14 @@ // 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::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 +33,49 @@ 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_uncompressed( + state: &CompactSketchState, + hash: impl Fn(&E) -> u64, +) -> Vec { + if state.is_empty() { + return EMPTY_SKETCH_BYTES.to_vec(); + } + + let entries = state.retained_entries(); + let pre_longs = if state.is_estimation_mode() { + 3 + } else if entries.len() == 1 { + 1 + } else { + 2 + }; + 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(hash(entry)); + } + bytes.into_bytes() +} diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index 4ec71b79..02c98149 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -532,18 +532,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. @@ -569,46 +557,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_uncompressed(&self.compact_state, |hash| *hash) } 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..500b277b 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_uncompressed(&state, ThetaEntry::hash) + } + /// 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..fee1c667 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,60 @@ 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_state(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 = if is_estimation_mode { + 3 + } else if retained_entries.len() == 1 { + 1 + } else { + 2 + }; + 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 0fee8218..431c5c9d 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_state; /// Read-only view of a mutable or compact Tuple sketch. /// @@ -546,16 +543,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 @@ -579,44 +566,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_state(&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..e121cd2f 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_state; 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_state(&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..3683fd6f 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,193 @@ 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), + (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() { + let seed = 123; + let sketch = |values: std::ops::Range, tag: &str| { + let mut sketch = TupleSketchBuilder::new(TagsPolicy) + .lg_k(5) + .seed(seed) + .build() + .unwrap(); + for value in values { + sketch.update(value, tag); + } + sketch + }; + let union_builder = || TupleUnionBuilder::new(TagsPolicy).lg_k(5).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)); + } +} From 73da92a1db587b09c15237e25504e617806fb6ae Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 10 Nov 2026 15:02:08 +0800 Subject: [PATCH 2/7] refactor(theta): use SketchEntry for compact serialization Implement the internal SketchEntry trait for retained u64 hashes so the compact encoder can read both compact-sketch hashes and union entries through the same interface without a hash projection callback. --- datasketches/src/thetafamily/common/mod.rs | 6 ++++++ datasketches/src/thetafamily/theta/serialization.rs | 8 +++----- datasketches/src/thetafamily/theta/sketch.rs | 2 +- datasketches/src/thetafamily/theta/union.rs | 2 +- 4 files changed, 11 insertions(+), 7 deletions(-) diff --git a/datasketches/src/thetafamily/common/mod.rs b/datasketches/src/thetafamily/common/mod.rs index 358fb756..e1a9ae89 100644 --- a/datasketches/src/thetafamily/common/mod.rs +++ b/datasketches/src/thetafamily/common/mod.rs @@ -33,6 +33,12 @@ 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) diff --git a/datasketches/src/thetafamily/theta/serialization.rs b/datasketches/src/thetafamily/theta/serialization.rs index 003e0a05..e4695b84 100644 --- a/datasketches/src/thetafamily/theta/serialization.rs +++ b/datasketches/src/thetafamily/theta/serialization.rs @@ -19,6 +19,7 @@ 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; @@ -34,10 +35,7 @@ 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_uncompressed( - state: &CompactSketchState, - hash: impl Fn(&E) -> u64, -) -> Vec { +pub(super) fn serialize_uncompressed(state: &CompactSketchState) -> Vec { if state.is_empty() { return EMPTY_SKETCH_BYTES.to_vec(); } @@ -75,7 +73,7 @@ pub(super) fn serialize_uncompressed( bytes.write_u64_le(state.theta()); } for entry in entries { - bytes.write_u64_le(hash(entry)); + 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 02c98149..d6aa630b 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -557,7 +557,7 @@ impl CompactThetaSketch { /// /// Empty sketches serialize with a zero seed hash. pub fn serialize(&self) -> Vec { - serialization::serialize_uncompressed(&self.compact_state, |hash| *hash) + serialization::serialize_uncompressed(&self.compact_state) } fn serialize_v4(&self) -> Vec { diff --git a/datasketches/src/thetafamily/theta/union.rs b/datasketches/src/thetafamily/theta/union.rs index 500b277b..95dfb5e0 100644 --- a/datasketches/src/thetafamily/theta/union.rs +++ b/datasketches/src/thetafamily/theta/union.rs @@ -67,7 +67,7 @@ impl ThetaUnion { /// singletons always use their canonical encoding. pub fn serialize(&self, ordered: bool) -> Vec { let state = self.state.to_compact_sketch_state(ordered); - serialization::serialize_uncompressed(&state, ThetaEntry::hash) + serialization::serialize_uncompressed(&state) } /// Resets the union to its empty state. From 8740495e809b124df2297cf6ecf98d37a9e42a88 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 10 Nov 2026 15:02:38 +0800 Subject: [PATCH 3/7] refactor(theta,tuple): align compact encoder names Name both family encoders serialize_compact_uncompressed to describe the compact representation and the uncompressed wire format consistently. Keep the family-specific headers and entry encoding in their respective serialization modules. --- datasketches/src/thetafamily/theta/serialization.rs | 4 +++- datasketches/src/thetafamily/theta/sketch.rs | 2 +- datasketches/src/thetafamily/theta/union.rs | 2 +- datasketches/src/thetafamily/tuple/serialization.rs | 2 +- datasketches/src/thetafamily/tuple/sketch.rs | 4 ++-- datasketches/src/thetafamily/tuple/union.rs | 4 ++-- 6 files changed, 10 insertions(+), 8 deletions(-) diff --git a/datasketches/src/thetafamily/theta/serialization.rs b/datasketches/src/thetafamily/theta/serialization.rs index e4695b84..9f30bd16 100644 --- a/datasketches/src/thetafamily/theta/serialization.rs +++ b/datasketches/src/thetafamily/theta/serialization.rs @@ -35,7 +35,9 @@ 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_uncompressed(state: &CompactSketchState) -> Vec { +pub(super) fn serialize_compact_uncompressed( + state: &CompactSketchState, +) -> Vec { if state.is_empty() { return EMPTY_SKETCH_BYTES.to_vec(); } diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index d6aa630b..44088273 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -557,7 +557,7 @@ impl CompactThetaSketch { /// /// Empty sketches serialize with a zero seed hash. pub fn serialize(&self) -> Vec { - serialization::serialize_uncompressed(&self.compact_state) + serialization::serialize_compact_uncompressed(&self.compact_state) } fn serialize_v4(&self) -> Vec { diff --git a/datasketches/src/thetafamily/theta/union.rs b/datasketches/src/thetafamily/theta/union.rs index 95dfb5e0..a8f21faf 100644 --- a/datasketches/src/thetafamily/theta/union.rs +++ b/datasketches/src/thetafamily/theta/union.rs @@ -67,7 +67,7 @@ impl ThetaUnion { /// singletons always use their canonical encoding. pub fn serialize(&self, ordered: bool) -> Vec { let state = self.state.to_compact_sketch_state(ordered); - serialization::serialize_uncompressed(&state) + serialization::serialize_compact_uncompressed(&state) } /// Resets the union to its empty state. diff --git a/datasketches/src/thetafamily/tuple/serialization.rs b/datasketches/src/thetafamily/tuple/serialization.rs index fee1c667..634eba97 100644 --- a/datasketches/src/thetafamily/tuple/serialization.rs +++ b/datasketches/src/thetafamily/tuple/serialization.rs @@ -52,7 +52,7 @@ pub const SKETCH_TYPE_LEGACY: u8 = 5; 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_state(compact_state: &CompactSketchState) -> Vec +pub(super) fn serialize_compact_uncompressed(compact_state: &CompactSketchState) -> Vec where S: TupleSummaryValue, E: Borrow>, diff --git a/datasketches/src/thetafamily/tuple/sketch.rs b/datasketches/src/thetafamily/tuple/sketch.rs index 431c5c9d..2aa93221 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -55,7 +55,7 @@ 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_state; +use crate::tuple::serialization::serialize_compact_uncompressed; /// Read-only view of a mutable or compact Tuple sketch. /// @@ -566,7 +566,7 @@ impl CompactTupleSketch { where S: TupleSummaryValue, { - serialize_compact_state(&self.compact_state) + 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 e121cd2f..f2d419a2 100644 --- a/datasketches/src/thetafamily/tuple/union.rs +++ b/datasketches/src/thetafamily/tuple/union.rs @@ -30,7 +30,7 @@ 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_state; +use crate::tuple::serialization::serialize_compact_uncompressed; use crate::tuple::sketch::CompactTupleSketch; use crate::tuple::sketch::TupleSketchView; @@ -144,7 +144,7 @@ where where P::Summary: TupleSummaryValue, { - serialize_compact_state(&self.state.to_borrowed_compact_sketch_state(ordered)) + serialize_compact_uncompressed(&self.state.to_borrowed_compact_sketch_state(ordered)) } /// Resets the union to its initial empty state. From 4a3ec3c8f85e92caef1afffc11614a78e920fad0 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 10 Nov 2026 15:04:06 +0800 Subject: [PATCH 4/7] refactor(theta,tuple): share uncompressed preamble sizing Derive the compact preamble length from the canonical sketch state in one place for both family encoders: one 8-byte word for empty images and exact singletons, two for other exact results, and three for estimation results. The helper is specific to uncompressed formats. Compressed Theta has a different header layout, so its size calculation belongs with that codec rather than behind a compression flag on the shared helper. --- datasketches/src/thetafamily/common/sketch_state.rs | 13 +++++++++++++ datasketches/src/thetafamily/theta/serialization.rs | 8 +------- datasketches/src/thetafamily/tuple/serialization.rs | 8 +------- 3 files changed, 15 insertions(+), 14 deletions(-) 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/theta/serialization.rs b/datasketches/src/thetafamily/theta/serialization.rs index 9f30bd16..b25edcf6 100644 --- a/datasketches/src/thetafamily/theta/serialization.rs +++ b/datasketches/src/thetafamily/theta/serialization.rs @@ -43,13 +43,7 @@ pub(super) fn serialize_compact_uncompressed( } let entries = state.retained_entries(); - let pre_longs = if state.is_estimation_mode() { - 3 - } else if entries.len() == 1 { - 1 - } else { - 2 - }; + 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); diff --git a/datasketches/src/thetafamily/tuple/serialization.rs b/datasketches/src/thetafamily/tuple/serialization.rs index 634eba97..1e7b9164 100644 --- a/datasketches/src/thetafamily/tuple/serialization.rs +++ b/datasketches/src/thetafamily/tuple/serialization.rs @@ -63,13 +63,7 @@ where let retained_entries = compact_state.retained_entries(); let is_estimation_mode = compact_state.is_estimation_mode(); - let pre_longs = if is_estimation_mode { - 3 - } else if retained_entries.len() == 1 { - 1 - } else { - 2 - }; + let pre_longs = compact_state.uncompressed_preamble_longs(); let entries_size: usize = retained_entries .iter() .map(|entry| 8 + entry.borrow().summary().serialize_size()) From 84fd913907eb9b494422d4bc6f75ae761a3adfcb Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 10 Nov 2026 15:27:46 +0800 Subject: [PATCH 5/7] bench(tuple): cover union serialization across retained sizes Add a larger exact result and estimation workloads with lg_k 8 and 16 alongside the default-capacity cases. Compare ordered and unordered serialization for scalar and heap-backed summaries with fixed rounds so changes to result preparation can be checked across retained sizes. --- benchmarks/tuple/serde.rs | 68 +++++++++++++++++++++++++-------------- 1 file changed, 43 insertions(+), 25 deletions(-) diff --git a/benchmarks/tuple/serde.rs b/benchmarks/tuple/serde.rs index 096c7d03..b49b0648 100644 --- a/benchmarks/tuple/serde.rs +++ b/benchmarks/tuple/serde.rs @@ -32,27 +32,29 @@ use datasketches::tuple::TupleSummaryValue; use datasketches::tuple::TupleUnion; use datasketches::tuple::TupleUnionBuilder; -fn u64_union(rows: u64) -> TupleUnion> { +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) { - let union = u64_union(rows); +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) { - let union = u64_union(rows); +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()); } @@ -103,67 +105,83 @@ impl SummaryCombinePolicy for LabelsPolicy { } } -fn labels_union(rows: u64) -> TupleUnion { - let mut sketch = TupleSketchBuilder::new(LabelsPolicy).build().unwrap(); +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).build().unwrap(); + 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) { - let union = labels_union(rows); +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) { - let union = labels_union(rows); +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. - for (rows, u64_rounds, labels_rounds) in [ - (0, 10_000_000, 10_000_000), - (4, 2_000_000, 500_000), - (100, 300_000, 20_000), - (80_000, 5_000, 500), + // 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, 30_000, 3_000), + (80_000, 12, 5_000, 500), + (80_000, 16, 250, 20), ] { - register_order::(runner, rows, u64_rounds, labels_rounds); - register_order::(runner, rows, u64_rounds, labels_rounds); + 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}/{rows}"), + &format!("tuple::serde::u64::{order}/{size}"), u64_rounds, &[ ("via_compact_sketch", &|bencher| { - serialize_union_via_sketch::(bencher, rows) + serialize_union_via_sketch::(bencher, rows, lg_k) }), ("direct", &|bencher| { - serialize_union::(bencher, rows) + serialize_union::(bencher, rows, lg_k) }), ], ); runner.compare( - &format!("tuple::serde::labels::{order}/{rows}"), + &format!("tuple::serde::labels::{order}/{size}"), labels_rounds, &[ ("via_compact_sketch", &|bencher| { - serialize_labels_union_via_sketch::(bencher, rows) + serialize_labels_union_via_sketch::(bencher, rows, lg_k) }), ("direct", &|bencher| { - serialize_labels_union::(bencher, rows) + serialize_labels_union::(bencher, rows, lg_k) }), ], ); From f21d6c2cdf5508fad3f81fe4fb8c106097b6f83e Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 10 Nov 2026 15:38:43 +0800 Subject: [PATCH 6/7] perf(tuple): bucket borrowed union entries before sorting Global comparison sorting repeatedly dereferences the union hash table when ordering borrowed entries. Partition large ordered results by their top active hash bits, scaled to the result theta, and sort within each bucket to reduce those indirect comparisons without cloning summaries. Keep small and unordered results on the existing compaction path and keep shared compaction inlineable. Large ordered results use one extra reference buffer, or 32 KiB for 4096 retained entries on 64-bit targets. Cover ordering with spread and clustered hashes, small-result boundaries, and larger serialized results whose heap-backed summaries must not be cloned. Size the lg_k 8 benchmark rounds for the measured workload. --- benchmarks/tuple/serde.rs | 2 +- datasketches/src/thetafamily/common/union.rs | 76 +++++++++++++++++++- tests-integration/tests/serde_tests/tuple.rs | 75 +++++++++---------- 3 files changed, 114 insertions(+), 39 deletions(-) diff --git a/benchmarks/tuple/serde.rs b/benchmarks/tuple/serde.rs index b49b0648..700675bf 100644 --- a/benchmarks/tuple/serde.rs +++ b/benchmarks/tuple/serde.rs @@ -139,7 +139,7 @@ pub fn register(runner: &mut Runner) { (4, 12, 2_000_000, 500_000), (100, 12, 300_000, 20_000), (1_000, 12, 30_000, 2_000), - (80_000, 8, 30_000, 3_000), + (80_000, 8, 100_000, 6_000), (80_000, 12, 5_000, 500), (80_000, 16, 250, 20), ] { diff --git a/datasketches/src/thetafamily/common/union.rs b/datasketches/src/thetafamily/common/union.rs index b9e5d1db..df735ad0 100644 --- a/datasketches/src/thetafamily/common/union.rs +++ b/datasketches/src/thetafamily/common/union.rs @@ -122,12 +122,28 @@ where /// Returns the union as canonical compact-sketch state that borrows the retained entries. #[cfg(feature = "tuple")] + #[inline] 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) + let entries = self.retained_entries(theta).collect::>(); + if !ordered || entries.len() < 256 { + return self.finish_compaction(entries, theta, ordered); + } + + let mut state = self.finish_compaction(entries, theta, false); + if let CompactSketchState::NonEmpty { + retained_entries, + theta, + ordered, + .. + } = &mut state + { + sort_borrowed_entries(retained_entries, *theta); + *ordered = true; + } + state } fn retained_entries(&self, theta: u64) -> impl Iterator { @@ -137,6 +153,7 @@ where .filter(move |entry| all_retained || entry.hash() < theta) } + #[inline] fn finish_compaction( &self, mut retained_entries: Vec, @@ -171,3 +188,58 @@ where size_of::

() + self.table.estimated_size() } } + +#[cfg(feature = "tuple")] +fn sort_borrowed_entries(entries: &mut [&E], theta: u64) { + // Partition by the top eight active hash bits before comparison sorting. Hashes are + // distributed below theta, so scaling the buckets to theta also works in estimation mode. + // Small buckets reduce comparisons that dereference entries in the original hash table. + let shift = (u64::BITS - theta.leading_zeros()).saturating_sub(8); + let mut offsets = [0usize; 256]; + for entry in entries.iter() { + offsets[(entry.hash() >> shift) as usize] += 1; + } + let mut total = 0; + for offset in &mut offsets { + let count = *offset; + *offset = total; + total += count; + } + let mut scratch = entries.to_vec(); + for &entry in entries.iter() { + let bucket = (entry.hash() >> shift) as usize; + scratch[offsets[bucket]] = entry; + offsets[bucket] += 1; + } + let mut start = 0; + for end in offsets { + scratch[start..end].sort_unstable_by_key(SketchEntry::hash); + start = end; + } + entries.copy_from_slice(&scratch); +} + +#[cfg(all(test, feature = "tuple"))] +mod tests { + use super::sort_borrowed_entries; + use crate::thetacommon::constants::MAX_THETA; + + #[test] + fn borrowed_ordering_matches_hash_order_for_spread_and_clustered_entries() { + for len in [255, 256, 257, 4096] { + for theta in [8192, 1 << 32, MAX_THETA] { + for spacing in [1, theta / (len as u64 + 1)] { + let hashes = (1..=len as u64) + .map(|value| value * spacing) + .collect::>(); + let mut entries = hashes.iter().rev().collect::>(); + entries.rotate_left(len / 3); + + sort_borrowed_entries(&mut entries, theta); + + assert!(entries.iter().copied().eq(hashes.iter())); + } + } + } + } +} diff --git a/tests-integration/tests/serde_tests/tuple.rs b/tests-integration/tests/serde_tests/tuple.rs index 3683fd6f..679923ce 100644 --- a/tests-integration/tests/serde_tests/tuple.rs +++ b/tests-integration/tests/serde_tests/tuple.rs @@ -270,6 +270,7 @@ fn union_serialization_matches_compact_result() { (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), @@ -376,43 +377,45 @@ impl SummaryCombinePolicy for TagsPolicy { #[test] fn union_serialization_with_heap_backed_summaries_supports_later_merges() { - let seed = 123; - let sketch = |values: std::ops::Range, tag: &str| { - let mut sketch = TupleSketchBuilder::new(TagsPolicy) - .lg_k(5) - .seed(seed) - .build() + 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(); - for value in values { - sketch.update(value, tag); - } - sketch - }; - let union_builder = || TupleUnionBuilder::new(TagsPolicy).lg_k(5).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"); + 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)); + 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)); + } } } From 1f8c314d1f614c66a00bf607b715d5569a27245f Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 10 Nov 2026 15:49:06 +0800 Subject: [PATCH 7/7] refactor(tuple): simplify borrowed union result ordering Use shared compaction to trim and sort borrowed entries in one path. Remove the bucket partition, size cutoff, temporary reference buffer, and associated inline hints to keep result construction straightforward. Retain summary borrowing, expanded serialization benchmarks, and integration coverage for large results, zero summary clones, and subsequent merges. --- datasketches/src/thetafamily/common/union.rs | 76 +------------------- 1 file changed, 2 insertions(+), 74 deletions(-) diff --git a/datasketches/src/thetafamily/common/union.rs b/datasketches/src/thetafamily/common/union.rs index df735ad0..b9e5d1db 100644 --- a/datasketches/src/thetafamily/common/union.rs +++ b/datasketches/src/thetafamily/common/union.rs @@ -122,28 +122,12 @@ where /// Returns the union as canonical compact-sketch state that borrows the retained entries. #[cfg(feature = "tuple")] - #[inline] 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::>(); - if !ordered || entries.len() < 256 { - return self.finish_compaction(entries, theta, ordered); - } - - let mut state = self.finish_compaction(entries, theta, false); - if let CompactSketchState::NonEmpty { - retained_entries, - theta, - ordered, - .. - } = &mut state - { - sort_borrowed_entries(retained_entries, *theta); - *ordered = true; - } - state + let entries = self.retained_entries(theta).collect(); + self.finish_compaction(entries, theta, ordered) } fn retained_entries(&self, theta: u64) -> impl Iterator { @@ -153,7 +137,6 @@ where .filter(move |entry| all_retained || entry.hash() < theta) } - #[inline] fn finish_compaction( &self, mut retained_entries: Vec, @@ -188,58 +171,3 @@ where size_of::

() + self.table.estimated_size() } } - -#[cfg(feature = "tuple")] -fn sort_borrowed_entries(entries: &mut [&E], theta: u64) { - // Partition by the top eight active hash bits before comparison sorting. Hashes are - // distributed below theta, so scaling the buckets to theta also works in estimation mode. - // Small buckets reduce comparisons that dereference entries in the original hash table. - let shift = (u64::BITS - theta.leading_zeros()).saturating_sub(8); - let mut offsets = [0usize; 256]; - for entry in entries.iter() { - offsets[(entry.hash() >> shift) as usize] += 1; - } - let mut total = 0; - for offset in &mut offsets { - let count = *offset; - *offset = total; - total += count; - } - let mut scratch = entries.to_vec(); - for &entry in entries.iter() { - let bucket = (entry.hash() >> shift) as usize; - scratch[offsets[bucket]] = entry; - offsets[bucket] += 1; - } - let mut start = 0; - for end in offsets { - scratch[start..end].sort_unstable_by_key(SketchEntry::hash); - start = end; - } - entries.copy_from_slice(&scratch); -} - -#[cfg(all(test, feature = "tuple"))] -mod tests { - use super::sort_borrowed_entries; - use crate::thetacommon::constants::MAX_THETA; - - #[test] - fn borrowed_ordering_matches_hash_order_for_spread_and_clustered_entries() { - for len in [255, 256, 257, 4096] { - for theta in [8192, 1 << 32, MAX_THETA] { - for spacing in [1, theta / (len as u64 + 1)] { - let hashes = (1..=len as u64) - .map(|value| value * spacing) - .collect::>(); - let mut entries = hashes.iter().rev().collect::>(); - entries.rotate_left(len / 3); - - sort_borrowed_entries(&mut entries, theta); - - assert!(entries.iter().copied().eq(hashes.iter())); - } - } - } - } -}