Serialize count_min_sketch sketch array in a single write call

Author: cv4g

count/include/count_min.hpp

@@ -21,6 +21,8 @@
 #define COUNT_MIN_HPP_
 
 #include <iterator>
+#include <type_traits>
+#include <vector>
 #include "common_defs.hpp"
 
 namespace datasketches {
@@ -36,6 +38,7 @@ template <typename W,
           typename Allocator = std::allocator<W>>
 class count_min_sketch{
   static_assert(std::is_arithmetic<W>::value, "Arithmetic type expected");
+  static_assert(!std::is_same<W, bool>::value, "Boolean weight type is not supported");
 public:
   using allocator_type = Allocator;
   using const_iterator = typename std::vector<W, Allocator>::const_iterator;

count/include/count_min_impl.hpp

@@ -313,7 +313,11 @@ size_t count_min_sketch<W,A>::serialize_to(WriteBytes&& write_bytes) const {
   write_count_min_value(write_bytes, bytes_written, t_weight);
 
   // Long 3 onwards: remaining bytes are consumed by writing the weight and the array values.
-  for (const auto& value: _sketch_array) write_count_min_value(write_bytes, bytes_written, value);
+  const size_t sketch_array_bytes = sizeof(W) * _sketch_array.size();
+  if (sketch_array_bytes > 0) {
+    write_bytes(_sketch_array.data(), sketch_array_bytes);
+    bytes_written += sketch_array_bytes;
+  }
 
   return bytes_written;
 }

count/test/count_min_test.cpp

@@ -301,26 +301,32 @@ 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 check_sink_serialize = [](const count_min_sketch<uint64_t>& c, size_t expected_fragments, size_t expected_last_fragment_size = 0) {
     auto bytes = c.serialize();
     std::vector<uint8_t> sink_bytes;
     sink_bytes.reserve(c.get_serialized_size_bytes());
+    std::vector<size_t> fragment_sizes;
 
-    const size_t bytes_written = c.serialize_to([&sink_bytes](const void* data, size_t size) {
+    const size_t bytes_written = c.serialize_to([&sink_bytes, &fragment_sizes](const void* data, size_t size) {
       const auto* begin = static_cast<const uint8_t*>(data);
       sink_bytes.insert(sink_bytes.end(), begin, begin + size);
+      fragment_sizes.push_back(size);
     });
 
     REQUIRE(bytes_written == bytes.size());
     REQUIRE(sink_bytes == bytes);
+    REQUIRE(fragment_sizes.size() == expected_fragments);
+    if (expected_last_fragment_size > 0) {
+      REQUIRE(fragment_sizes.back() == expected_last_fragment_size);
+    }
   };
 
   count_min_sketch<uint64_t> empty(3, 32);
-  check_sink_serialize(empty);
+  check_sink_serialize(empty, 9);
 
   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);
+  check_sink_serialize(non_empty, 11, sizeof(uint64_t) * non_empty.get_num_hashes() * non_empty.get_num_buckets());
 }
 
 } /* namespace datasketches */