Merge pull request #509 from cv4g/countmin-serialize-to

Introduce `count_min_sketch::serialize_to`

Author: leerho

benchmarks/CMakeLists.txt

@@ -48,4 +48,5 @@ set_target_properties(count_min_sketch_benchmark PROPERTIES
 target_sources(count_min_sketch_benchmark
   PRIVATE
     benchmark_count_min_sketch.cpp
+    benchmark_count_min_sketch_serialization.cpp
 )

benchmarks/benchmark_count_min_sketch_serialization.cpp

@@ -0,0 +1,161 @@
+/*
+ * 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.
+ */
+
+#include <benchmark/benchmark.h>
+#include <count_min.hpp>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <ostream>
+#include <streambuf>
+#include <vector>
+
+namespace
+{
+
+using Sketch = datasketches::count_min_sketch<uint64_t>;
+
+// Roughly 99.9% confidence. Serialization cost then scales with the bucket count.
+constexpr uint8_t NUM_HASHES = 7;
+
+Sketch makeSketch(uint32_t num_buckets)
+{
+    Sketch sketch(NUM_HASHES, num_buckets);
+    for (uint64_t i = 0; i < NUM_HASHES; ++i)
+        sketch.update(i, i + 1);
+    return sketch;
+}
+
+class fixed_buffer_streambuf final: public std::streambuf
+{
+public:
+    explicit fixed_buffer_streambuf(std::vector<uint8_t>& buffer): buffer_(buffer), position_(0) {}
+
+    void reset()
+    {
+        position_ = 0;
+    }
+
+    size_t bytes_written() const
+    {
+        return position_;
+    }
+
+protected:
+    std::streamsize xsputn(const char* data, std::streamsize size) override
+    {
+        const size_t requested = static_cast<size_t>(size);
+        const size_t available = buffer_.size() - position_;
+        const size_t bytes_to_write = std::min(requested, available);
+        if (bytes_to_write > 0) {
+            std::memcpy(buffer_.data() + position_, data, bytes_to_write);
+            position_ += bytes_to_write;
+        }
+        return static_cast<std::streamsize>(bytes_to_write);
+    }
+
+    int_type overflow(int_type ch) override
+    {
+        if (traits_type::eq_int_type(ch, traits_type::eof()))
+            return traits_type::not_eof(ch);
+        if (position_ == buffer_.size())
+            return traits_type::eof();
+        buffer_[position_++] = static_cast<uint8_t>(traits_type::to_char_type(ch));
+        return ch;
+    }
+
+private:
+    std::vector<uint8_t>& buffer_;
+    size_t position_;
+};
+
+void BM_CountMinGetSerializedSizeBytes(benchmark::State& state)
+{
+    const auto sketch = makeSketch(static_cast<uint32_t>(state.range(0)));
+
+    for (auto _ : state) {
+        const auto size = sketch.get_serialized_size_bytes();
+        benchmark::DoNotOptimize(size);
+    }
+
+    state.SetItemsProcessed(state.iterations());
+}
+
+void BM_CountMinSerializeVector(benchmark::State& state)
+{
+    const auto sketch = makeSketch(static_cast<uint32_t>(state.range(0)));
+    const auto serialized_size = sketch.get_serialized_size_bytes();
+
+    for (auto _ : state) {
+        auto bytes = sketch.serialize();
+        benchmark::DoNotOptimize(bytes.data());
+        benchmark::DoNotOptimize(bytes.size());
+        benchmark::ClobberMemory();
+    }
+
+    state.SetBytesProcessed(state.iterations() * static_cast<int64_t>(serialized_size));
+}
+
+void BM_CountMinSerializeOStream(benchmark::State& state)
+{
+    const auto sketch = makeSketch(static_cast<uint32_t>(state.range(0)));
+    const auto serialized_size = sketch.get_serialized_size_bytes();
+    std::vector<uint8_t> bytes(serialized_size);
+    fixed_buffer_streambuf stream_buffer(bytes);
+    std::ostream os(&stream_buffer);
+
+    for (auto _ : state) {
+        stream_buffer.reset();
+        os.clear();
+        sketch.serialize(os);
+        benchmark::DoNotOptimize(stream_buffer.bytes_written());
+        benchmark::ClobberMemory();
+    }
+
+    state.SetBytesProcessed(state.iterations() * static_cast<int64_t>(serialized_size));
+}
+
+void BM_CountMinSerializeToSink(benchmark::State& state)
+{
+    const auto sketch = makeSketch(static_cast<uint32_t>(state.range(0)));
+    const auto serialized_size = sketch.get_serialized_size_bytes();
+    std::vector<uint8_t> bytes(serialized_size);
+
+    for (auto _ : state) {
+        uint8_t* ptr = bytes.data();
+        const auto bytes_written = sketch.serialize_to([&ptr](const void* data, size_t size) {
+            std::memcpy(ptr, data, size);
+            ptr += size;
+        });
+        benchmark::DoNotOptimize(ptr);
+        benchmark::DoNotOptimize(bytes_written);
+        benchmark::ClobberMemory();
+    }
+
+    state.SetBytesProcessed(state.iterations() * static_cast<int64_t>(serialized_size));
+}
+
+}
+
+BENCHMARK(BM_CountMinGetSerializedSizeBytes)->RangeMultiplier(8)->Range(1024, 65536);
+BENCHMARK(BM_CountMinSerializeVector)->RangeMultiplier(8)->Range(1024, 65536);
+BENCHMARK(BM_CountMinSerializeOStream)->RangeMultiplier(8)->Range(1024, 65536);
+BENCHMARK(BM_CountMinSerializeToSink)->RangeMultiplier(8)->Range(1024, 65536);

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;
@@ -326,6 +329,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 Sink>
+  size_t serialize_to(Sink&& sink) 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

@@ -262,39 +262,58 @@ 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 Sink, typename T>
+static inline void write_count_min_value(Sink& sink, size_t& bytes_written, const T& value) {
+  sink(&value, sizeof(value));
+  bytes_written += sizeof(value);
+}
+
+template<typename W, typename A>
+template<typename Sink>
+size_t count_min_sketch<W,A>::serialize_to(Sink&& sink) const {
+  size_t bytes_written = 0;
+
   // 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_count_min_value(sink, bytes_written, preamble_longs);
+  write_count_min_value(sink, bytes_written, ser_ver);
+  write_count_min_value(sink, bytes_written, family_id);
+  write_count_min_value(sink, bytes_written, flags_byte);
+  write_count_min_value(sink, bytes_written, 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_count_min_value(sink, bytes_written, nbuckets);
+  write_count_min_value(sink, bytes_written, nhashes);
+  write_count_min_value(sink, bytes_written, seed_hash);
+  write_count_min_value(sink, bytes_written, 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_count_min_value(sink, bytes_written, 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;
+  const size_t sketch_array_bytes = sizeof(W) * _sketch_array.size();
+  if (sketch_array_bytes > 0) {
+    sink(_sketch_array.data(), sketch_array_bytes);
+    bytes_written += sketch_array_bytes;
   }
+
+  return bytes_written;
 }
 
 template<typename W, typename A>
@@ -345,40 +364,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,40 @@ TEST_CASE("CountMin sketch: bytes serialize-deserialize non-empty", "[cm_sketch]
 
 }
 
+TEST_CASE("CountMin sketch: sink serialize-deserialize round trip", "[cm_sketch]") {
+  auto check_sink_serialize = [](const count_min_sketch<uint64_t>& c) {
+    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 == c.get_serialized_size_bytes());
+    REQUIRE(sink_bytes.size() == bytes_written);
+
+    auto d = count_min_sketch<uint64_t>::deserialize(sink_bytes.data(), sink_bytes.size(), c.get_seed());
+    REQUIRE(c.get_num_hashes() == d.get_num_hashes());
+    REQUIRE(c.get_num_buckets() == d.get_num_buckets());
+    REQUIRE(c.get_seed() == d.get_seed());
+    REQUIRE(c.get_total_weight() == d.get_total_weight());
+
+    auto c_it = c.begin();
+    auto d_it = d.begin();
+    while (c_it != c.end()) {
+      REQUIRE(*c_it == *d_it);
+      ++c_it;
+      ++d_it;
+    }
+  };
+
+  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 */