Add GetEntityForUpdate to optimistic and WriteCommitted pessimistic transactions (#12668)

Summary: Pull Request resolved: https://github.com/facebook/rocksdb/pull/12668 The patch adds a new `GetEntityForUpdate` API to optimistic and WriteCommitted pessimistic transactions, which provides transactional wide-column point lookup functionality with concurrency control. For WriteCommitted transactions, user-defined timestamps are also supported similarly to the `GetForUpdate` API. Reviewed By: jaykorean Differential Revision: D57458304 fbshipit-source-id: 7eadbac531ca5446353e494abbd0635d63f62d24
2026-10-23 08:09:02 +03:00 · 2024-05-20 10:43:05 -07:00
parent f910a0c025
commit c87f5cf91c
10 changed files with 488 additions and 20 deletions
--- a/include/rocksdb/utilities/transaction.h
+++ b/include/rocksdb/utilities/transaction.h
@@ -443,6 +443,13 @@ class Transaction {
    }
  }

+  virtual Status GetEntityForUpdate(const ReadOptions& read_options,
+                                    ColumnFamilyHandle* column_family,
+                                    const Slice& key,
+                                    PinnableWideColumns* columns,
+                                    bool exclusive = true,
+                                    bool do_validate = true) = 0;
+
  virtual std::vector<Status> MultiGetForUpdate(
      const ReadOptions& options,
      const std::vector<ColumnFamilyHandle*>& column_family,
--- a/unreleased_history/new_features/get_entity_for_update.md
+++ b/unreleased_history/new_features/get_entity_for_update.md
@@ -0,0 +1 @@
+Optimistic transactions and pessimistic transactions with the WriteCommitted policy now support the `GetEntityForUpdate` API.
--- a/utilities/transactions/optimistic_transaction_test.cc
+++ b/utilities/transactions/optimistic_transaction_test.cc
@@ -1719,9 +1719,10 @@ TEST_P(OptimisticTransactionTest, PutEntitySuccess) {
    }

    {
-      PinnableSlice value;
-      ASSERT_OK(txn->GetForUpdate(ReadOptions(), foo, &value));
-      ASSERT_EQ(value, foo_columns[0].value());
+      PinnableWideColumns columns;
+      ASSERT_OK(txn->GetEntityForUpdate(
+          ReadOptions(), txn_db->DefaultColumnFamily(), foo, &columns));
+      ASSERT_EQ(columns.columns(), foo_columns);
    }

    ASSERT_OK(
@@ -1737,9 +1738,10 @@ TEST_P(OptimisticTransactionTest, PutEntitySuccess) {
    }

    {
-      PinnableSlice value;
-      ASSERT_OK(txn->GetForUpdate(ReadOptions(), foo, &value));
-      ASSERT_EQ(value, foo_new_columns[0].value());
+      PinnableWideColumns columns;
+      ASSERT_OK(txn->GetEntityForUpdate(
+          ReadOptions(), txn_db->DefaultColumnFamily(), foo, &columns));
+      ASSERT_EQ(columns.columns(), foo_new_columns);
    }

    ASSERT_OK(txn->Commit());
@@ -2034,6 +2036,60 @@ TEST_P(OptimisticTransactionTest, PutEntityWriteConflictTxnTxn) {
  }
 }

+TEST_P(OptimisticTransactionTest, PutEntityReadConflict) {
+  constexpr char foo[] = "foo";
+  const WideColumns foo_columns{
+      {kDefaultWideColumnName, "bar"}, {"col1", "val1"}, {"col2", "val2"}};
+
+  ASSERT_OK(txn_db->PutEntity(WriteOptions(), txn_db->DefaultColumnFamily(),
+                              foo, foo_columns));
+
+  std::unique_ptr<Transaction> txn(txn_db->BeginTransaction(WriteOptions()));
+  ASSERT_NE(txn, nullptr);
+
+  txn->SetSnapshot();
+
+  ReadOptions snapshot_read_options;
+  snapshot_read_options.snapshot = txn->GetSnapshot();
+
+  {
+    PinnableWideColumns columns;
+    ASSERT_OK(txn->GetEntityForUpdate(
+        snapshot_read_options, txn_db->DefaultColumnFamily(), foo, &columns));
+    ASSERT_EQ(columns.columns(), foo_columns);
+  }
+
+  // This PutEntity outside of a transaction will conflict with the previous
+  // write
+  const WideColumns foo_conflict_columns{{kDefaultWideColumnName, "X"},
+                                         {"conflicting", "write"}};
+  ASSERT_OK(txn_db->PutEntity(WriteOptions(), txn_db->DefaultColumnFamily(),
+                              foo, foo_conflict_columns));
+
+  {
+    PinnableWideColumns columns;
+    ASSERT_OK(txn_db->GetEntity(ReadOptions(), txn_db->DefaultColumnFamily(),
+                                foo, &columns));
+    ASSERT_EQ(columns.columns(), foo_conflict_columns);
+  }
+
+  {
+    PinnableWideColumns columns;
+    ASSERT_OK(txn->GetEntity(ReadOptions(), txn_db->DefaultColumnFamily(), foo,
+                             &columns));
+    ASSERT_EQ(columns.columns(), foo_conflict_columns);
+  }
+
+  ASSERT_TRUE(txn->Commit().IsBusy());  // Txn should not commit
+
+  {
+    PinnableWideColumns columns;
+    ASSERT_OK(txn_db->GetEntity(ReadOptions(), txn_db->DefaultColumnFamily(),
+                                foo, &columns));
+    ASSERT_EQ(columns.columns(), foo_conflict_columns);
+  }
+}
+
 TEST_P(OptimisticTransactionTest, EntityReadSanityChecks) {
  constexpr char foo[] = "foo";
  constexpr char bar[] = "bar";
@@ -2121,6 +2177,49 @@ TEST_P(OptimisticTransactionTest, EntityReadSanityChecks) {
    ASSERT_TRUE(statuses[0].IsInvalidArgument());
    ASSERT_TRUE(statuses[1].IsInvalidArgument());
  }
+
+  {
+    constexpr ColumnFamilyHandle* column_family = nullptr;
+    PinnableWideColumns columns;
+    ASSERT_TRUE(
+        txn->GetEntityForUpdate(ReadOptions(), column_family, foo, &columns)
+            .IsInvalidArgument());
+  }
+
+  {
+    constexpr PinnableWideColumns* columns = nullptr;
+    ASSERT_TRUE(txn->GetEntityForUpdate(ReadOptions(),
+                                        txn_db->DefaultColumnFamily(), foo,
+                                        columns)
+                    .IsInvalidArgument());
+  }
+
+  {
+    ReadOptions read_options;
+    read_options.io_activity = Env::IOActivity::kGet;
+
+    PinnableWideColumns columns;
+    ASSERT_TRUE(txn->GetEntityForUpdate(read_options,
+                                        txn_db->DefaultColumnFamily(), foo,
+                                        &columns)
+                    .IsInvalidArgument());
+  }
+
+  {
+    txn->SetSnapshot();
+
+    ReadOptions read_options;
+    read_options.snapshot = txn->GetSnapshot();
+
+    PinnableWideColumns columns;
+    constexpr bool exclusive = true;
+    constexpr bool do_validate = false;
+
+    ASSERT_TRUE(txn->GetEntityForUpdate(read_options,
+                                        txn_db->DefaultColumnFamily(), foo,
+                                        &columns, exclusive, do_validate)
+                    .IsInvalidArgument());
+  }
 }

 INSTANTIATE_TEST_CASE_P(
--- a/utilities/transactions/pessimistic_transaction.cc
+++ b/utilities/transactions/pessimistic_transaction.cc
@@ -217,6 +217,60 @@ inline Status WriteCommittedTxn::GetForUpdateImpl(
                                           value, exclusive, do_validate);
 }

+Status WriteCommittedTxn::GetEntityForUpdate(const ReadOptions& read_options,
+                                             ColumnFamilyHandle* column_family,
+                                             const Slice& key,
+                                             PinnableWideColumns* columns,
+                                             bool exclusive, bool do_validate) {
+  if (!column_family) {
+    return Status::InvalidArgument(
+        "Cannot call GetEntityForUpdate without a column family handle");
+  }
+
+  const Comparator* const ucmp = column_family->GetComparator();
+  assert(ucmp);
+  const size_t ts_sz = ucmp->timestamp_size();
+
+  if (ts_sz == 0) {
+    return TransactionBaseImpl::GetEntityForUpdate(
+        read_options, column_family, key, columns, exclusive, do_validate);
+  }
+
+  assert(ts_sz > 0);
+
+  if (!do_validate) {
+    if (read_timestamp_ != kMaxTxnTimestamp) {
+      return Status::InvalidArgument(
+          "Read timestamp must not be set if validation is disabled");
+    }
+  } else {
+    if (read_timestamp_ == kMaxTxnTimestamp) {
+      return Status::InvalidArgument(
+          "Read timestamp must be set for validation");
+    }
+  }
+
+  std::string ts_buf;
+  PutFixed64(&ts_buf, read_timestamp_);
+  Slice ts(ts_buf);
+
+  if (!read_options.timestamp) {
+    ReadOptions read_options_copy = read_options;
+    read_options_copy.timestamp = &ts;
+
+    return TransactionBaseImpl::GetEntityForUpdate(
+        read_options_copy, column_family, key, columns, exclusive, do_validate);
+  }
+
+  assert(read_options.timestamp);
+  if (*read_options.timestamp != ts) {
+    return Status::InvalidArgument("Must read from the same read timestamp");
+  }
+
+  return TransactionBaseImpl::GetEntityForUpdate(
+      read_options, column_family, key, columns, exclusive, do_validate);
+}
+
 Status WriteCommittedTxn::PutEntityImpl(ColumnFamilyHandle* column_family,
                                        const Slice& key,
                                        const WideColumns& columns,
--- a/utilities/transactions/pessimistic_transaction.h
+++ b/utilities/transactions/pessimistic_transaction.h
@@ -234,6 +234,11 @@ class WriteCommittedTxn : public PessimisticTransaction {
                      PinnableSlice* pinnable_val, bool exclusive,
                      const bool do_validate) override;

+  Status GetEntityForUpdate(const ReadOptions& read_options,
+                            ColumnFamilyHandle* column_family, const Slice& key,
+                            PinnableWideColumns* columns, bool exclusive,
+                            bool do_validate) override;
+
  using TransactionBaseImpl::Put;
  // `key` does NOT include timestamp even when it's enabled.
  Status Put(ColumnFamilyHandle* column_family, const Slice& key,
--- a/utilities/transactions/transaction_base.cc
+++ b/utilities/transactions/transaction_base.cc
@@ -351,6 +351,24 @@ Status TransactionBaseImpl::GetForUpdate(const ReadOptions& read_options,
  return s;
 }

+Status TransactionBaseImpl::GetEntityForUpdate(
+    const ReadOptions& read_options, ColumnFamilyHandle* column_family,
+    const Slice& key, PinnableWideColumns* columns, bool exclusive,
+    bool do_validate) {
+  if (!do_validate && read_options.snapshot != nullptr) {
+    return Status::InvalidArgument(
+        "Snapshot must not be set if validation is disabled");
+  }
+
+  const Status s =
+      TryLock(column_family, key, true /* read_only */, exclusive, do_validate);
+  if (!s.ok()) {
+    return s;
+  }
+
+  return GetEntityImpl(read_options, column_family, key, columns);
+}
+
 std::vector<Status> TransactionBaseImpl::MultiGet(
    const ReadOptions& _read_options,
    const std::vector<ColumnFamilyHandle*>& column_family,
--- a/utilities/transactions/transaction_base.h
+++ b/utilities/transactions/transaction_base.h
@@ -95,6 +95,11 @@ class TransactionBaseImpl : public Transaction {
                        exclusive, do_validate);
  }

+  Status GetEntityForUpdate(const ReadOptions& read_options,
+                            ColumnFamilyHandle* column_family, const Slice& key,
+                            PinnableWideColumns* columns, bool exclusive = true,
+                            bool do_validate = true) override;
+
  using Transaction::MultiGet;
  std::vector<Status> MultiGet(
      const ReadOptions& _read_options,
--- a/utilities/transactions/transaction_test.cc
+++ b/utilities/transactions/transaction_test.cc
@@ -7025,9 +7025,10 @@ TEST_P(TransactionTest, PutEntitySuccess) {
    }

    {
-      PinnableSlice value;
-      ASSERT_OK(txn->GetForUpdate(ReadOptions(), foo, &value));
-      ASSERT_EQ(value, foo_columns[0].value());
+      PinnableWideColumns columns;
+      ASSERT_OK(txn->GetEntityForUpdate(
+          ReadOptions(), db->DefaultColumnFamily(), foo, &columns));
+      ASSERT_EQ(columns.columns(), foo_columns);
    }

    ASSERT_OK(txn->PutEntity(db->DefaultColumnFamily(), foo, foo_new_columns));
@@ -7042,9 +7043,10 @@ TEST_P(TransactionTest, PutEntitySuccess) {
    }

    {
-      PinnableSlice value;
-      ASSERT_OK(txn->GetForUpdate(ReadOptions(), foo, &value));
-      ASSERT_EQ(value, foo_new_columns[0].value());
+      PinnableWideColumns columns;
+      ASSERT_OK(txn->GetEntityForUpdate(
+          ReadOptions(), db->DefaultColumnFamily(), foo, &columns));
+      ASSERT_EQ(columns.columns(), foo_new_columns);
    }

    ASSERT_OK(txn->Commit());
@@ -7227,6 +7229,67 @@ TEST_P(TransactionTest, PutEntityWriteConflict) {
  }
 }

+TEST_P(TransactionTest, PutEntityReadConflict) {
+  const TxnDBWritePolicy write_policy = std::get<2>(GetParam());
+  if (write_policy != TxnDBWritePolicy::WRITE_COMMITTED) {
+    ROCKSDB_GTEST_BYPASS("Test only WriteCommitted for now");
+    return;
+  }
+
+  constexpr char foo[] = "foo";
+  const WideColumns foo_columns{
+      {kDefaultWideColumnName, "bar"}, {"col1", "val1"}, {"col2", "val2"}};
+
+  ASSERT_OK(db->PutEntity(WriteOptions(), db->DefaultColumnFamily(), foo,
+                          foo_columns));
+
+  std::unique_ptr<Transaction> txn(db->BeginTransaction(WriteOptions()));
+  ASSERT_NE(txn, nullptr);
+
+  txn->SetSnapshot();
+
+  ReadOptions snapshot_read_options;
+  snapshot_read_options.snapshot = txn->GetSnapshot();
+
+  {
+    PinnableWideColumns columns;
+    ASSERT_OK(txn->GetEntityForUpdate(
+        snapshot_read_options, db->DefaultColumnFamily(), foo, &columns));
+    ASSERT_EQ(columns.columns(), foo_columns);
+  }
+
+  // This PutEntity outside of a transaction will conflict with the previous
+  // write
+  const WideColumns foo_conflict_columns{{kDefaultWideColumnName, "X"},
+                                         {"conflicting", "write"}};
+  ASSERT_TRUE(db->PutEntity(WriteOptions(), db->DefaultColumnFamily(), foo,
+                            foo_conflict_columns)
+                  .IsTimedOut());
+
+  {
+    PinnableWideColumns columns;
+    ASSERT_OK(
+        db->GetEntity(ReadOptions(), db->DefaultColumnFamily(), foo, &columns));
+    ASSERT_EQ(columns.columns(), foo_columns);
+  }
+
+  {
+    PinnableWideColumns columns;
+    ASSERT_OK(txn->GetEntity(ReadOptions(), db->DefaultColumnFamily(), foo,
+                             &columns));
+    ASSERT_EQ(columns.columns(), foo_columns);
+  }
+
+  ASSERT_OK(txn->Commit());
+
+  {
+    PinnableWideColumns columns;
+    ASSERT_OK(
+        db->GetEntity(ReadOptions(), db->DefaultColumnFamily(), foo, &columns));
+    ASSERT_EQ(columns.columns(), foo_columns);
+  }
+}
+
 TEST_P(TransactionTest, EntityReadSanityChecks) {
  constexpr char foo[] = "foo";
  constexpr char bar[] = "bar";
@@ -7314,6 +7377,46 @@ TEST_P(TransactionTest, EntityReadSanityChecks) {
    ASSERT_TRUE(statuses[0].IsInvalidArgument());
    ASSERT_TRUE(statuses[1].IsInvalidArgument());
  }
+
+  {
+    constexpr ColumnFamilyHandle* column_family = nullptr;
+    PinnableWideColumns columns;
+    ASSERT_TRUE(
+        txn->GetEntityForUpdate(ReadOptions(), column_family, foo, &columns)
+            .IsInvalidArgument());
+  }
+
+  {
+    constexpr PinnableWideColumns* columns = nullptr;
+    ASSERT_TRUE(txn->GetEntityForUpdate(ReadOptions(),
+                                        db->DefaultColumnFamily(), foo, columns)
+                    .IsInvalidArgument());
+  }
+
+  {
+    ReadOptions read_options;
+    read_options.io_activity = Env::IOActivity::kGet;
+
+    PinnableWideColumns columns;
+    ASSERT_TRUE(txn->GetEntityForUpdate(read_options, db->DefaultColumnFamily(),
+                                        foo, &columns)
+                    .IsInvalidArgument());
+  }
+
+  {
+    txn->SetSnapshot();
+
+    ReadOptions read_options;
+    read_options.snapshot = txn->GetSnapshot();
+
+    PinnableWideColumns columns;
+    constexpr bool exclusive = true;
+    constexpr bool do_validate = false;
+
+    ASSERT_TRUE(txn->GetEntityForUpdate(read_options, db->DefaultColumnFamily(),
+                                        foo, &columns, exclusive, do_validate)
+                    .IsInvalidArgument());
+  }
 }

 TEST_F(TransactionDBTest, CollapseKey) {
--- a/utilities/transactions/write_committed_transaction_ts_test.cc
+++ b/utilities/transactions/write_committed_transaction_ts_test.cc
@@ -717,6 +717,147 @@ TEST_P(WriteCommittedTxnWithTsTest, CheckKeysForConflicts) {
  SyncPoint::GetInstance()->ClearAllCallBacks();
 }

+TEST_P(WriteCommittedTxnWithTsTest, GetEntityForUpdate) {
+  ASSERT_OK(ReOpenNoDelete());
+
+  ColumnFamilyOptions cf_options;
+  cf_options.comparator = test::BytewiseComparatorWithU64TsWrapper();
+
+  const std::string test_cf_name = "test_cf";
+
+  ColumnFamilyHandle* cfh = nullptr;
+  ASSERT_OK(db->CreateColumnFamily(cf_options, test_cf_name, &cfh));
+  std::unique_ptr<ColumnFamilyHandle> cfh_guard(cfh);
+
+  constexpr char foo[] = "foo";
+  constexpr char bar[] = "bar";
+  constexpr char baz[] = "baz";
+  constexpr char quux[] = "quux";
+
+  {
+    std::unique_ptr<Transaction> txn0(
+        NewTxn(WriteOptions(), TransactionOptions()));
+
+    {
+      std::unique_ptr<Transaction> txn1(
+          NewTxn(WriteOptions(), TransactionOptions()));
+      ASSERT_OK(txn1->Put(cfh, foo, bar));
+      ASSERT_OK(txn1->Put(cfh, baz, quux));
+      ASSERT_OK(txn1->SetCommitTimestamp(24));
+      ASSERT_OK(txn1->Commit());
+    }
+
+    ASSERT_OK(txn0->SetReadTimestampForValidation(23));
+
+    // Validation fails: timestamp from db(24) > validation timestamp(23)
+    PinnableWideColumns columns;
+    ASSERT_TRUE(
+        txn0->GetEntityForUpdate(ReadOptions(), cfh, foo, &columns).IsBusy());
+
+    ASSERT_OK(txn0->Rollback());
+  }
+
+  {
+    std::unique_ptr<Transaction> txn2(
+        NewTxn(WriteOptions(), TransactionOptions()));
+
+    ASSERT_OK(txn2->SetReadTimestampForValidation(25));
+
+    // Validation successful: timestamp from db(24) < validation timestamp (25)
+    {
+      PinnableWideColumns columns;
+      ASSERT_OK(txn2->GetEntityForUpdate(ReadOptions(), cfh, foo, &columns));
+    }
+
+    // Using a different read timestamp in ReadOptions while doing validation is
+    // not allowed
+    {
+      ReadOptions read_options;
+      std::string read_timestamp;
+      Slice diff_read_ts = EncodeU64Ts(24, &read_timestamp);
+      read_options.timestamp = &diff_read_ts;
+
+      PinnableWideColumns columns;
+      ASSERT_TRUE(txn2->GetEntityForUpdate(read_options, cfh, foo, &columns)
+                      .IsInvalidArgument());
+
+      ASSERT_OK(txn2->SetCommitTimestamp(26));
+      ASSERT_OK(txn2->Commit());
+    }
+  }
+
+  // GetEntityForUpdate with validation timestamp set but no validation is not
+  // allowed
+  {
+    std::unique_ptr<Transaction> txn3(
+        NewTxn(WriteOptions(), TransactionOptions()));
+
+    ASSERT_OK(txn3->SetReadTimestampForValidation(27));
+
+    PinnableWideColumns columns;
+    ASSERT_TRUE(txn3->GetEntityForUpdate(ReadOptions(), cfh, foo, &columns,
+                                         /*exclusive=*/true,
+                                         /*do_validate=*/false)
+                    .IsInvalidArgument());
+
+    ASSERT_OK(txn3->Rollback());
+  }
+
+  // GetEntityForUpdate with validation but no validation timestamp is not
+  // allowed
+  {
+    std::unique_ptr<Transaction> txn4(
+        NewTxn(WriteOptions(), TransactionOptions()));
+
+    // ReadOptions.timestamp is not set
+    {
+      PinnableWideColumns columns;
+      ASSERT_TRUE(txn4->GetEntityForUpdate(ReadOptions(), cfh, foo, &columns)
+                      .IsInvalidArgument());
+    }
+
+    // ReadOptions.timestamp is set
+    {
+      ReadOptions read_options;
+      std::string read_timestamp;
+      Slice read_ts = EncodeU64Ts(27, &read_timestamp);
+      read_options.timestamp = &read_ts;
+
+      PinnableWideColumns columns;
+      ASSERT_TRUE(txn4->GetEntityForUpdate(read_options, cfh, foo, &columns)
+                      .IsInvalidArgument());
+    }
+
+    ASSERT_OK(txn4->Rollback());
+  }
+
+  // Validation disabled
+  {
+    std::unique_ptr<Transaction> txn5(
+        NewTxn(WriteOptions(), TransactionOptions()));
+
+    // ReadOptions.timestamp is not set => success
+    {
+      PinnableWideColumns columns;
+      ASSERT_OK(txn5->GetEntityForUpdate(ReadOptions(), cfh, foo, &columns,
+                                         /*exclusive=*/true,
+                                         /*do_validate=*/false));
+    }
+
+    // ReadOptions.timestamp explicitly set to max timestamp => success
+    {
+      ReadOptions read_options;
+      Slice max_ts = MaxU64Ts();
+      read_options.timestamp = &max_ts;
+
+      PinnableWideColumns columns;
+      ASSERT_OK(txn5->GetEntityForUpdate(read_options, cfh, baz, &columns,
+                                         /*exclusive=*/true,
+                                         /*do_validate=*/false));
+    }
+  }
+}
+
 }  // namespace ROCKSDB_NAMESPACE

 int main(int argc, char** argv) {
@@ -724,4 +865,3 @@ int main(int argc, char** argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
 }
-
--- a/utilities/write_batch_with_index/write_batch_with_index.cc
+++ b/utilities/write_batch_with_index/write_batch_with_index.cc
@@ -849,8 +849,22 @@ Status WriteBatchWithIndex::GetEntityFromBatchAndDB(

  const Comparator* const ucmp = rep->comparator.GetComparator(column_family);
  size_t ts_sz = ucmp ? ucmp->timestamp_size() : 0;
-  if (ts_sz > 0 && !read_options.timestamp) {
-    return Status::InvalidArgument("Must specify timestamp");
+  if (ts_sz > 0) {
+    if (!read_options.timestamp) {
+      return Status::InvalidArgument("Must specify timestamp");
+    }
+
+    if (read_options.timestamp->size() != ts_sz) {
+      return Status::InvalidArgument(
+          "Timestamp size does not match the timestamp size of the "
+          "column family");
+    }
+  } else {
+    if (read_options.timestamp) {
+      return Status::InvalidArgument(
+          "Cannot specify timestamp since the column family does not have "
+          "timestamps enabled");
+    }
  }

  if (!columns) {
@@ -948,12 +962,34 @@ void WriteBatchWithIndex::MultiGetEntityFromBatchAndDB(

  const Comparator* const ucmp = rep->comparator.GetComparator(column_family);
  const size_t ts_sz = ucmp ? ucmp->timestamp_size() : 0;
-  if (ts_sz > 0 && !read_options.timestamp) {
-    const Status s = Status::InvalidArgument("Must specify timestamp");
-    for (size_t i = 0; i < num_keys; ++i) {
-      statuses[i] = s;
+  if (ts_sz > 0) {
+    if (!read_options.timestamp) {
+      const Status s = Status::InvalidArgument("Must specify timestamp");
+      for (size_t i = 0; i < num_keys; ++i) {
+        statuses[i] = s;
+      }
+      return;
+    }
+
+    if (read_options.timestamp->size() != ts_sz) {
+      const Status s = Status::InvalidArgument(
+          "Timestamp size does not match the timestamp size of the "
+          "column family");
+      for (size_t i = 0; i < num_keys; ++i) {
+        statuses[i] = s;
+      }
+      return;
+    }
+  } else {
+    if (read_options.timestamp) {
+      const Status s = Status::InvalidArgument(
+          "Cannot specify timestamp since the column family does not have "
+          "timestamps enabled");
+      for (size_t i = 0; i < num_keys; ++i) {
+        statuses[i] = s;
+      }
+      return;
    }
-    return;
  }

  if (!keys) {
				`@@ -0,0 +1 @@`
				Optimistic transactions and pessimistic transactions with the WriteCommitted policy now support the `GetEntityForUpdate` API.