[tf.data / Bigtable] Fix server-set timestamp mutations.

Use SetCell with no timestamp arg when using server-set timestamps instead of using std::chrono::milliseconds timestamp(-1), because this results in a timestamp_micros value of -1000 instead of -1, which causes the server to become unhappy.

PiperOrigin-RevId: 203165785

2 files changed, 14 insertions, 6 deletions

diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc b/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc
index 94c0dec812..f43b44f2cb 100644
--- a/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc
+++ b/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc
@@ -212,7 +212,6 @@ class ToBigtableOp : public AsyncOpKernel {
       OP_REQUIRES_ASYNC(ctx, timestamp_int >= -1,
                         errors::InvalidArgument("timestamp must be >= -1"),
                         done);
-      std::chrono::milliseconds timestamp(timestamp_int);
 
       BigtableTableResource* resource;
       OP_REQUIRES_OK_ASYNC(
@@ -233,7 +232,7 @@ class ToBigtableOp : public AsyncOpKernel {
             OP_REQUIRES_OK_ASYNC(
                 ctx,
                 CreateMutation(std::move(components), column_families, columns,
-                               timestamp, &mutation),
+                               timestamp_int, &mutation),
                 done);
           }
           components.clear();
@@ -282,7 +281,7 @@ class ToBigtableOp : public AsyncOpKernel {
 
   Status CreateMutation(
       std::vector<Tensor> tensors, const std::vector<string>& column_families,
-      const std::vector<string>& columns, std::chrono::milliseconds timestamp,
+      const std::vector<string>& columns, int64 timestamp_int,
       ::google::cloud::bigtable::BulkMutation* bulk_mutation) {
     if (tensors.size() != column_families.size() + 1) {
       return errors::InvalidArgument(
@@ -290,13 +289,20 @@ class ToBigtableOp : public AsyncOpKernel {
     }
     ::google::cloud::bigtable::SingleRowMutation mutation(
         std::move(tensors[0].scalar<string>()()));
+    std::chrono::milliseconds timestamp(timestamp_int);
     for (size_t i = 1; i < tensors.size(); ++i) {
       if (!TensorShapeUtils::IsScalar(tensors[i].shape())) {
         return errors::Internal("Output tensor ", i, " was not a scalar");
       }
-      mutation.emplace_back(::google::cloud::bigtable::SetCell(
-          column_families[i - 1], columns[i - 1], timestamp,
-          std::move(tensors[i].scalar<string>()())));
+      if (timestamp_int == -1) {
+        mutation.emplace_back(::google::cloud::bigtable::SetCell(
+            column_families[i - 1], columns[i - 1],
+            std::move(tensors[i].scalar<string>()())));
+      } else {
+        mutation.emplace_back(::google::cloud::bigtable::SetCell(
+            column_families[i - 1], columns[i - 1], timestamp,
+            std::move(tensors[i].scalar<string>()())));
+      }
     }
     bulk_mutation->emplace_back(std::move(mutation));
     return Status::OK();
diff --git a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc
index 0f107f169c..c164682508 100644
--- a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc
+++ b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc
@@ -28,6 +28,8 @@ namespace {
 void UpdateRow(const ::google::bigtable::v2::Mutation& mut,
                std::map<string, string>* row) {
   if (mut.has_set_cell()) {
+    CHECK(mut.set_cell().timestamp_micros() >= -1)
+        << "Timestamp_micros: " << mut.set_cell().timestamp_micros();
     auto col =
         strings::Printf("%s:%s", mut.set_cell().family_name().c_str(),
                         string(mut.set_cell().column_qualifier()).c_str());