perf(tuple): serialize union results without cloning summaries

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 `TupleUnion::serialize` to serialize union results without cloning summaries, which avoids per-entry allocations for heap-backed summaries.
 * 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/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,162 @@
+// 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;
+
+const ROWS: [u64; 4] = [0, 4, 100, 80_000];
+
+fn u64_union(rows: u64) -> TupleUnion<DefaultUnionPolicy<u64>> {
+    let mut sketch = TupleSketchBuilder::new(DefaultUpdatePolicy::<u64>::default())
+        .build()
+        .unwrap();
+    for value in 0..rows {
+        sketch.update(value, 1);
+    }
+    let mut union = TupleUnionBuilder::new(DefaultUnionPolicy::<u64>::default())
+        .build()
+        .unwrap();
+    union.update(&sketch).unwrap();
+    union
+}
+
+fn serialize_union<const ORDERED: bool>(bencher: Bencher, rows: u64) {
+    let union = u64_union(rows);
+    bencher.bench_local(|| black_box(&union).serialize(ORDERED));
+}
+
+fn serialize_union_via_sketch<const ORDERED: bool>(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<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) -> TupleUnion<LabelsPolicy> {
+    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<const ORDERED: bool>(bencher: Bencher, rows: u64) {
+    let union = labels_union(rows);
+    bencher.bench_local(|| black_box(&union).serialize(ORDERED));
+}
+
+fn serialize_labels_union_via_sketch<const ORDERED: bool>(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) {
+    for rows in ROWS {
+        register_order::<false>(runner, rows);
+        register_order::<true>(runner, rows);
+    }
+}
+
+fn register_order<const ORDERED: bool>(runner: &mut Runner, rows: u64) {
+    let order = if ORDERED { "ordered" } else { "unordered" };
+    let rounds = if rows >= 100 { 500 } else { 10_000 };
+    runner.compare(
+        &format!("tuple::serde::u64::{order}::{rows}"),
+        rounds,
+        &[
+            ("via_compact_sketch", &|bencher| {
+                serialize_union_via_sketch::<ORDERED>(bencher, rows)
+            }),
+            ("direct", &|bencher| {
+                serialize_union::<ORDERED>(bencher, rows)
+            }),
+        ],
+    );
+    runner.compare(
+        &format!("tuple::serde::labels::{order}::{rows}"),
+        rounds,
+        &[
+            ("via_compact_sketch", &|bencher| {
+                serialize_labels_union_via_sketch::<ORDERED>(bencher, rows)
+            }),
+            ("direct", &|bencher| {
+                serialize_labels_union::<ORDERED>(bencher, rows)
+            }),
+        ],
+    );
+}

datasketches/src/thetafamily/common/mod.rs

@@ -33,6 +33,12 @@ pub(super) trait SketchEntry {
     fn hash(&self) -> u64;
 }
 
+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/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, _) =

datasketches/src/thetafamily/tuple/sketch.rs

@@ -22,6 +22,7 @@
 //! created by a [`SummaryPolicy`] and updated through one or more [`SummaryUpdatePolicy`]
 //! implementations.
 
+use std::borrow::Borrow;
 use std::hash::Hash;
 use std::slice;
 
@@ -442,6 +443,60 @@ where
     }
 }
 
+/// Serializes compact Tuple state whose retained entries are owned or borrowed.
+pub(super) fn serialize_compact_state<S, E>(compact_state: &CompactSketchState<E>) -> Vec<u8>
+where
+    S: TupleSummaryValue,
+    E: Borrow<TupleEntry<S>>,
+{
+    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()
+}
+
 /// Compact (immutable) Tuple sketch.
 ///
 /// This is the serialization-friendly form: a compact array of retained hash-summary pairs plus
@@ -546,16 +601,6 @@ impl<S> CompactTupleSketch<S> {
         self.compact_state.retained_entries_capacity() * size_of::<TupleEntry<S>>()
     }
 
-    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 +624,7 @@ impl<S> CompactTupleSketch<S> {
     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.