Introduced const_iterator for tdigest

Author: MaheshGPai

tdigest/include/tdigest.hpp

@@ -143,11 +143,6 @@ class tdigest {
    */
   uint64_t get_total_weight() const;
 
-  /**
-   * @return centroids
-   */
-  vector_centroid get_centroids() const;
-
   /**
    * Returns an instance of the allocator for this t-Digest.
    * @return allocator
@@ -262,6 +257,21 @@ class tdigest {
    */
   static tdigest deserialize(const void* bytes, size_t size, const Allocator& allocator = Allocator());
 
+  class const_iterator;
+
+  /**
+   * Iterator pointing to the first centroid in the sketch.
+   * If the sketch is empty, the returned iterator must not be dereferenced or incremented.
+   * @return iterator pointing to the first centroid in the sketch
+   */
+  const_iterator begin() const;
+
+  /**
+   * Iterator pointing to the past-the-end centroid in the sketch.
+   * It does not point to any centroid, and must not be dereferenced or incremented.
+   * @return iterator pointing to the past-the-end centroid in the sketch
+   */
+  const_iterator end() const;
 private:
   bool reverse_merge_;
   uint16_t k_;
@@ -302,6 +312,28 @@ class tdigest {
   static inline void check_split_points(const T* values, uint32_t size);
 };
 
+template<typename T, typename A>
+class tdigest<T, A>::const_iterator {
+public:
+  using iterator_category = std::input_iterator_tag;
+  using value_type = std::pair<const T&, const W>;
+  using difference_type = void;
+  using pointer = const return_value_holder<value_type>;
+  using reference = const value_type;
+
+  const_iterator(const tdigest<T, A> &tdigest_, bool is_end);
+
+  const_iterator& operator++();
+  const_iterator& operator++(int);
+  bool operator==(const const_iterator& other) const;
+  bool operator!=(const const_iterator& other) const;
+  reference operator*() const;
+  pointer operator->() const;
+private:
+  friend class tdigest<T, A>;
+  uint32_t index_;
+  vector_centroid centroids_;
+};
 } /* namespace datasketches */
 
 #include "tdigest_impl.hpp"

tdigest/include/tdigest_impl.hpp

@@ -85,11 +85,6 @@ uint64_t tdigest<T, A>::get_total_weight() const {
   return centroids_weight_ + buffer_.size();
 }
 
-template<typename T, typename A>
-auto tdigest<T, A>::get_centroids() const -> vector_centroid{
-  return centroids_;
-}
-
 template<typename T, typename A>
 A tdigest<T, A>::get_allocator() const {
   return buffer_.get_allocator();
@@ -632,6 +627,65 @@ void tdigest<T, A>::check_split_points(const T* values, uint32_t size) {
   }
 }
 
+template <typename T, typename A>
+typename tdigest<T, A>::const_iterator tdigest<T, A>::begin() const {
+  return tdigest<T, A>::const_iterator(*this, false);
+}
+
+template <typename T, typename A>
+  typename tdigest<T, A>::const_iterator tdigest<T, A>::end() const {
+  return tdigest::const_iterator(*this, true);
+}
+
+template<typename T, typename A>
+tdigest<T, A>::const_iterator::const_iterator(const tdigest& tdigest_, const bool is_end):
+  centroids_()
+{
+  // Create a copy of the tdigest to generate the centroids after processing the buffered values
+  tdigest tmp(tdigest_);
+  tmp.compress();
+  centroids_.insert(centroids_.end(), tmp.centroids_.begin(), tmp.centroids_.end());
+
+  if (is_end) {
+    index_ = centroids_.size();
+  } else {
+    index_ = 0;
+  }
+}
+
+template<typename T, typename A>
+typename tdigest<T, A>::const_iterator& tdigest<T, A>::const_iterator::operator++() {
+  ++index_;
+  return *this;
+}
+
+template<typename T, typename A>
+typename tdigest<T, A>::const_iterator& tdigest<T, A>::const_iterator::operator++(int) {
+  const_iterator tmp(*this);
+  operator++();
+  return tmp;
+}
+
+template<typename T, typename A>
+bool tdigest<T, A>::const_iterator::operator==(const const_iterator& other) const {
+  return index_ == other.index_;
+}
+
+template<typename T, typename A>
+bool tdigest<T, A>::const_iterator::operator!=(const const_iterator& other) const {
+  return !operator==(other);
+}
+
+template<typename T, typename A>
+auto tdigest<T, A>::const_iterator::operator*() const -> reference {
+  return value_type(centroids_[index_].get_mean(), centroids_[index_].get_weight());
+}
+
+template<typename T, typename A>
+auto tdigest<T, A>::const_iterator::operator->() const -> pointer {
+  return **this;
+}
+
 } /* namespace datasketches */
 
 #endif // _TDIGEST_IMPL_HPP_

tdigest/test/tdigest_test.cpp

@@ -453,4 +453,18 @@ TEST_CASE("deserialize from reference implementation bytes float", "[tdigest]")
   REQUIRE(td.get_rank(n) == 1);
 }
 
+TEST_CASE("iterate centroids", "[tdigest]") {
+  tdigest_double td(100);
+  for (int i = 0; i < 10; i++) {
+    td.update(i);
+  }
+
+  auto centroid_count = 0;
+  for (const auto &centroid: td) {
+    centroid_count++;
+  }
+  // Ensure that centroids are retrieved for a case where there is buffered values
+  REQUIRE(centroid_count == 10);
+}
+
 } /* namespace datasketches */