Merge branch 'sample_row_keys' into for_google

Author: marete

google/cloud/bigtable/emulator/column_family.h

@@ -30,6 +30,7 @@
 #include <google/bigtable/v2/types.pb.h>
 #include <absl/strings/str_format.h>
 #include <chrono>
+#include <cstddef>
 #include <cstdint>
 #include <map>
 #include <memory>

google/cloud/bigtable/emulator/server.cc

@@ -27,6 +27,7 @@
 #include <grpcpp/server_builder.h>
 #include <grpcpp/support/client_callback.h>
 #include <grpcpp/support/status.h>
+#include <cstddef>
 #include <cstdint>
 
 namespace google {
@@ -56,10 +57,16 @@ class EmulatorService final : public btproto::Bigtable::Service {
 
   grpc::Status SampleRowKeys(
       grpc::ServerContext* /* context */,
-      btproto::SampleRowKeysRequest const* /* request */,
-      grpc::ServerWriter<btproto::SampleRowKeysResponse>* /* writer */)
-      override {
-    return grpc::Status::OK;
+      btproto::SampleRowKeysRequest const* request,
+      grpc::ServerWriter<btproto::SampleRowKeysResponse>* writer) override {
+    auto maybe_table = cluster_->FindTable(request->table_name());
+    if (!maybe_table) {
+      return ToGrpcStatus(maybe_table.status());
+    }
+
+    auto& table = maybe_table.value();
+
+    return ToGrpcStatus(table->SampleRowKeys(0.0001, writer));
   }
 
   grpc::Status MutateRow(grpc::ServerContext* /* context */,

google/cloud/bigtable/emulator/table.cc

@@ -28,12 +28,18 @@
 #include <absl/strings/match.h>
 #include <absl/strings/str_format.h>
 #include <absl/types/optional.h>
+#include <grpcpp/support/sync_stream.h>
 #include <re2/re2.h>
 #include <chrono>
+#include <climits>
+#include <cstddef>
 #include <cstdint>
 #include <cstdlib>
+#include <exception>
 #include <memory>
 #include <mutex>
+#include <random>
+#include <ratio>
 #include <optional>
 #include <string>
 #include <type_traits>
@@ -593,6 +599,119 @@ bool Table::IsDeleteProtectedNoLock() const {
   return schema_.deletion_protection();
 }
 
+Status Table::SampleRowKeys(
+    double pass_probability,
+    grpc::ServerWriter<google::bigtable::v2::SampleRowKeysResponse>* writer) {
+  std::lock_guard<std::mutex> lock(mu_);
+
+  // First, stream all rows and cells and compute the offsets.
+  auto all_rows_set = std::make_shared<StringRangeSet>(StringRangeSet::All());
+  auto maybe_all_rows_steam = CreateCellStream(all_rows_set, absl::nullopt);
+  if (!maybe_all_rows_steam) {
+    return maybe_all_rows_steam.status();
+  }
+
+  auto& stream = *maybe_all_rows_steam;
+
+  std::map<std::string, std::size_t> row_offset_map;
+  size_t row_offset = 0;
+
+  std::string current_row_key;
+  bool first_row = true;
+
+  std::map<std::string, std::size_t> column_family_size_map;
+  std::map<std::string, std::size_t> column_qualifier_size_map;
+  size_t timestamp_total_row_size = 0;
+  size_t value_total_row_size = 0;
+
+  for (; stream; ++stream) {
+    auto row_key = stream->row_key();
+
+    if ((row_key != current_row_key) || first_row) {
+      row_offset += current_row_key.size();
+
+      for (auto const& cf : column_family_size_map) {
+        row_offset += cf.second;
+      }
+
+      for (auto const& cq : column_qualifier_size_map) {
+        row_offset += cq.second;
+      }
+
+      row_offset += timestamp_total_row_size;
+      row_offset += value_total_row_size;
+
+      // The rows before this (row_key) have this size in total.
+      row_offset_map[row_key] = row_offset;
+
+      current_row_key = row_key;
+
+      first_row = false;
+
+      column_family_size_map.clear();
+      column_qualifier_size_map.clear();
+      timestamp_total_row_size = 0;
+      value_total_row_size = 0;
+    }
+
+    column_family_size_map.emplace(stream->column_family(),
+                                   stream->column_family().size());
+    column_qualifier_size_map.emplace(stream->column_qualifier(),
+                                      stream->column_qualifier().size());
+    timestamp_total_row_size += sizeof(stream->timestamp());
+    value_total_row_size += stream->value().size();
+  }
+
+  google::bigtable::v2::RowFilter sample_filter;
+  sample_filter.set_row_sample_filter(pass_probability);
+
+  auto maybe_stream = CreateCellStream(all_rows_set, sample_filter);
+  if (!maybe_stream) {
+    return maybe_stream.status();
+  }
+
+  auto& sampled_stream = *maybe_stream;
+
+  bool wrote_a_sample = false;
+
+  for (; sampled_stream; ++sampled_stream) {
+    google::bigtable::v2::SampleRowKeysResponse resp;
+    resp.set_row_key(sampled_stream->row_key());
+    resp.set_offset_bytes(row_offset_map[sampled_stream->row_key()]);
+
+    writer->Write(std::move(resp));
+
+    wrote_a_sample = true;
+  }
+
+  // Cloud bigtable client tests expect that, if they populated the
+  // table with at least one row, then at least one row sampele is
+  // returned.
+  //
+  // In such a case, return the last row key.
+  if (!wrote_a_sample && !row_offset_map.empty()) {
+    auto it = std::prev(row_offset_map.end());
+
+    google::bigtable::v2::SampleRowKeysResponse resp;
+    resp.set_row_key(it->first);
+    resp.set_offset_bytes(it->second);
+  }
+
+  // Client code expects the last response to be an empty row key
+  // and moreover it also expects the offset for the last response
+  // to be more than every other offset.
+  google::bigtable::v2::SampleRowKeysResponse resp;
+  resp.set_row_key("");
+  // Client test code expects offset_bytes to be strictly
+  // increasing.
+  resp.set_offset_bytes(row_offset + 1);
+  auto opts = grpc::WriteOptions();
+  opts.set_last_message();
+  writer->WriteLast(std::move(resp), opts);
+
+  return Status();
+}
+
 Status Table::DropRowRange(
     ::google::bigtable::admin::v2::DropRowRangeRequest const& request) {
   std::lock_guard<std::mutex> lock(mu_);

google/cloud/bigtable/emulator/table.h

@@ -32,11 +32,13 @@
 #include <google/protobuf/util/time_util.h>
 #include <absl/types/optional.h>
 #include <chrono>
+#include <cstddef>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <stack>
 #include <string>
+#include <utility>
 
 namespace google {
 namespace cloud {
@@ -93,6 +95,10 @@ class Table : public std::enable_shared_from_this<Table> {
     return column_families_.find(column_family);
   }
 
+  Status SampleRowKeys(
+      double pass_probability,
+      grpc::ServerWriter<google::bigtable::v2::SampleRowKeysResponse>* writer);
+
   std::shared_ptr<Table> get() { return shared_from_this(); }
 
   Status DropRowRange(