Add batch to VirtualChunked ReadParameters

.github/workflows/build.yml

@@ -10,7 +10,9 @@ jobs:
           - os: "ubuntu-latest"
             cibw_build: "*"
             wheel_identifier: linux_x86_64
-          - os: "arm-ubuntu-arm-22.04-8core"
+          - os: "ubuntu-22.04-arm"
+            # TODO: Revert back to the following before merge
+            # - os: "arm-ubuntu-arm-22.04-8core"
             cibw_build: "*"
             wheel_identifier: linux_arm64
           - os: "windows-2022"

python/tensorstore/__init__.pyi

@@ -12185,6 +12185,12 @@ class VirtualChunkedReadParameters:
       Virtual views
     """
 
+    @property
+    def batch(self) -> Batch | None:
+        """
+        Batch associated with the read request.
+        """
+
     @property
     def if_not_equal(self) -> bytes:
         """

python/tensorstore/tests/virtual_chunked_test.py

@@ -309,3 +309,51 @@ def do_read(
   do_read.future = future  # type: ignore
 
   gc_tester(t)
+
+
+def test_read_batch() -> None:
+  """Tests reading both with and without a batch."""
+  array = ts.array(np.arange(np.prod(shape := (4, 5))).reshape(shape))
+  # Avoid sharing the common threadpool with array
+  context = ts.Context({"data_copy_concurrency": {"limit": 1}})
+
+  def do_batched_read(
+      domain: ts.IndexDomain,
+      chunk: np.ndarray,
+      read_params: ts.VirtualChunkedReadParameters,
+  ) -> None:
+    assert isinstance(read_params.batch, ts.Batch)
+    chunk[...] = array[domain].read(batch=read_params.batch).result()
+
+  t = ts.virtual_chunked(
+      do_batched_read,
+      None,
+      dtype=array.dtype,
+      shape=array.shape,
+      chunk_layout=ts.ChunkLayout(read_chunk_shape=(2, 3)),
+      context=context,
+  )
+
+  with ts.Batch() as b:
+    f = t[1:3, 1:3].read(batch=b)
+
+  np.testing.assert_array_equal(f.result(), array[1:3, 1:3])
+
+  def do_unbatched_read(
+      domain: ts.IndexDomain,
+      chunk: np.ndarray,
+      read_params: ts.VirtualChunkedReadParameters,
+  ) -> None:
+    assert read_params.batch is None
+    chunk[...] = array[domain].read(batch=read_params.batch).result()
+
+  t2 = ts.virtual_chunked(
+      do_unbatched_read,
+      None,
+      dtype=array.dtype,
+      shape=array.shape,
+      chunk_layout=ts.ChunkLayout(read_chunk_shape=(2, 3)),
+      context=context
+  )
+
+  np.testing.assert_array_equal(t2[1:3, 1:3].read().result(), array[1:3, 1:3])

python/tensorstore/virtual_chunked.cc

@@ -42,6 +42,7 @@
 #include "python/tensorstore/time.h"
 #include "python/tensorstore/type_name_override.h"
 #include "tensorstore/array.h"
+#include "tensorstore/batch.h"
 #include "tensorstore/container_kind.h"
 #include "tensorstore/context.h"
 #include "tensorstore/index_space/index_domain.h"
@@ -96,6 +97,16 @@ Cached generation, read request can be skipped if no newer data is available.
       },
       R"(
 Read may be fulfilled with cached data no older than the specified bound.
+)");
+
+  cls.def_property_readonly(
+      "batch",
+      [](const Self& self) {
+        if (auto b = self.batch(); b) return std::optional<Batch>{b};
+        return std::optional<Batch>{std::nullopt};
+      },
+      R"(
+Batch associated with the read request.
 )");
 }
 

tensorstore/driver/virtual_chunked/BUILD

@@ -27,6 +27,7 @@ tensorstore_cc_library(
         "//tensorstore:open_mode",
         "//tensorstore:open_options",
         "//tensorstore:rank",
+        "//tensorstore:schema",
         "//tensorstore:staleness_bound",
         "//tensorstore:strided_layout",
         "//tensorstore:transaction",

tensorstore/driver/virtual_chunked/virtual_chunked.cc

@@ -30,6 +30,7 @@
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "tensorstore/array.h"
+#include "tensorstore/batch.h"
 #include "tensorstore/box.h"
 #include "tensorstore/chunk_layout.h"
 #include "tensorstore/codec_spec.h"
@@ -61,6 +62,7 @@
 #include "tensorstore/open_mode.h"
 #include "tensorstore/open_options.h"
 #include "tensorstore/rank.h"
+#include "tensorstore/schema.h"
 #include "tensorstore/serialization/absl_time.h"  // IWYU pragma: keep
 #include "tensorstore/serialization/std_optional.h"  // IWYU pragma: keep
 #include "tensorstore/staleness_bound.h"
@@ -237,7 +239,8 @@ void VirtualChunkedCache::DoRead(EntryOrNode& node,
   auto& executor = cache.executor();
   // `node` is guaranteed to remain valid until `ReadSuccess` or `ReadError`
   // is called.  Therefore we don't need to separately hold a reference.
-  executor([&node, staleness_bound = request.staleness_bound] {
+  executor([&node, staleness_bound = request.staleness_bound,
+            batch = Batch(request.batch)] {
     auto& entry = GetOwningEntry(node);
     auto& cache = GetOwningCache(entry);
     const auto& component_spec = cache.grid().components.front();
@@ -267,6 +270,7 @@ void VirtualChunkedCache::DoRead(EntryOrNode& node,
       read_params.if_not_equal_ = lock.stamp().generation;
     }
     read_params.staleness_bound_ = staleness_bound;
+    read_params.batch_ = std::move(batch);
     auto read_future =
         cache.read_function_(ConstDataTypeCast<void>(std::move(partial_array)),
                              std::move(read_params));

tensorstore/driver/virtual_chunked/virtual_chunked_test.cc

@@ -15,6 +15,7 @@
 #include "tensorstore/virtual_chunked.h"
 
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -25,6 +26,7 @@
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "tensorstore/array.h"
+#include "tensorstore/batch.h"
 #include "tensorstore/chunk_layout.h"
 #include "tensorstore/context.h"
 #include "tensorstore/data_type.h"
@@ -51,6 +53,7 @@
 
 namespace {
 
+using ::tensorstore::Batch;
 using ::tensorstore::DimensionIndex;
 using ::tensorstore::dynamic_rank;
 using ::tensorstore::Future;
@@ -587,4 +590,76 @@ TEST(VirtualChunkedTest, NonAtomicSingleChunk) {
   TENSORSTORE_ASSERT_OK(future);
 }
 
+template <typename... Option>
+Result<tensorstore::TensorStore<void, dynamic_rank,
+                                tensorstore::ReadWriteMode::read>>
+BatchSettingView(std::optional<Batch::View>& output_batch, Option&&... option) {
+  auto mutex = std::make_shared<absl::Mutex>();
+  return tensorstore::VirtualChunked(
+      tensorstore::NonSerializable{
+          [mutex, &output_batch](auto output, auto read_params)
+              -> Future<TimestampedStorageGeneration> {
+            tensorstore::InitializeArray(output);
+            absl::MutexLock lock(*mutex.get());
+            output_batch = read_params.batch();
+            return TimestampedStorageGeneration{
+                StorageGeneration::FromString("abc"), absl::Now()};
+          }},
+      std::forward<Option>(option)...);
+}
+
+// Test passing of Read batch into VirtualChunked ReadParameters
+// No Batch argument used in tensorstore::Read
+// results in output_batch not set
+TEST(VirtualChunkedTest, ReadNoBatch) {
+  std::optional<Batch::View> output_batch{std::nullopt};
+  auto virtual_chunked =
+      BatchSettingView(output_batch, tensorstore::dtype_v<int>,
+                       tensorstore::Schema::Shape({2, 3}),
+                       tensorstore::ChunkLayout::ReadChunkShape({2, 1}));
+
+  EXPECT_FALSE(output_batch.has_value());
+  auto data = tensorstore::Read(virtual_chunked).result();
+  EXPECT_TRUE(output_batch.has_value());
+  EXPECT_FALSE(*output_batch);
+}
+
+// No Batch argument passed to tensorstore::Read
+// results in output_batch not set
+TEST(VirtualChunkedTest, ReadNoBatchArgument) {
+  std::optional<Batch::View> output_batch{std::nullopt};
+  auto batch = Batch{tensorstore::no_batch};
+  auto virtual_chunked =
+      BatchSettingView(output_batch, tensorstore::dtype_v<int>,
+                       tensorstore::Schema::Shape({2, 3}),
+                       tensorstore::ChunkLayout::ReadChunkShape({2, 1}));
+
+  EXPECT_FALSE(output_batch.has_value());
+  auto read_future = tensorstore::Read(virtual_chunked, batch);
+  batch.Release();
+  auto data = read_future.result();
+
+  EXPECT_TRUE(output_batch.has_value());
+  EXPECT_FALSE(*output_batch);
+}
+
+// Batch argument passed to tensorstore::Read
+// is reflected in output_batch
+TEST(VirtualChunkedTest, ReadBatchArgument) {
+  std::optional<Batch::View> output_batch{std::nullopt};
+  auto virtual_chunked =
+      BatchSettingView(output_batch, tensorstore::dtype_v<int>,
+                       tensorstore::Schema::Shape({2, 3}),
+                       tensorstore::ChunkLayout::ReadChunkShape({2, 1}));
+  auto batch = Batch::New();
+  Batch::View batch_view{batch};
+  EXPECT_FALSE(output_batch.has_value());
+  auto read_future = tensorstore::Read(virtual_chunked, batch);
+  batch.Release();
+  auto data = read_future.result();
+
+  EXPECT_TRUE(output_batch.has_value());
+  EXPECT_TRUE(*output_batch);
+}
+
 }  // namespace

tensorstore/virtual_chunked.h

@@ -243,22 +243,32 @@
 /// no different than binding the transaction to an existing virtual chunked
 /// view.
 
-#include <functional>
 #include <type_traits>
+#include <utility>
 
 #include "absl/base/attributes.h"
 #include "absl/meta/type_traits.h"
 #include "absl/status/status.h"
+#include "absl/time/time.h"
 #include "tensorstore/array.h"
-#include "tensorstore/box.h"
+#include "tensorstore/batch.h"
 #include "tensorstore/context.h"
+#include "tensorstore/driver/driver.h"
+#include "tensorstore/index.h"
 #include "tensorstore/kvstore/generation.h"
+#include "tensorstore/open_mode.h"
+#include "tensorstore/rank.h"
+#include "tensorstore/schema.h"
 #include "tensorstore/serialization/function.h"
 #include "tensorstore/staleness_bound.h"
+#include "tensorstore/static_cast.h"
+#include "tensorstore/strided_layout.h"
 #include "tensorstore/tensorstore.h"
 #include "tensorstore/transaction.h"
 #include "tensorstore/util/executor.h"
-#include "tensorstore/util/option.h"
+#include "tensorstore/util/future.h"
+#include "tensorstore/util/result.h"
+#include "tensorstore/util/status.h"
 
 namespace tensorstore {
 namespace virtual_chunked {
@@ -278,11 +288,14 @@ class ReadParameters {
   /// Read may be fulfilled with cached data no older than the specified bound.
   absl::Time staleness_bound() const { return staleness_bound_; }
 
-  // Treat as private:
+  /// Batch associated with the read request.
+  Batch::View batch() const { return batch_; }
 
+  // Treat as private:
   Executor executor_;
   StorageGeneration if_not_equal_;
   absl::Time staleness_bound_;
+  Batch batch_{no_batch};
 };
 
 /// Type-erased function called to read a single chunk.