query: add calculate_row_count function to query::result

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;

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;