From e42136f704ff1245f9382f748985a863857885ff Mon Sep 17 00:00:00 2001 From: Nash Aschenbath Date: Fri, 6 Nov 2026 12:36:20 +0800 Subject: [PATCH] feat(theta): serialize union results directly --- CHANGELOG.md | 1 + benchmarks/theta/mod.rs | 2 + benchmarks/theta/serde.rs | 48 ++++++++++++++++ .../src/thetafamily/theta/serialization.rs | 55 ++++++++++++++++++- datasketches/src/thetafamily/theta/sketch.rs | 50 +---------------- datasketches/src/thetafamily/theta/union.rs | 10 ++++ tests-integration/tests/serde_tests/theta.rs | 50 +++++++++++++++++ 7 files changed, 167 insertions(+), 49 deletions(-) create mode 100644 benchmarks/theta/serde.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 8fc9c254..22d874e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ All significant changes to this project will be documented in this file. ### New features * Add `HllUnion::serialize` to serialize union results directly. +* Add `ThetaUnion::serialize` to serialize ordered or unordered union results directly in the uncompressed compact Theta format. * Add immutable `XorFilter` behind the `xor` feature, with 8-, 16-, and 32-bit fingerprints, construction from values with automatic deduplication or borrowed distinct precomputed hashes, bounded seeded retries and compact temporary storage, optional cardinality metadata for imported filters, standard serialization for all three widths (Java-compatible for 8/16-bit filters), and validated imports for application-defined storage. ### 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..6cc7378b --- /dev/null +++ b/benchmarks/theta/serde.rs @@ -0,0 +1,48 @@ +// 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::Runner; +use datasketches::theta::ThetaSketchBuilder; +use datasketches::theta::ThetaUnionBuilder; + +pub fn register(runner: &mut Runner) { + for (rows, rounds) in [(40, 200_000), (8_000, 4_000)] { + for ordered in [false, true] { + let order = if ordered { "ordered" } else { "unordered" }; + 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(); + runner.compare( + &format!("theta::serde::{order}/{rows}"), + rounds, + &[ + ("via_compact", &|b| { + b.bench_local(|| black_box(&union).to_sketch(ordered).serialize()); + }), + ("direct", &|b| { + b.bench_local(|| black_box(&union).serialize(ordered)); + }), + ], + ); + } + } +} 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..3e1c59ff 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,15 @@ impl ThetaUnion { CompactThetaSketch::from_compact_state(compact_state) } + /// Serializes this union result into the uncompressed compact Theta format. + /// + /// 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/tests-integration/tests/serde_tests/theta.rs b/tests-integration/tests/serde_tests/theta.rs index 4f6dd0f9..4c711ad7 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,52 @@ 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(); + 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] { + assert_eq!( + union.serialize(ordered), + union.to_sketch(ordered).serialize(), + "rows={rows}, probability={probability}, mixed={mixed_precision}, ordered={ordered}", + ); + } + } +}