feat(theta): serialize union results directly

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

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);
 }

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));
+                    }),
+                ],
+            );
+        }
+    }
+}

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<E>(
+    state: &CompactSketchState<E>,
+    hash: impl Fn(&E) -> u64,
+) -> Vec<u8> {
+    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()
+}

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<u8> {
-        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<u8> {
         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());
 

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<u8> {
+        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();

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}",
+            );
+        }
+    }
+}