From 15ebe5e4e5d0ea5fc58c7dae9d1d25a105797af5 Mon Sep 17 00:00:00 2001
From: Gleb Natapov <gleb@scylladb.com>
Date: Thu, 14 Apr 2016 17:02:05 +0300
Subject: [PATCH] query: add calculate_row_count function to query::result

---
 query-result.hh |  2 ++
 query.cc        | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/query-result.hh b/query-result.hh
index 07c6f1de33..0935151a34 100644
--- a/query-result.hh
+++ b/query-result.hh
@@ -117,6 +117,8 @@ public:
         return _digest;
     }
 
+    uint32_t calculate_row_count(const query::partition_slice&);
+
     struct printer {
         schema_ptr s;
         const query::partition_slice& slice;
diff --git a/query.cc b/query.cc
index bbbbc93335..e90631b6e7 100644
--- a/query.cc
+++ b/query.cc
@@ -163,6 +163,40 @@ std::ostream& operator<<(std::ostream& os, const query::result::printer& p) {
     return os;
 }
 
+uint32_t result::calculate_row_count(const query::partition_slice& slice) {
+    struct {
+        uint32_t total_count = 0;
+        uint32_t current_partition_count = 0;
+        void accept_new_partition(const partition_key& key, uint32_t row_count) {
+            accept_new_partition(row_count);
+        }
+        void accept_new_partition(uint32_t row_count) {
+            total_count += row_count;
+            current_partition_count = row_count;
+        }
+        void accept_new_row(...) {}
+        void accept_partition_end(const query::result_row_view& static_row) {
+            if (current_partition_count == 0) {
+                total_count++;
+            }
+        }
+    } counter;
+
+    bytes_view v;
+
+    if (buf().is_linearized()) {
+        v = buf().view();
+    } else {
+        // FIXME: make result_view::consume() work on fragments to avoid linearization.
+        bytes_ostream w(buf());
+        v = w.linearize();
+    }
+
+    query::result_view view(v);
+    view.consume(slice, counter);
+    return counter.total_count;
+}
+
 result::result()
     : result([] {
         bytes_ostream out;