Add serialize_to sink API for count-min sketch

Author: cv4g

count/include/count_min.hpp

@@ -326,6 +326,15 @@ class count_min_sketch{
    */
   void serialize(std::ostream& os) const;
 
+  /**
+   * This method serializes the sketch by passing binary fragments to a callback.
+   * The callback must be callable as write_bytes(const void* data, size_t size).
+   * The data pointer is valid only until the callback returns.
+   * @return size in bytes written to the callback
+   */
+  template<typename WriteBytes>
+  size_t serialize_to(WriteBytes&& write_bytes) const;
+
   // This is a convenience alias for users
   // The type returned by the following serialize method
   using vector_bytes = std::vector<uint8_t, typename std::allocator_traits<Allocator>::template rebind_alloc<uint8_t>>;

count/include/count_min_impl.hpp

@@ -268,39 +268,55 @@ return _sketch_array.end();
 
 template<typename W, typename A>
 void count_min_sketch<W,A>::serialize(std::ostream& os) const {
+  serialize_to([&os](const void* data, size_t size) {
+    os.write(static_cast<const char*>(data), size);
+  });
+}
+
+template<typename W, typename A>
+template<typename WriteBytes>
+size_t count_min_sketch<W,A>::serialize_to(WriteBytes&& write_bytes) const {
+  size_t bytes_written = 0;
+  auto write_raw = [&write_bytes, &bytes_written](const void* data, size_t size) {
+    write_bytes(data, size);
+    bytes_written += size;
+  };
+  auto write_value = [&write_raw](const auto& value) {
+    write_raw(&value, sizeof(value));
+  };
+
   // Long 0
   //const uint8_t preamble_longs = is_empty() ? PREAMBLE_LONGS_SHORT : PREAMBLE_LONGS_FULL;
   const uint8_t preamble_longs = PREAMBLE_LONGS_SHORT;
   const uint8_t ser_ver = SERIAL_VERSION_1;
   const uint8_t family_id = FAMILY_ID;
   const uint8_t flags_byte = (is_empty() ? 1 << flags::IS_EMPTY : 0);
   const uint32_t unused32 = NULL_32;
-  write(os, preamble_longs);
-  write(os, ser_ver);
-  write(os, family_id);
-  write(os, flags_byte);
-  write(os, unused32);
+  write_value(preamble_longs);
+  write_value(ser_ver);
+  write_value(family_id);
+  write_value(flags_byte);
+  write_value(unused32);
 
   // Long 1
   const uint32_t nbuckets = _num_buckets;
   const uint8_t nhashes = _num_hashes;
   const uint16_t seed_hash(compute_seed_hash(_seed));
   const uint8_t unused8 =  NULL_8;
-  write(os, nbuckets);
-  write(os, nhashes);
-  write(os, seed_hash);
-  write(os, unused8);
-  if (is_empty()) return; // sketch is empty, no need to write further bytes.
+  write_value(nbuckets);
+  write_value(nhashes);
+  write_value(seed_hash);
+  write_value(unused8);
+  if (is_empty()) return bytes_written; // sketch is empty, no need to write further bytes.
 
   // Long 2
-  write(os, _total_weight);
+  const W t_weight = _total_weight;
+  write_value(t_weight);
 
   // Long 3 onwards: remaining bytes are consumed by writing the weight and the array values.
-  auto it = _sketch_array.begin();
-  while (it != _sketch_array.end()) {
-    write(os, *it);
-    ++it;
-  }
+  for (const auto& value: _sketch_array) write_value(value);
+
+  return bytes_written;
 }
 
 template<typename W, typename A>
@@ -351,40 +367,9 @@ template<typename W, typename A>
 auto count_min_sketch<W,A>::serialize(unsigned header_size_bytes) const -> vector_bytes {
   vector_bytes bytes(header_size_bytes + get_serialized_size_bytes(), 0, _allocator);
   uint8_t *ptr = bytes.data() + header_size_bytes;
-
-  // Long 0
-  const uint8_t preamble_longs = PREAMBLE_LONGS_SHORT;
-  ptr += copy_to_mem(preamble_longs, ptr);
-  const uint8_t ser_ver = SERIAL_VERSION_1;
-  ptr += copy_to_mem(ser_ver, ptr);
-  const uint8_t family_id = FAMILY_ID;
-  ptr += copy_to_mem(family_id, ptr);
-  const uint8_t flags_byte = (is_empty() ? 1 << flags::IS_EMPTY : 0);
-  ptr += copy_to_mem(flags_byte, ptr);
-  const uint32_t unused32 = NULL_32;
-  ptr += copy_to_mem(unused32, ptr);
-
-  // Long 1
-  const uint32_t nbuckets = _num_buckets;
-  const uint8_t nhashes = _num_hashes;
-  const uint16_t seed_hash(compute_seed_hash(_seed));
-  const uint8_t null_characters_8 =  NULL_8;
-  ptr += copy_to_mem(nbuckets, ptr);
-  ptr += copy_to_mem(nhashes, ptr);
-  ptr += copy_to_mem(seed_hash, ptr);
-  ptr += copy_to_mem(null_characters_8, ptr);
-  if (is_empty()) return bytes; // sketch is empty, no need to write further bytes.
-
-  // Long 2
-  const W t_weight = _total_weight;
-  ptr += copy_to_mem(t_weight, ptr);
-
-  // Long  3 onwards: remaining bytes are consumed by writing the weight and the array values.
-  auto it = _sketch_array.begin();
-  while (it != _sketch_array.end()) {
-    ptr += copy_to_mem(*it, ptr);
-    ++it;
-  }
+  serialize_to([&ptr](const void* data, size_t size) {
+    ptr += copy_to_mem(data, ptr, size);
+  });
 
   return bytes;
 }

count/test/count_min_test.cpp

@@ -300,4 +300,27 @@ TEST_CASE("CountMin sketch: bytes serialize-deserialize non-empty", "[cm_sketch]
 
 }
 
+TEST_CASE("CountMin sketch: sink serialize matches bytes", "[cm_sketch]") {
+  auto check_sink_serialize = [](const count_min_sketch<uint64_t>& c) {
+    auto bytes = c.serialize();
+    std::vector<uint8_t> sink_bytes;
+    sink_bytes.reserve(c.get_serialized_size_bytes());
+
+    const size_t bytes_written = c.serialize_to([&sink_bytes](const void* data, size_t size) {
+      const auto* begin = static_cast<const uint8_t*>(data);
+      sink_bytes.insert(sink_bytes.end(), begin, begin + size);
+    });
+
+    REQUIRE(bytes_written == bytes.size());
+    REQUIRE(sink_bytes == bytes);
+  };
+
+  count_min_sketch<uint64_t> empty(3, 32);
+  check_sink_serialize(empty);
+
+  count_min_sketch<uint64_t> non_empty(5, 64);
+  for (uint64_t i=0; i < 10; ++i) non_empty.update(i, 10 * i * i);
+  check_sink_serialize(non_empty);
+}
+
 } /* namespace datasketches */