feat(theta,tuple): serialize union results directly

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

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

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

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<DefaultUnionPolicy<u64>> {
+    let mut sketch = TupleSketchBuilder::new(DefaultUpdatePolicy::<u64>::default())
+        .lg_k(lg_k)
+        .build()
+        .unwrap();
+    for value in 0..rows {
+        sketch.update(value, 1);
+    }
+    let mut union = TupleUnionBuilder::new(DefaultUnionPolicy::<u64>::default())
+        .lg_k(lg_k)
+        .build()
+        .unwrap();
+    union.update(&sketch).unwrap();
+    union
+}
+
+fn serialize_union<const ORDERED: bool>(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<const ORDERED: bool>(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<String>);
+
+impl TupleSummaryValue for Labels {
+    fn serialize_size(&self) -> usize {
+        4 + self.0.iter().map(|label| 4 + label.len()).sum::<usize>()
+    }
+
+    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<Self, Error> {
+        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<u64> 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<LabelsPolicy> {
+    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<const ORDERED: bool>(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<const ORDERED: bool>(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::<false>(runner, rows, lg_k, u64_rounds, labels_rounds);
+        register_order::<true>(runner, rows, lg_k, u64_rounds, labels_rounds);
+    }
+}
+
+fn register_order<const ORDERED: bool>(
+    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::<ORDERED>(bencher, rows, lg_k)
+            }),
+            ("direct", &|bencher| {
+                serialize_union::<ORDERED>(bencher, rows, lg_k)
+            }),
+        ],
+    );
+    runner.compare(
+        &format!("tuple::serde::labels::{order}/{size}"),
+        labels_rounds,
+        &[
+            ("via_compact_sketch", &|bencher| {
+                serialize_labels_union_via_sketch::<ORDERED>(bencher, rows, lg_k)
+            }),
+            ("direct", &|bencher| {
+                serialize_labels_union::<ORDERED>(bencher, rows, lg_k)
+            }),
+        ],
+    );
+}

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<E: SketchEntry> SketchEntry for &E {
+    fn hash(&self) -> u64 {
+        E::hash(self)
+    }
+}
+
 pub(super) trait KeySketch: Copy {
     fn metadata(self) -> sketch_state::ThetaFamilySketchMetadata;
 

datasketches/src/thetafamily/common/sketch_state.rs

@@ -82,6 +82,19 @@ impl<E> CompactSketchState<E> {
         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,

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::<Vec<_>>()
-        } else {
-            self.table
-                .iter_entries()
-                .filter(|entry| entry.hash() < theta)
-                .cloned()
-                .collect::<Vec<_>>()
+    /// 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<Item = &E> {
+        let all_retained = theta >= self.table.retention_theta();
+        self.table
+            .iter_entries()
+            .filter(move |entry| all_retained || entry.hash() < theta)
+    }
+
+    fn finish_compaction<T: SketchEntry>(
+        &self,
+        mut retained_entries: Vec<T>,
+        result_theta: u64,
+        ordered: bool,
+    ) -> CompactSketchState<T> {
+        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, _) =