Add kvstore::Driver::GetImplicitTransaction method

This facilities better automatic coalescing of writes for the
neuroglancer_precomputed and zarr v3 sharded formats.

PiperOrigin-RevId: 547299354
Change-Id: Ie85443f5fe5fb93880f53d98d4214f2b6b08347d

tensorstore/kvstore/driver.h

@@ -194,6 +194,12 @@ class Driver {
       internal::OpenTransactionPtr& transaction, size_t& phase, Key key,
       ReadModifyWriteSource& source);
 
+  /// Returns an implicit transaction for the specified key.
+  ///
+  /// This may return either an existing or new implicit transaction.
+  virtual Result<internal::OpenTransactionPtr> GetImplicitTransaction(
+      const Key& key);
+
   /// Registers a transactional delete range operation.
   ///
   /// The actual deletion will not occur until the transaction is committed.

tensorstore/kvstore/neuroglancer_uint64_sharded/neuroglancer_uint64_sharded.cc

@@ -1096,6 +1096,21 @@ class ShardedKeyValueStore
     return absl::OkStatus();
   }
 
+  Result<internal::OpenTransactionPtr> GetImplicitTransaction(
+      const Key& key) override {
+    TENSORSTORE_ASSIGN_OR_RETURN(ChunkId chunk_id, KeyToChunkIdOrError(key));
+    const auto& sharding_spec = this->sharding_spec();
+    const auto shard_info = GetSplitShardInfo(
+        sharding_spec, GetChunkShardInfo(sharding_spec, chunk_id));
+    const std::uint64_t shard = shard_info.shard;
+    auto entry = GetCacheEntry(
+        write_cache_, ShardedKeyValueStoreWriteCache::ShardToKey(shard));
+    internal::OpenTransactionPtr transaction;
+    TENSORSTORE_ASSIGN_OR_RETURN(auto node,
+                                 GetTransactionNode(*entry, transaction));
+    return transaction;
+  }
+
   absl::Status TransactionalDeleteRange(
       const internal::OpenTransactionPtr& transaction,
       KeyRange range) override {

tensorstore/kvstore/test_util.cc

@@ -38,13 +38,15 @@
 #include "tensorstore/internal/json_gtest.h"
 #include "tensorstore/json_serialization_options.h"
 #include "tensorstore/kvstore/byte_range.h"
+#include "tensorstore/kvstore/driver.h"
 #include "tensorstore/kvstore/generation.h"
 #include "tensorstore/kvstore/generation_testutil.h"
 #include "tensorstore/kvstore/key_range.h"
 #include "tensorstore/kvstore/kvstore.h"
 #include "tensorstore/kvstore/operations.h"
 #include "tensorstore/kvstore/read_result.h"
 #include "tensorstore/kvstore/spec.h"
+#include "tensorstore/kvstore/transaction.h"
 #include "tensorstore/util/execution/execution.h"
 #include "tensorstore/util/execution/sender_testutil.h"
 #include "tensorstore/util/future.h"
@@ -513,6 +515,46 @@ void TestKeyValueStoreStalenessBoundOps(
                   MatchesKvsReadResult(value2, write_result2->generation)));
 }
 
+void TestKeyValueStoreGetImplicitTransaction(
+    const KvStore& store,
+    absl::FunctionRef<std::string(std::string key)> get_key) {
+  std::vector<std::string> keys;
+  constexpr size_t kNumKeys = 4;
+  for (size_t i = 0; i < kNumKeys; ++i) {
+    keys.push_back(get_key(absl::StrFormat("testImplicit%d", i)));
+  }
+  Cleanup cleanup(store, keys);
+
+  std::vector<internal::OpenTransactionPtr> implicit_txns_direct;
+  std::vector<internal::OpenTransactionPtr>
+      implicit_txns_from_read_modify_write;
+  for (const auto& key : keys) {
+    auto full_key = store.path + key;
+    {
+      TENSORSTORE_ASSERT_OK_AND_ASSIGN(
+          auto txn, store.driver->GetImplicitTransaction(full_key));
+      implicit_txns_direct.push_back(std::move(txn));
+    }
+    {
+      internal::OpenTransactionPtr txn;
+      size_t phase;
+      auto future = internal_kvstore::WriteViaExistingTransaction(
+          store.driver.get(), txn, phase, full_key, std::nullopt,
+          kvstore::WriteOptions{});
+      ASSERT_TRUE(txn);
+      implicit_txns_from_read_modify_write.push_back(std::move(txn));
+    }
+  }
+
+  for (size_t i = 0; i < kNumKeys; ++i) {
+    for (size_t j = i + 1; j < kNumKeys; ++j) {
+      EXPECT_EQ((implicit_txns_direct[i] == implicit_txns_direct[j]),
+                (implicit_txns_from_read_modify_write[i] ==
+                 implicit_txns_from_read_modify_write[j]));
+    }
+  }
+}
+
 }  // namespace
 
 void TestKeyValueStoreBasicFunctionality(
@@ -523,6 +565,7 @@ void TestKeyValueStoreBasicFunctionality(
   TestKeyValueStoreConditionalWriteOps(store, get_key);
   TestKeyValueStoreConditionalDeleteOps(store, get_key);
   TestKeyValueStoreStalenessBoundOps(store, get_key);
+  TestKeyValueStoreGetImplicitTransaction(store, get_key);
 }
 
 /// Tests List on `store`, which should be empty.

tensorstore/kvstore/transaction.cc

@@ -1606,6 +1606,11 @@ absl::Status Driver::ReadModifyWrite(internal::OpenTransactionPtr& transaction,
   return internal_kvstore::GetNonAtomicReadModifyWriteError(*node, rmw_status);
 }
 
+Result<internal::OpenTransactionPtr> Driver::GetImplicitTransaction(
+    const Key& key) {
+  return internal::TransactionState::MakeImplicit();
+}
+
 absl::Status Driver::TransactionalDeleteRange(
     const internal::OpenTransactionPtr& transaction, KeyRange range) {
   if (range.empty()) return absl::OkStatus();