Merge pull request #362 from jmalkin/density_serialization

Density serialization

Author: jmalkin

common/include/memory_operations.hpp

@@ -24,6 +24,7 @@
 #include <exception>
 #include <iostream>
 #include <string>
+#include <cstring>
 
 namespace datasketches {
 

density/include/density_sketch.hpp

@@ -126,6 +126,44 @@ class density_sketch {
    */
   Allocator get_allocator() const;
 
+  /**
+   * This method serializes the sketch into a given stream in a binary form
+   * @param os output stream
+   */
+  void serialize(std::ostream& os) const;
+
+  using vector_bytes = std::vector<uint8_t, typename std::allocator_traits<Allocator>::template rebind_alloc<uint8_t>>;
+
+  /**
+   * This method serializes the sketch as a vector of bytes.
+   * An optional header can be reserved in front of the sketch.
+   * It is an uninitialized space of a given size.
+   * This header is used in Datasketches PostgreSQL extension.
+   * @param header_size_bytes space to reserve in front of the sketch
+   */
+  vector_bytes serialize(unsigned header_size_bytes = 0) const;
+
+  /**
+   * This method deserializes a sketch from a given stream.
+   * @param is input stream
+   * @param kernel the kernel function to use for this sketch
+   * @param allocator the memory allocator to use with this sketch
+   * @return an instance of the sketch
+   */
+  static density_sketch deserialize(std::istream& is,
+      const Kernel& kernel=Kernel(), const Allocator& allocator = Allocator());
+
+  /**
+   * This method deserializes a sketch from a given array of bytes.
+   * @param bytes pointer to the array of bytes
+   * @param size the size of the array
+   * @param kernel the kernel function to use for this sketch
+   * @param allocator the memory allocator to use with this sketch
+   * @return an instance of the sketch
+   */
+  static density_sketch deserialize(const void* bytes, size_t size,
+      const Kernel& kernel=Kernel(), const Allocator& allocator = Allocator());
+
   /**
    * Prints a summary of the sketch.
    * @param print_levels if true include information about levels
@@ -138,6 +176,14 @@ class density_sketch {
   const_iterator end() const;
 
 private:
+  enum flags { RESERVED0, RESERVED1, IS_EMPTY };
+  static const uint8_t PREAMBLE_INTS_SHORT = 3;
+  static const uint8_t PREAMBLE_INTS_LONG = 6;
+  static const uint8_t FAMILY_ID = 19;
+  static const uint8_t SERIAL_VERSION = 1;
+  static const size_t LEVELS_ARRAY_START = 5;
+
+  Allocator allocator_;
   Kernel kernel_;
   uint16_t k_;
   uint32_t dim_;
@@ -147,6 +193,14 @@ class density_sketch {
 
   void compact();
   void compact_level(unsigned height);
+
+  static void check_k(uint16_t k);
+  static void check_serial_version(uint8_t serial_version);
+  static void check_family_id(uint8_t family_id);
+  static void check_header_validity(uint8_t preamble_ints, uint8_t flags_byte, uint8_t serial_version);
+
+  density_sketch(uint16_t k, uint32_t dim, uint32_t num_retained, uint64_t n, Levels&& levels,
+                 const Kernel& kernel = Kernel());
 };
 
 template<typename T, typename K, typename A>

density/include/density_sketch_impl.hpp

@@ -23,6 +23,7 @@
 #include <algorithm>
 #include <sstream>
 
+#include "memory_operations.hpp"
 #include "conditional_forward.hpp"
 
 namespace datasketches {
@@ -35,7 +36,22 @@ dim_(dim),
 num_retained_(0),
 n_(0),
 levels_(1, Level(allocator), allocator)
-{}
+{
+  check_k(k);
+}
+
+template<typename T, typename K, typename A>
+density_sketch<T, K, A>::density_sketch(uint16_t k, uint32_t dim, uint32_t num_retained, uint64_t n,
+                                        Levels&& levels, const K& kernel):
+kernel_(kernel),
+k_(k),
+dim_(dim),
+num_retained_(num_retained),
+n_(n),
+levels_(std::move(levels))
+{
+  check_k(k);
+}
 
 template<typename T, typename K, typename A>
 uint16_t density_sketch<T, K, A>::get_k() const {
@@ -146,6 +162,259 @@ void density_sketch<T, K, A>::compact_level(unsigned height) {
   level.clear();
 }
 
+/* Serialized sketch layout:
+ * Int  || Start Byte Addr:
+ * Addr:
+ *      ||       0        |    1   |    2   |    3   |    4   |    5   |    6   |    7   |
+ *  0   || Preamble_Ints  | SerVer | FamID  |  Flags |------- K -------|---- unused -----|
+ *
+ *      ||       8        |    9   |   10   |   11   |   12   |   13   |   14   |   15   |
+ *  2   ||------------- Num Dimensions --------------|------ Num Retained Items ---------|
+ *
+ *      ||       16       |   17   |   18   |   19   |   20   |   21   |   22   |   23   |
+ *  4   ||---------------------------Items Seen Count (N)--------------------------------|
+ *
+ * Ints 2 and 3 are omitted when the sketch is empty, meaning Num Dimensions is stored at
+ * offset 8 in that case. Otherwise, Int 5 is the start of level data, consisting of the
+ * size of the level (as a uint32 value) followed by that number of points, with 
+ * Num Dimensions per point.
+ */
+
+template<typename T, typename K, typename A>
+void density_sketch<T, K, A>::serialize(std::ostream& os) const {
+  const uint8_t preamble_ints = is_empty() ? PREAMBLE_INTS_SHORT : PREAMBLE_INTS_LONG;
+  write(os, preamble_ints);
+  const uint8_t ser_ver = SERIAL_VERSION;
+  write(os, ser_ver);
+  const uint8_t family = FAMILY_ID;
+  write(os, family);
+
+  // only empty is a valid flag
+  const uint8_t flags_byte = (is_empty() ? 1 << flags::IS_EMPTY : 0);
+  write(os, flags_byte);
+  write(os, k_);
+  const uint16_t unused = 0;
+  write(os, unused);
+  write(os, dim_);
+
+  if (is_empty())
+    return;
+
+  write(os, num_retained_);
+  write(os, n_);
+
+  // levels array -- uint32_t since a single level may be larger than k
+  size_t pt_size = sizeof(T) * dim_;
+  for (const Level& lvl : levels_) {
+    const uint32_t level_size = static_cast<uint32_t>(lvl.size());
+    write(os, level_size);
+    for (const Vector& pt : lvl) {
+      write(os, pt.data(), pt_size);
+    }
+  }
+}
+
+template<typename T, typename K, typename A>
+auto density_sketch<T, K, A>::serialize(unsigned header_size_bytes) const -> vector_bytes {
+  const uint8_t preamble_ints = (is_empty() ? PREAMBLE_INTS_SHORT : PREAMBLE_INTS_LONG);
+  
+  // pre-compute size
+  size_t size = header_size_bytes + preamble_ints * sizeof(uint32_t);
+  if (!is_empty())
+    for (const Level& lvl : levels_)
+      size += sizeof(uint32_t) + (lvl.size() * dim_ * sizeof(T));
+
+  vector_bytes bytes(size, 0, levels_.get_allocator());
+  uint8_t* ptr = bytes.data() + header_size_bytes;
+  const uint8_t* end_ptr = ptr + size;
+  
+  ptr += copy_to_mem(preamble_ints, ptr);
+  const uint8_t ser_ver = SERIAL_VERSION;
+  ptr += copy_to_mem(ser_ver, ptr);
+  const uint8_t family = FAMILY_ID;
+  ptr += copy_to_mem(family, ptr);
+
+  // empty is the only valid flat
+  const uint8_t flags_byte = (is_empty() ? 1 << flags::IS_EMPTY : 0);
+  ptr += copy_to_mem(flags_byte, ptr);
+  ptr += copy_to_mem(k_, ptr);
+  ptr += sizeof(uint16_t); // 2 unused bytes
+  ptr += copy_to_mem(dim_, ptr);
+  
+  if (is_empty())
+    return bytes;
+
+  ptr += copy_to_mem(num_retained_, ptr);
+  ptr += copy_to_mem(n_, ptr);
+
+  // levels array -- uint32_t since a single level may be larger than k
+  size_t pt_size = sizeof(T) * dim_;
+  for (const Level& lvl : levels_) {
+    ptr += copy_to_mem(static_cast<uint32_t>(lvl.size()), ptr);
+    for (const Vector& pt : lvl) {
+      ptr += copy_to_mem(pt.data(), ptr, pt_size);
+    }
+  }
+
+  if (ptr != end_ptr)
+    throw std::runtime_error("Actual output size does not equal expected output size");
+
+  return bytes;
+}
+
+template<typename T, typename K, typename A>
+density_sketch<T, K, A> density_sketch<T, K, A>::deserialize(std::istream& is, const K& kernel, const A& allocator) {
+  const auto preamble_ints = read<uint8_t>(is);
+  const auto serial_version = read<uint8_t>(is);
+  const auto family_id = read<uint8_t>(is);
+  const auto flags_byte = read<uint8_t>(is);
+  const auto k = read<uint16_t>(is);
+  read<uint16_t>(is); // unused
+  const auto dim = read<uint32_t>(is);
+
+  check_k(k); // do we have constraints?
+  check_serial_version(serial_version); // a little redundant with the header check
+  check_family_id(family_id);
+  check_header_validity(preamble_ints, flags_byte, serial_version);
+
+  if (!is.good()) throw std::runtime_error("error reading from std::istream");
+  const bool is_empty = (flags_byte & (1 << flags::IS_EMPTY)) > 0;
+  if (is_empty) {
+    return density_sketch(k, dim, kernel, allocator);
+  }
+
+  const auto num_retained = read<uint32_t>(is);
+  const auto n = read<uint64_t>(is);
+
+  // levels arrays
+  size_t pt_size = sizeof(T) * dim;
+  Levels levels(allocator);
+  int64_t num_to_read = num_retained; // num_retrained is uint32_t so this allows error checking
+  while (num_to_read > 0) {
+    const auto level_size = read<uint32_t>(is);
+    Level lvl(allocator);
+    lvl.reserve(level_size);
+    for (uint32_t i = 0; i < level_size; ++i) {
+      Vector pt(dim, 0, allocator);
+      read(is, pt.data(), pt_size);
+      lvl.push_back(pt);
+    }
+    levels.push_back(lvl);
+    num_to_read -= lvl.size();
+  }
+
+  if (num_to_read != 0)
+    throw std::runtime_error("Error deserializing sketch: Incorrect number of items read");
+  if (!is.good()) throw std::runtime_error("error reading from std::istream");
+
+  return density_sketch(k, dim, num_retained, n, std::move(levels), kernel);
+}
+
+template<typename T, typename K, typename A>
+density_sketch<T, K, A> density_sketch<T, K, A>::deserialize(const void* bytes, size_t size, const K& kernel, const A& allocator) {
+  ensure_minimum_memory(size, PREAMBLE_INTS_SHORT * sizeof(uint32_t));
+  const char* ptr = static_cast<const char*>(bytes);
+  const char* end_ptr = static_cast<const char*>(bytes) + size;
+  uint8_t preamble_ints;
+  ptr += copy_from_mem(ptr, preamble_ints);
+  uint8_t serial_version;
+  ptr += copy_from_mem(ptr, serial_version);
+  uint8_t family_id;
+  ptr += copy_from_mem(ptr, family_id);
+  uint8_t flags_byte;
+  ptr += copy_from_mem(ptr, flags_byte);
+  uint16_t k;
+  ptr += copy_from_mem(ptr, k);
+  uint16_t unused;
+  ptr += copy_from_mem(ptr, unused);
+  uint32_t dim;
+  ptr += copy_from_mem(ptr, dim);
+
+  check_k(k);
+  check_serial_version(serial_version); // a little redundant with the header check
+  check_family_id(family_id);
+  check_header_validity(preamble_ints, flags_byte, serial_version);
+
+  const bool is_empty = (flags_byte & (1 << flags::IS_EMPTY)) > 0;
+  if (is_empty) {
+    return density_sketch(k, dim, kernel, allocator);
+  }
+
+  ensure_minimum_memory(size, PREAMBLE_INTS_LONG * sizeof(uint32_t));
+  uint32_t num_retained;
+  ptr += copy_from_mem(ptr, num_retained);
+  uint64_t n;
+  ptr += copy_from_mem(ptr, n);
+
+  // Predicting the number of levels seems hard so determining the exact remaining
+  // size is also hard. But we need at least num_retained * dim * sizeof(T)
+  // bytes for the points so we can check that.
+  size_t pt_size = sizeof(T) * dim;
+  ensure_minimum_memory(end_ptr - ptr, num_retained * pt_size);
+
+  // levels arrays
+  Levels levels(allocator);
+  int64_t num_to_read = num_retained; // num_retrained is uint32_t so this allows error checking
+  while (num_to_read > 0) {
+    uint32_t level_size;
+    ptr += copy_from_mem(ptr, level_size);
+    Level lvl(allocator);
+    lvl.reserve(level_size);
+    for (uint32_t i = 0; i < level_size; ++i) {
+      Vector pt(dim, 0, allocator);
+      ptr += copy_from_mem(ptr, pt.data(), pt_size);
+      lvl.push_back(pt);
+    }
+    levels.push_back(lvl);
+    num_to_read -= lvl.size();
+  }
+
+  if (num_to_read != 0)
+    throw std::runtime_error("Error deserializing sketch: Incorrect number of items read");
+  if (ptr > end_ptr) throw std::runtime_error("Error deserializing sketch: Read beyond provided memory");
+
+  return density_sketch(k, dim, num_retained, n, std::move(levels), kernel);
+}
+
+template<typename T, typename K, typename A>
+void density_sketch<T, K, A>::check_k(uint16_t k) {
+  if (k < 2)
+    throw std::invalid_argument("k must be > 1. Found: " + std::to_string(k));
+}
+
+template<typename T, typename K, typename A>
+void density_sketch<T, K, A>::check_serial_version(uint8_t serial_version) {
+  if (serial_version == SERIAL_VERSION)
+    return;
+  else
+    throw std::invalid_argument("Possible corruption. Unrecognized serialization version: " + std::to_string(serial_version));
+}
+
+template<typename T, typename K, typename A>
+void density_sketch<T, K, A>::check_family_id(uint8_t family_id) {
+  if (family_id == FAMILY_ID)
+    return;
+  else
+    throw std::invalid_argument("Possible corruption. Family id does not indicate density sketch: " + std::to_string(family_id));
+}
+
+template<typename T, typename K, typename A>
+void density_sketch<T, K, A>::check_header_validity(uint8_t preamble_ints, uint8_t flags_byte, uint8_t serial_version) {
+  const bool empty = (flags_byte & (1 << flags::IS_EMPTY)) > 0;
+
+  if ((empty && preamble_ints == PREAMBLE_INTS_SHORT)
+      || (!empty && preamble_ints == PREAMBLE_INTS_LONG))
+      return;
+  else {
+    std::ostringstream os;
+    os << "Possible sketch corruption. Inconsistent state: "
+       << "preamble_ints = " << preamble_ints
+       << ", empty = " << (empty ? "true" : "false")
+       << ", serialization_version = " << serial_version;
+    throw std::invalid_argument(os.str());
+  }
+}
+
 template<typename T, typename K, typename A>
 string<A> density_sketch<T, K, A>::to_string(bool print_levels, bool print_items) const {
   // Using a temporary stream for implementation here does not comply with AllocatorAwareContainer requirements.