diff --git a/docs/conf.py b/docs/conf.py index 4d38d7ef9..33ad5827e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -217,7 +217,11 @@ """ -python_module_names_to_strip_from_xrefs = ['tensorstore'] +python_module_names_to_strip_from_xrefs = [ + 'tensorstore', + 'collections.abc', + 'numbers', +] python_type_aliases = { 'dtype': 'numpy.dtype', diff --git a/python/tensorstore/BUILD b/python/tensorstore/BUILD index cf54dff1e..945dedd3d 100644 --- a/python/tensorstore/BUILD +++ b/python/tensorstore/BUILD @@ -50,6 +50,7 @@ pybind11_py_extension( ":python_imports", ":serialization", ":spec", + ":stack", ":tensorstore_class", ":tensorstore_module_components", ":transaction", @@ -182,6 +183,16 @@ tensorstore_pytest_test( ], ) +tensorstore_pytest_test( + name = "stack_test", + size = "small", + srcs = ["tests/stack_test.py"], + deps = [ + ":tensorstore", + "@pypa_numpy//:numpy", + ], +) + pybind11_cc_library( name = "subscript_method", hdrs = ["subscript_method.h"], @@ -830,6 +841,29 @@ pybind11_cc_library( alwayslink = True, ) +pybind11_cc_library( + name = "stack", + srcs = ["stack.cc"], + deps = [ + ":dim_expression", + ":index", + ":keyword_arguments", + ":result_type_caster", + ":sequence_parameter", + ":spec", + ":tensorstore_class", + ":tensorstore_module_components", + "//tensorstore", + "//tensorstore:index", + "//tensorstore:spec", + "//tensorstore:stack", + "//tensorstore/internal:global_initializer", + "//tensorstore/util:executor", + "@com_github_pybind_pybind11//:pybind11", + ], + alwayslink = True, +) + pybind11_cc_library( name = "cast", srcs = ["cast.cc"], diff --git a/python/tensorstore/stack.cc b/python/tensorstore/stack.cc new file mode 100644 index 000000000..08a5bad6a --- /dev/null +++ b/python/tensorstore/stack.cc @@ -0,0 +1,444 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +// Other headers must be included after pybind11 to ensure header-order +// inclusion constraints are satisfied. + +// Other headers +#include + +#include +#include +#include +#include + +#include "python/tensorstore/dim_expression.h" +#include "python/tensorstore/index.h" // IWYU pragma: keep +#include "python/tensorstore/keyword_arguments.h" +#include "python/tensorstore/result_type_caster.h" +#include "python/tensorstore/sequence_parameter.h" +#include "python/tensorstore/spec.h" +#include "python/tensorstore/tensorstore_class.h" +#include "python/tensorstore/tensorstore_module_components.h" +#include "tensorstore/index.h" +#include "tensorstore/internal/global_initializer.h" +#include "tensorstore/spec.h" +#include "tensorstore/stack.h" +#include "tensorstore/tensorstore.h" +#include "tensorstore/util/executor.h" + +namespace tensorstore { +namespace internal_python { +namespace { + +namespace py = ::pybind11; + +constexpr auto ForwardStackSetters = [](auto callback, auto... other_param) { + callback(other_param..., open_setters::SetRead{}, open_setters::SetWrite{}, + open_setters::SetContext{}, open_setters::SetTransaction{}, + schema_setters::SetRank{}, schema_setters::SetDtype{}, + schema_setters::SetDomain{}, schema_setters::SetShape{}, + schema_setters::SetDimensionUnits{}, schema_setters::SetSchema{}); +}; + +void RegisterStackBindings(pybind11::module m, Executor defer) { + defer([m]() mutable { + ForwardStackSetters([&](auto... param_def) { + { + std::string doc = R"( +Virtually overlays a sequence of :py:obj:`TensorStore` layers within a common domain. + + >>> store = ts.overlay([ + ... ts.array([1, 2, 3, 4], dtype=ts.uint32), + ... ts.array([5, 6, 7, 8], dtype=ts.uint32).translate_to[3] + ... ]) + >>> store + TensorStore({ + 'context': {'data_copy_concurrency': {}}, + 'driver': 'stack', + 'dtype': 'uint32', + 'layers': [ + { + 'array': [1, 2, 3, 4], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': {'input_exclusive_max': [4], 'input_inclusive_min': [0]}, + }, + { + 'array': [5, 6, 7, 8], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [7], + 'input_inclusive_min': [3], + 'output': [{'input_dimension': 0, 'offset': -3}], + }, + }, + ], + 'schema': {'domain': {'exclusive_max': [7], 'inclusive_min': [0]}}, + 'transform': {'input_exclusive_max': [7], 'input_inclusive_min': [0]}, + }) + >>> await store.read() + array([1, 2, 3, 5, 6, 7, 8], dtype=uint32) + +Args: + + layers: Sequence of layers to overlay. Later layers take precedence. If a + layer is specified as a :py:obj:`Spec` rather than a :py:obj:`TensorStore`, + it must have a known :py:obj:`~Spec.domain` and will be opened on-demand as + neneded for individual read and write operations. + +)"; + AppendKeywordArgumentDocs(doc, param_def...); + doc += R"( + +See also: + - :ref:`stack-driver` + - :py:obj:`tensorstore.stack` + - :py:obj:`tensorstore.concat` + +Group: + Views +)"; + m.def( + "overlay", + [](SequenceParameter< + std::variant> + python_layers, + KeywordArgument... kwarg) -> TensorStore<> { + StackOpenOptions options; + ApplyKeywordArguments(options, kwarg...); + std::vector, Spec>> layers( + python_layers.size()); + for (size_t i = 0; i < layers.size(); ++i) { + std::visit([&](auto* obj) { layers[i] = obj->value; }, + python_layers[i]); + } + return ValueOrThrow( + tensorstore::Overlay(layers, std::move(options))); + }, + doc.c_str(), py::arg("layers"), py::kw_only(), + MakeKeywordArgumentPyArg(param_def)...); + } + + { + std::string doc = R"( +Virtually stacks a sequence of :py:obj:`TensorStore` layers along a new dimension. + + >>> store = ts.stack([ + ... ts.array([1, 2, 3, 4], dtype=ts.uint32), + ... ts.array([5, 6, 7, 8], dtype=ts.uint32) + ... ]) + >>> store + TensorStore({ + 'context': {'data_copy_concurrency': {}}, + 'driver': 'stack', + 'dtype': 'uint32', + 'layers': [ + { + 'array': [1, 2, 3, 4], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [1, 4], + 'input_inclusive_min': [0, 0], + 'output': [{'input_dimension': 1}], + }, + }, + { + 'array': [5, 6, 7, 8], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [2, 4], + 'input_inclusive_min': [1, 0], + 'output': [{'input_dimension': 1}], + }, + }, + ], + 'schema': {'domain': {'exclusive_max': [2, 4], 'inclusive_min': [0, 0]}}, + 'transform': {'input_exclusive_max': [2, 4], 'input_inclusive_min': [0, 0]}, + }) + >>> await store.read() + array([[1, 2, 3, 4], + [5, 6, 7, 8]], dtype=uint32) + >>> store = ts.stack([ + ... ts.array([1, 2, 3, 4], dtype=ts.uint32), + ... ts.array([5, 6, 7, 8], dtype=ts.uint32) + ... ], + ... axis=-1) + >>> store + TensorStore({ + 'context': {'data_copy_concurrency': {}}, + 'driver': 'stack', + 'dtype': 'uint32', + 'layers': [ + { + 'array': [1, 2, 3, 4], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [4, 1], + 'input_inclusive_min': [0, 0], + 'output': [{'input_dimension': 0}], + }, + }, + { + 'array': [5, 6, 7, 8], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [4, 2], + 'input_inclusive_min': [0, 1], + 'output': [{'input_dimension': 0}], + }, + }, + ], + 'schema': {'domain': {'exclusive_max': [4, 2], 'inclusive_min': [0, 0]}}, + 'transform': {'input_exclusive_max': [4, 2], 'input_inclusive_min': [0, 0]}, + }) + >>> await store.read() + array([[1, 5], + [2, 6], + [3, 7], + [4, 8]], dtype=uint32) + +Args: + + layers: Sequence of layers to stack. If a layer is specified as a + :py:obj:`Spec` rather than a :py:obj:`TensorStore`, it must have a known + :py:obj:`~Spec.domain` and will be opened on-demand as needed for individual + read and write operations. + + axis: New dimension along which to stack. A negative number counts from the end. +)"; + AppendKeywordArgumentDocs(doc, param_def...); + doc += R"( + +See also: + - :py:obj:`numpy.stack` + - :ref:`stack-driver` + - :py:obj:`tensorstore.overlay` + - :py:obj:`tensorstore.concat` + +Group: + Views +)"; + m.def( + "stack", + [](SequenceParameter< + std::variant> + python_layers, + DimensionIndex stack_dimension, + KeywordArgument... kwarg) -> TensorStore<> { + StackOpenOptions options; + ApplyKeywordArguments(options, kwarg...); + std::vector, Spec>> layers( + python_layers.size()); + for (size_t i = 0; i < layers.size(); ++i) { + std::visit([&](auto* obj) { layers[i] = obj->value; }, + python_layers[i]); + } + return ValueOrThrow(tensorstore::Stack(layers, stack_dimension, + std::move(options))); + }, + doc.c_str(), py::arg("layers"), py::arg("axis") = 0, py::kw_only(), + MakeKeywordArgumentPyArg(param_def)...); + } + + { + std::string doc = R"( +Virtually concatenates a sequence of :py:obj:`TensorStore` layers along an existing dimension. + + >>> store = ts.concat([ + ... ts.array([1, 2, 3, 4], dtype=ts.uint32), + ... ts.array([5, 6, 7, 8], dtype=ts.uint32) + ... ], + ... axis=0) + >>> store + TensorStore({ + 'context': {'data_copy_concurrency': {}}, + 'driver': 'stack', + 'dtype': 'uint32', + 'layers': [ + { + 'array': [1, 2, 3, 4], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': {'input_exclusive_max': [4], 'input_inclusive_min': [0]}, + }, + { + 'array': [5, 6, 7, 8], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [8], + 'input_inclusive_min': [4], + 'output': [{'input_dimension': 0, 'offset': -4}], + }, + }, + ], + 'schema': {'domain': {'exclusive_max': [8], 'inclusive_min': [0]}}, + 'transform': {'input_exclusive_max': [8], 'input_inclusive_min': [0]}, + }) + >>> await store.read() + array([1, 2, 3, 4, 5, 6, 7, 8], dtype=uint32) + >>> store = ts.concat([ + ... ts.array([[1, 2, 3], [4, 5, 6]], dtype=ts.uint32), + ... ts.array([[7, 8, 9], [10, 11, 12]], dtype=ts.uint32) + ... ], + ... axis=0) + >>> store + TensorStore({ + 'context': {'data_copy_concurrency': {}}, + 'driver': 'stack', + 'dtype': 'uint32', + 'layers': [ + { + 'array': [[1, 2, 3], [4, 5, 6]], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [2, 3], + 'input_inclusive_min': [0, 0], + }, + }, + { + 'array': [[7, 8, 9], [10, 11, 12]], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [4, 3], + 'input_inclusive_min': [2, 0], + 'output': [ + {'input_dimension': 0, 'offset': -2}, + {'input_dimension': 1}, + ], + }, + }, + ], + 'schema': {'domain': {'exclusive_max': [4, 3], 'inclusive_min': [0, 0]}}, + 'transform': {'input_exclusive_max': [4, 3], 'input_inclusive_min': [0, 0]}, + }) + >>> await store.read() + array([[ 1, 2, 3], + [ 4, 5, 6], + [ 7, 8, 9], + [10, 11, 12]], dtype=uint32) + >>> store = ts.concat([ + ... ts.array([[1, 2, 3], [4, 5, 6]], dtype=ts.uint32), + ... ts.array([[7, 8, 9], [10, 11, 12]], dtype=ts.uint32) + ... ], + ... axis=-1) + >>> store + TensorStore({ + 'context': {'data_copy_concurrency': {}}, + 'driver': 'stack', + 'dtype': 'uint32', + 'layers': [ + { + 'array': [[1, 2, 3], [4, 5, 6]], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [2, 3], + 'input_inclusive_min': [0, 0], + }, + }, + { + 'array': [[7, 8, 9], [10, 11, 12]], + 'driver': 'array', + 'dtype': 'uint32', + 'transform': { + 'input_exclusive_max': [2, 6], + 'input_inclusive_min': [0, 3], + 'output': [ + {'input_dimension': 0}, + {'input_dimension': 1, 'offset': -3}, + ], + }, + }, + ], + 'schema': {'domain': {'exclusive_max': [2, 6], 'inclusive_min': [0, 0]}}, + 'transform': {'input_exclusive_max': [2, 6], 'input_inclusive_min': [0, 0]}, + }) + >>> await store.read() + array([[ 1, 2, 3, 7, 8, 9], + [ 4, 5, 6, 10, 11, 12]], dtype=uint32) + >>> await ts.concat([ + ... ts.array([[1, 2, 3], [4, 5, 6]], dtype=ts.uint32).label["x", "y"], + ... ts.array([[7, 8, 9], [10, 11, 12]], dtype=ts.uint32) + ... ], + ... axis="y").read() + array([[ 1, 2, 3, 7, 8, 9], + [ 4, 5, 6, 10, 11, 12]], dtype=uint32) + +Args: + + layers: Sequence of layers to concatenate. If a layer is specified as a + :py:obj:`Spec` rather than a :py:obj:`TensorStore`, it must have a known + :py:obj:`~Spec.domain` and will be opened on-demand as needed for individual + read and write operations. + + axis: Existing dimension along which to concatenate. A negative number counts + from the end. May also be specified by a + :ref:`dimension label`. +)"; + AppendKeywordArgumentDocs(doc, param_def...); + doc += R"( + +See also: + - :py:obj:`numpy.concatenate` + - :ref:`stack-driver` + - :py:obj:`tensorstore.overlay` + - :py:obj:`tensorstore.stack` + +Group: + Views +)"; + m.def( + "concat", + [](SequenceParameter< + std::variant> + python_layers, + PythonDimensionIdentifier concat_dimension, + KeywordArgument... kwarg) -> TensorStore<> { + StackOpenOptions options; + ApplyKeywordArguments(options, kwarg...); + std::vector, Spec>> layers( + python_layers.size()); + for (size_t i = 0; i < layers.size(); ++i) { + std::visit([&](auto* obj) { layers[i] = obj->value; }, + python_layers[i]); + } + return ValueOrThrow(tensorstore::Concat( + layers, + internal_python::ToDimensionIdentifier(concat_dimension), + std::move(options))); + }, + doc.c_str(), py::arg("layers"), py::arg("axis"), py::kw_only(), + MakeKeywordArgumentPyArg(param_def)...); + } + }); + }); +} + +TENSORSTORE_GLOBAL_INITIALIZER { + RegisterPythonComponent(RegisterStackBindings, /*priority=*/-360); +} + +} // namespace +} // namespace internal_python +} // namespace tensorstore diff --git a/python/tensorstore/tensorstore_class.cc b/python/tensorstore/tensorstore_class.cc index 6d19d83fe..cd51fbd1a 100644 --- a/python/tensorstore/tensorstore_class.cc +++ b/python/tensorstore/tensorstore_class.cc @@ -104,85 +104,24 @@ WriteFutures IssueCopyOrWrite( } } -namespace open_setters { - -struct SetRead : public spec_setters::SetModeBase { - static constexpr const char* name = "read"; - static constexpr const char* doc = R"( -Allow read access. Defaults to `True` if neither ``read`` nor ``write`` is specified. -)"; -}; - -struct SetWrite : public spec_setters::SetModeBase { - static constexpr const char* name = "write"; - static constexpr const char* doc = R"( -Allow write access. Defaults to `True` if neither ``read`` nor ``write`` is specified. -)"; -}; - -using spec_setters::SetAssumeMetadata; -using spec_setters::SetCreate; -using spec_setters::SetDeleteExisting; -using spec_setters::SetOpen; -using spec_setters::SetOpenMode; - -struct SetContext { - using type = internal_context::ContextImplPtr; - static constexpr const char* name = "context"; - static constexpr const char* doc = R"( - -Shared resource context. Defaults to a new (unshared) context with default -options, as returned by :py:meth:`tensorstore.Context`. To share resources, -such as cache pools, between multiple open TensorStores, you must specify a -context. - -)"; - template - static absl::Status Apply(Self& self, type value) { - return self.Set(WrapImpl(std::move(value))); - } -}; - -struct SetTransaction { - using type = internal::TransactionState::CommitPtr; - static constexpr const char* name = "transaction"; - static constexpr const char* doc = R"( - -Transaction to use for opening/creating, and for subsequent operations. By -default, the open is non-transactional. - -.. note:: - - To perform transactional operations using a :py:obj:`TensorStore` that was - previously opened without a transaction, use - :py:obj:`TensorStore.with_transaction`. - -)"; - template - static absl::Status Apply(Self& self, type value) { - return self.Set( - internal::TransactionState::ToTransaction(std::move(value))); - } -}; - -} // namespace open_setters - constexpr auto ForwardOpenSetters = [](auto callback, auto... other_param) { WithSchemaKeywordArguments( callback, other_param..., open_setters::SetRead{}, open_setters::SetWrite{}, open_setters::SetOpenMode{}, open_setters::SetOpen{}, open_setters::SetCreate{}, open_setters::SetDeleteExisting{}, open_setters::SetAssumeMetadata{}, - open_setters::SetContext{}, open_setters::SetTransaction{}, - spec_setters::SetKvstore{}); + open_setters::SetAssumeCachedMetadata{}, open_setters::SetContext{}, + open_setters::SetTransaction{}, spec_setters::SetKvstore{}); }; constexpr auto ForwardSpecRequestSetters = [](auto callback, auto... other_param) { callback(other_param..., spec_setters::SetOpenMode{}, spec_setters::SetOpen{}, spec_setters::SetCreate{}, spec_setters::SetDeleteExisting{}, - spec_setters::SetAssumeMetadata{}, spec_setters::SetMinimalSpec{}, - spec_setters::SetRetainContext{}, spec_setters::SetUnbindContext{}); + spec_setters::SetAssumeMetadata{}, + spec_setters::SetAssumeCachedMetadata{}, + spec_setters::SetMinimalSpec{}, spec_setters::SetRetainContext{}, + spec_setters::SetUnbindContext{}); }; using TensorStoreCls = py::class_; diff --git a/python/tensorstore/tensorstore_class.h b/python/tensorstore/tensorstore_class.h index 4967ba61a..6543ead4b 100644 --- a/python/tensorstore/tensorstore_class.h +++ b/python/tensorstore/tensorstore_class.h @@ -25,6 +25,8 @@ // inclusion constraints are satisfied. #include "python/tensorstore/garbage_collection.h" +#include "python/tensorstore/spec.h" +#include "python/tensorstore/transaction.h" #include "tensorstore/tensorstore.h" #include "tensorstore/util/executor.h" @@ -40,6 +42,70 @@ struct PythonTensorStoreObject using PythonTensorStore = PythonTensorStoreObject::Handle; +namespace open_setters { + +struct SetRead : public spec_setters::SetModeBase { + static constexpr const char* name = "read"; + static constexpr const char* doc = R"( +Allow read access. Defaults to `True` if neither ``read`` nor ``write`` is specified. +)"; +}; + +struct SetWrite : public spec_setters::SetModeBase { + static constexpr const char* name = "write"; + static constexpr const char* doc = R"( +Allow write access. Defaults to `True` if neither ``read`` nor ``write`` is specified. +)"; +}; + +using spec_setters::SetAssumeCachedMetadata; +using spec_setters::SetAssumeMetadata; +using spec_setters::SetCreate; +using spec_setters::SetDeleteExisting; +using spec_setters::SetOpen; +using spec_setters::SetOpenMode; + +struct SetContext { + using type = internal_context::ContextImplPtr; + static constexpr const char* name = "context"; + static constexpr const char* doc = R"( + +Shared resource context. Defaults to a new (unshared) context with default +options, as returned by :py:meth:`tensorstore.Context`. To share resources, +such as cache pools, between multiple open TensorStores, you must specify a +context. + +)"; + template + static absl::Status Apply(Self& self, type value) { + return self.Set(WrapImpl(std::move(value))); + } +}; + +struct SetTransaction { + using type = internal::TransactionState::CommitPtr; + static constexpr const char* name = "transaction"; + static constexpr const char* doc = R"( + +Transaction to use for opening/creating, and for subsequent operations. By +default, the open is non-transactional. + +.. note:: + + To perform transactional operations using a :py:obj:`TensorStore` that was + previously opened without a transaction, use + :py:obj:`TensorStore.with_transaction`. + +)"; + template + static absl::Status Apply(Self& self, type value) { + return self.Set( + internal::TransactionState::ToTransaction(std::move(value))); + } +}; + +} // namespace open_setters + } // namespace internal_python } // namespace tensorstore diff --git a/python/tensorstore/tests/stack_test.py b/python/tensorstore/tests/stack_test.py new file mode 100644 index 000000000..b2f842680 --- /dev/null +++ b/python/tensorstore/tests/stack_test.py @@ -0,0 +1,43 @@ +# Copyright 2020 The TensorStore Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for tensorstore.overlay.""" + +import numpy as np +import tensorstore as ts + + +def test_overlay(): + a = ts.array([1, 2, 3, 4], dtype=ts.int32) + b = ts.array([5, 6, 7, 8], dtype=ts.int32) + c = ts.overlay([a, b.translate_to[3]]) + np.testing.assert_equal(np.array([1, 2, 3, 5, 6, 7, 8], dtype=np.int32), c) + + +def test_stack(): + a = ts.array([1, 2, 3, 4], dtype=ts.int32) + b = ts.array([5, 6, 7, 8], dtype=ts.int32) + c = ts.stack([a, b]) + np.testing.assert_equal( + np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=np.int32), c + ) + + +def test_concat(): + a = ts.array([[1, 2, 3], [4, 5, 6]], dtype=ts.int32) + b = ts.array([[7, 8, 9], [10, 11, 12]], dtype=ts.int32) + c = ts.concat([a, b], axis=0) + np.testing.assert_equal( + np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=np.int32), + c, + ) diff --git a/python/tensorstore/tests/tensorstore_test.py b/python/tensorstore/tests/tensorstore_test.py index 11d2e8a4d..f9d5eed03 100644 --- a/python/tensorstore/tests/tensorstore_test.py +++ b/python/tensorstore/tests/tensorstore_test.py @@ -435,3 +435,41 @@ async def test_tensorstore_ocdbt_zarr_repr(): create=True, ).result() repr(arr) + + +async def test_spec_open_mode(): + spec = ts.Spec({ + "driver": "zarr", + "kvstore": "memory://", + "schema": {"dtype": "uint32", "domain": {"shape": [100, 200]}}, + }) + + for open_mode_kwargs in [ + {"create": True}, + {"delete_existing": True, "create": True}, + {"open": True}, + {"open": True, "create": True}, + {"open": True, "assume_metadata": True}, + {"open": True, "assume_cached_metadata": True}, + ]: + spec_copy = spec.copy() + open_mode = ts.OpenMode(**open_mode_kwargs) + spec_copy.update(**open_mode_kwargs) + assert spec_copy.open_mode == open_mode + + spec_copy = spec.copy() + spec_copy.update(open_mode=open_mode) + assert spec_copy.open_mode == open_mode + + context = None + if open_mode == ts.OpenMode(open=True): + context = ts.Context() + await ts.open(spec, create=True, context=context) + store = await ts.open(spec, context=context, **open_mode_kwargs) + + requested_spec = store.spec(**open_mode_kwargs) + assert requested_spec.open_mode == open_mode + + store = await ts.open(spec, context=context, open_mode=open_mode) + requested_spec = store.spec(**open_mode_kwargs) + assert requested_spec.open_mode == open_mode diff --git a/tensorstore/BUILD b/tensorstore/BUILD index 9078e85fb..69a27b1a3 100644 --- a/tensorstore/BUILD +++ b/tensorstore/BUILD @@ -243,34 +243,14 @@ tensorstore_cc_test( ], ) -tensorstore_cc_library( +exports_files( + ["virtual_chunked.h"], + visibility = [":internal_packages"], +) + +alias( name = "virtual_chunked", - srcs = ["//tensorstore/driver/virtual_chunked:virtual_chunked.cc"], - hdrs = ["virtual_chunked.h"], - deps = [ - ":array", - ":box", - ":context", - ":staleness_bound", - ":tensorstore", - ":transaction", - "//tensorstore/driver", - "//tensorstore/driver:chunk_cache_driver", - "//tensorstore/index_space:index_transform", - "//tensorstore/internal:data_copy_concurrency_resource", - "//tensorstore/internal/cache:cache_pool_resource", - "//tensorstore/internal/cache:chunk_cache", - "//tensorstore/kvstore:generation", - "//tensorstore/serialization", - "//tensorstore/serialization:absl_time", - "//tensorstore/serialization:function", - "//tensorstore/util:executor", - "//tensorstore/util:option", - "//tensorstore/util/garbage_collection", - "@com_google_absl//absl/base:core_headers", - "@com_google_absl//absl/status", - "@com_google_absl//absl/time", - ], + actual = "//tensorstore/driver/virtual_chunked", ) tensorstore_cc_library( @@ -999,3 +979,13 @@ tensorstore_cc_test( "@com_google_googletest//:gtest_main", ], ) + +exports_files( + ["stack.h"], + visibility = ["//tensorstore:internal_packages"], +) + +alias( + name = "stack", + actual = "//tensorstore/driver/stack", +) diff --git a/tensorstore/chunk_layout.cc b/tensorstore/chunk_layout.cc index dbb45809f..2bd5b6b46 100644 --- a/tensorstore/chunk_layout.cc +++ b/tensorstore/chunk_layout.cc @@ -179,7 +179,7 @@ struct ChunkLayoutData { Index chunk_elements_[kNumUsages] = {kImplicit, kImplicit, kImplicit}; }; -bool IsHardConstraint(ChunkLayoutData& impl, HardConstraintBit bit) { +bool IsHardConstraint(const ChunkLayoutData& impl, HardConstraintBit bit) { return impl.hard_constraint_[static_cast(bit)]; } @@ -321,6 +321,15 @@ void ClearHardConstraintBits(Storage& impl) { } } +bool HasAnyHardConstraints(const Storage& impl) { + if (IsHardConstraint(impl, HardConstraintBit::inner_order)) return true; + if (impl.grid_origin_hard_constraint_) return true; + for (int i = 0; i < kNumUsages; ++i) { + if (impl.chunk_shape_hard_constraint_[i]) return true; + } + return false; +} + absl::Status RankMismatchError(DimensionIndex new_rank, DimensionIndex existing_rank) { return absl::InvalidArgumentError(tensorstore::StrCat( @@ -698,6 +707,11 @@ DimensionIndex ChunkLayout::rank() const { return dynamic_rank; } +bool ChunkLayout::HasHardConstraints() const { + if (!storage_) return false; + return HasAnyHardConstraints(*storage_); +} + absl::Status ChunkLayout::Set(RankConstraint value) { if (value.rank == dynamic_rank) return absl::OkStatus(); StoragePtr storage_to_be_destroyed; diff --git a/tensorstore/chunk_layout.h b/tensorstore/chunk_layout.h index 62cd2799f..907b88465 100644 --- a/tensorstore/chunk_layout.h +++ b/tensorstore/chunk_layout.h @@ -703,6 +703,9 @@ class ChunkLayout { /// Returns the rank constraint, or `dynamic_rank` if unspecified. DimensionIndex rank() const; + /// Returns `true` if any hard constraints, other than rank, are specified. + bool HasHardConstraints() const; + /// Sets `box` to the precise write/read chunk template. /// /// For the purpose of this method, only hard constraints on `grid_origin()` diff --git a/tensorstore/chunk_layout_test.cc b/tensorstore/chunk_layout_test.cc index 1abff4f1b..88ffb3f78 100644 --- a/tensorstore/chunk_layout_test.cc +++ b/tensorstore/chunk_layout_test.cc @@ -1790,4 +1790,68 @@ TEST(ChooseReadWriteChunkShapesTest, WriteShapeConstrained) { EXPECT_THAT(write_chunk_shape, ::testing::ElementsAre(10, 14)); } +TEST(HasHardConstraints, Basic) { + ChunkLayout layout; + EXPECT_FALSE(layout.HasHardConstraints()); + TENSORSTORE_ASSERT_OK(layout.Set(tensorstore::RankConstraint{2})); + EXPECT_FALSE(layout.HasHardConstraints()); + + { + auto layout1 = layout; + TENSORSTORE_ASSERT_OK(layout1.Set( + tensorstore::ChunkLayout::InnerOrder({0, 1}, + /*hard_constraint=*/false))); + EXPECT_FALSE(layout1.HasHardConstraints()); + } + { + auto layout1 = layout; + TENSORSTORE_ASSERT_OK(layout1.Set( + tensorstore::ChunkLayout::InnerOrder({0, 1}, + /*hard_constraint=*/true))); + EXPECT_TRUE(layout1.HasHardConstraints()); + } + + { + auto layout1 = layout; + TENSORSTORE_ASSERT_OK(layout1.Set( + tensorstore::ChunkLayout::GridOrigin({100, 200}, + /*hard_constraint=*/false))); + EXPECT_FALSE(layout1.HasHardConstraints()); + } + { + auto layout1 = layout; + TENSORSTORE_ASSERT_OK( + layout1.Set(tensorstore::ChunkLayout::GridOrigin({100, 200}))); + EXPECT_TRUE(layout1.HasHardConstraints()); + } + + { + auto layout1 = layout; + TENSORSTORE_ASSERT_OK(layout1.Set( + tensorstore::ChunkLayout::ReadChunkShape({100, 200}, + /*hard_constraint=*/false))); + EXPECT_FALSE(layout1.HasHardConstraints()); + } + { + auto layout1 = layout; + TENSORSTORE_ASSERT_OK( + layout1.Set(tensorstore::ChunkLayout::ReadChunkShape({100, 200}))); + EXPECT_TRUE(layout1.HasHardConstraints()); + } + + { + auto layout1 = layout; + TENSORSTORE_ASSERT_OK( + layout1.Set(tensorstore::ChunkLayout::ReadChunkAspectRatio({1, 1}))); + EXPECT_FALSE(layout1.HasHardConstraints()); + } + + { + auto layout1 = layout; + TENSORSTORE_ASSERT_OK( + layout1.Set(tensorstore::ChunkLayout::ReadChunkElements(200))); + EXPECT_FALSE(layout1.HasHardConstraints()); + } +} + } // namespace diff --git a/tensorstore/driver/driver.cc b/tensorstore/driver/driver.cc index fb428e5c8..4edbb0475 100644 --- a/tensorstore/driver/driver.cc +++ b/tensorstore/driver/driver.cc @@ -273,6 +273,16 @@ Result GetTransformedDriverSpec( return transformed_driver_spec; } +absl::Status SetReadWriteMode(DriverHandle& handle, ReadWriteMode new_mode) { + if (new_mode != ReadWriteMode::dynamic) { + auto existing_mode = handle.driver.read_write_mode(); + TENSORSTORE_RETURN_IF_ERROR( + internal::ValidateSupportsModes(existing_mode, new_mode)); + handle.driver.set_read_write_mode(new_mode); + } + return absl::OkStatus(); +} + bool DriverHandleNonNullSerializer::Encode(serialization::EncodeSink& sink, const DriverHandle& value) { assert(value.driver); diff --git a/tensorstore/driver/driver.h b/tensorstore/driver/driver.h index 69a330f5a..0c718e35d 100644 --- a/tensorstore/driver/driver.h +++ b/tensorstore/driver/driver.h @@ -302,6 +302,12 @@ Future GetStorageStatistics( Result GetTransformedDriverSpec( const DriverHandle& handle, SpecRequestOptions&& options); +// Updates the read-write mode, or returns an error if incompatible with the +// existing mode. +// +// If `new_mode == ReadWriteMode::dynamic`, the existing mode is unchanged. +absl::Status SetReadWriteMode(DriverHandle& handle, ReadWriteMode new_mode); + struct DriverHandleNonNullSerializer { [[nodiscard]] static bool Encode(serialization::EncodeSink& sink, const DriverHandle& value); diff --git a/tensorstore/driver/driver_spec.cc b/tensorstore/driver/driver_spec.cc index 9cd1d72dd..f712212c3 100644 --- a/tensorstore/driver/driver_spec.cc +++ b/tensorstore/driver/driver_spec.cc @@ -102,25 +102,25 @@ absl::Status ApplyOptions(DriverSpec::Ptr& spec, SpecOptions&& options) { } namespace { -absl::Status MaybeDeriveTransform(TransformedDriverSpec& spec) { - TENSORSTORE_ASSIGN_OR_RETURN(auto domain, spec.driver_spec->GetDomain()); +absl::Status MaybeDeriveTransform(DriverSpec::Ptr& driver_spec, + IndexTransform<>& transform) { + TENSORSTORE_ASSIGN_OR_RETURN(auto domain, driver_spec->GetDomain()); if (domain.valid()) { - spec.transform = IdentityTransform(domain); + transform = IdentityTransform(domain); } return absl::OkStatus(); } } // namespace -absl::Status TransformAndApplyOptions(TransformedDriverSpec& spec, +absl::Status TransformAndApplyOptions(DriverSpec::Ptr& driver_spec, + IndexTransform<>& transform, SpecOptions&& options) { const bool should_get_transform = - !spec.transform.valid() && options.domain().valid(); - TENSORSTORE_RETURN_IF_ERROR( - options.TransformInputSpaceSchema(spec.transform)); - TENSORSTORE_RETURN_IF_ERROR( - ApplyOptions(spec.driver_spec, std::move(options))); + !transform.valid() && options.domain().valid(); + TENSORSTORE_RETURN_IF_ERROR(options.TransformInputSpaceSchema(transform)); + TENSORSTORE_RETURN_IF_ERROR(ApplyOptions(driver_spec, std::move(options))); if (should_get_transform) { - TENSORSTORE_RETURN_IF_ERROR(MaybeDeriveTransform(spec)); + TENSORSTORE_RETURN_IF_ERROR(MaybeDeriveTransform(driver_spec, transform)); } return absl::OkStatus(); } @@ -307,7 +307,7 @@ TENSORSTORE_DEFINE_JSON_BINDER( registry.RegisteredObjectBinder()), jb::Initialize([](auto* obj) { if (obj->transform.valid()) return absl::OkStatus(); - return MaybeDeriveTransform(*obj); + return MaybeDeriveTransform(obj->driver_spec, obj->transform); })))(is_loading, options, obj, j); }) diff --git a/tensorstore/driver/driver_spec.h b/tensorstore/driver/driver_spec.h index 6185d5526..ddff35d2f 100644 --- a/tensorstore/driver/driver_spec.h +++ b/tensorstore/driver/driver_spec.h @@ -257,8 +257,14 @@ struct TransformedDriverSpec { }; absl::Status ApplyOptions(DriverSpec::Ptr& spec, SpecOptions&& options); -absl::Status TransformAndApplyOptions(TransformedDriverSpec& spec, +absl::Status TransformAndApplyOptions(DriverSpec::Ptr& driver_spec, + IndexTransform<>& transform, SpecOptions&& options); +inline absl::Status TransformAndApplyOptions(TransformedDriverSpec& spec, + SpecOptions&& options) { + return internal::TransformAndApplyOptions(spec.driver_spec, spec.transform, + std::move(options)); +} OpenMode GetOpenMode(const TransformedDriverSpec& spec); diff --git a/tensorstore/driver/stack/BUILD b/tensorstore/driver/stack/BUILD index 983ccf881..a6ce32a3d 100644 --- a/tensorstore/driver/stack/BUILD +++ b/tensorstore/driver/stack/BUILD @@ -22,21 +22,31 @@ filegroup( tensorstore_cc_library( name = "stack", - srcs = ["driver.cc"], + srcs = [ + "driver.cc", + "driver.h", + ], + hdrs = ["//tensorstore:stack.h"], deps = [ - "//tensorstore:array", + "//tensorstore", "//tensorstore:box", "//tensorstore:context", "//tensorstore:data_type", "//tensorstore:index", "//tensorstore:index_interval", "//tensorstore:open_mode", + "//tensorstore:open_options", + "//tensorstore:rank", "//tensorstore:resize_options", "//tensorstore:schema", + "//tensorstore:spec", "//tensorstore:transaction", "//tensorstore/driver", "//tensorstore/driver:chunk", "//tensorstore/driver:chunk_receiver_utils", + "//tensorstore/index_space:dim_expression", + "//tensorstore/index_space:dimension_identifier", + "//tensorstore/index_space:dimension_units", "//tensorstore/index_space:index_transform", "//tensorstore/internal:concurrency_resource", "//tensorstore/internal:context_binding", @@ -48,32 +58,31 @@ tensorstore_cc_library( "//tensorstore/internal:type_traits", "//tensorstore/internal/cache:cache_pool_resource", "//tensorstore/internal/json_binding", - "//tensorstore/internal/json_binding:bindable", "//tensorstore/internal/json_binding:staleness_bound", "//tensorstore/serialization", "//tensorstore/util:executor", "//tensorstore/util:future", - "//tensorstore/util:iterate", "//tensorstore/util:iterate_over_index_range", + "//tensorstore/util:option", "//tensorstore/util:result", "//tensorstore/util:span", "//tensorstore/util:status", "//tensorstore/util:str_cat", - "//tensorstore/util/execution", "//tensorstore/util/execution:any_receiver", + "//tensorstore/util/execution:flow_sender_operation_state", "//tensorstore/util/garbage_collection", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/hash", "@com_google_absl//absl/log:absl_log", "@com_google_absl//absl/status", + "@com_google_absl//absl/strings:str_format", ], alwayslink = True, ) tensorstore_cc_test( - name = "stack_test", + name = "driver_test", size = "small", srcs = ["driver_test.cc"], deps = [ @@ -81,21 +90,29 @@ tensorstore_cc_test( "//tensorstore", "//tensorstore:array", "//tensorstore:array_testutil", + "//tensorstore:box", "//tensorstore:context", "//tensorstore:data_type", "//tensorstore:index", "//tensorstore:open", + "//tensorstore:open_mode", "//tensorstore:progress", + "//tensorstore:schema", "//tensorstore:strided_layout", - "//tensorstore/driver/array", # build_cleaner: keep - "//tensorstore/driver/n5", # build_cleaner: keep + "//tensorstore:transaction", + "//tensorstore/driver/array", + "//tensorstore/driver/n5", + "//tensorstore/driver/zarr3", "//tensorstore/index_space:dim_expression", + "//tensorstore/index_space:index_transform", "//tensorstore/internal:json_gtest", - "//tensorstore/kvstore/memory", # build_cleaner: keep - "//tensorstore/util:future", + "//tensorstore/kvstore", + "//tensorstore/kvstore/memory", "//tensorstore/util:result", "//tensorstore/util:status_testutil", + "//tensorstore/util:unit", "@com_github_nlohmann_json//:nlohmann_json", + "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest_main", ], @@ -107,22 +124,25 @@ tensorstore_cc_test( deps = [ ":stack", # build_cleaner: keep "//tensorstore", + "//tensorstore:array", "//tensorstore:array_testutil", + "//tensorstore:box", "//tensorstore:context", "//tensorstore:data_type", "//tensorstore:index", "//tensorstore:open", + "//tensorstore:open_mode", "//tensorstore:progress", "//tensorstore/driver/image/png", # build_cleaner: keep "//tensorstore/driver/neuroglancer_precomputed", # build_cleaner: keep "//tensorstore/index_space:dim_expression", "//tensorstore/kvstore", "//tensorstore/kvstore/memory", # build_cleaner: keep - "//tensorstore/util:future", "//tensorstore/util:result", "//tensorstore/util:status", "//tensorstore/util:status_testutil", "@com_github_nlohmann_json//:nlohmann_json", + "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:cord", "@com_google_googletest//:gtest_main", diff --git a/tensorstore/driver/stack/driver.cc b/tensorstore/driver/stack/driver.cc index 2727ee09c..9b33fa8d1 100644 --- a/tensorstore/driver/stack/driver.cc +++ b/tensorstore/driver/stack/driver.cc @@ -18,19 +18,17 @@ #include #include -#include #include -#include +#include #include #include #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" -#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/absl_log.h" #include "absl/status/status.h" -#include "tensorstore/array.h" +#include "absl/strings/str_format.h" #include "tensorstore/box.h" #include "tensorstore/context.h" #include "tensorstore/data_type.h" @@ -39,8 +37,12 @@ #include "tensorstore/driver/driver_handle.h" #include "tensorstore/driver/driver_spec.h" #include "tensorstore/driver/registry.h" +#include "tensorstore/driver/stack/driver.h" #include "tensorstore/index.h" #include "tensorstore/index_interval.h" +#include "tensorstore/index_space/dim_expression.h" +#include "tensorstore/index_space/dimension_identifier.h" +#include "tensorstore/index_space/dimension_units.h" #include "tensorstore/index_space/index_domain.h" #include "tensorstore/index_space/index_transform.h" #include "tensorstore/index_space/internal/propagate_bounds.h" @@ -50,21 +52,21 @@ #include "tensorstore/internal/grid_partition.h" #include "tensorstore/internal/intrusive_ptr.h" #include "tensorstore/internal/irregular_grid.h" -#include "tensorstore/internal/json_binding/bindable.h" #include "tensorstore/internal/json_binding/json_binding.h" #include "tensorstore/internal/json_binding/staleness_bound.h" // IWYU pragma: keep #include "tensorstore/internal/json_binding/std_array.h" // IWYU pragma: keep #include "tensorstore/internal/tagged_ptr.h" #include "tensorstore/internal/type_traits.h" #include "tensorstore/open_mode.h" +#include "tensorstore/open_options.h" +#include "tensorstore/rank.h" #include "tensorstore/resize_options.h" #include "tensorstore/schema.h" #include "tensorstore/transaction.h" #include "tensorstore/util/execution/any_receiver.h" -#include "tensorstore/util/execution/execution.h" +#include "tensorstore/util/execution/flow_sender_operation_state.h" #include "tensorstore/util/executor.h" #include "tensorstore/util/future.h" -#include "tensorstore/util/iterate.h" #include "tensorstore/util/iterate_over_index_range.h" #include "tensorstore/util/result.h" #include "tensorstore/util/span.h" @@ -75,6 +77,7 @@ #include "tensorstore/internal/context_binding_vector.h" // IWYU pragma: keep #include "tensorstore/serialization/std_vector.h" // IWYU pragma: keep #include "tensorstore/util/garbage_collection/fwd.h" // IWYU pragma: keep +#include "tensorstore/util/garbage_collection/std_optional.h" // IWYU pragma: keep #include "tensorstore/util/garbage_collection/std_vector.h" // IWYU pragma: keep namespace tensorstore { @@ -90,6 +93,41 @@ using ::tensorstore::internal::ReadChunk; using ::tensorstore::internal::TransformedDriverSpec; using ::tensorstore::internal::WriteChunk; +// Specifies a stack layer as either a `DriverSpec` (to be opened on demand) or +// an open `Driver`, along with a transform. +struct StackLayer { + // Index transform, must not be null. + IndexTransform<> transform; + + // Driver spec, if layer is to be opened on demand. + internal::DriverSpecPtr driver_spec; + + // Driver, if layer is already open. + internal::ReadWritePtr driver; + + bool is_open() const { return static_cast(driver); } + + internal::DriverHandle GetDriverHandle(Transaction transaction) const { + assert(is_open()); + return {driver, transform, transaction}; + } + + internal::DriverHandle GetDriverHandle( + internal::OpenTransactionPtr& transaction) const { + return GetDriverHandle( + internal::TransactionState::ToTransaction(transaction)); + } + + internal::TransformedDriverSpec GetTransformedDriverSpec() const { + assert(!is_open()); + return internal::TransformedDriverSpec{driver_spec, transform}; + } + + constexpr static auto ApplyMembers = [](auto&& x, auto f) { + return f(x.transform, x.driver_spec, x.driver); + }; +}; + // NOTES: // // The stack driver would benefit from being able to hold a weak pointer to a @@ -110,11 +148,6 @@ using ::tensorstore::internal::WriteChunk; namespace jb = tensorstore::internal_json_binding; -absl::Status TransactionError() { - return absl::UnimplementedError( - "\"stack\" driver does not support transactions"); -} - /// Used to index individual cells struct Cell { std::vector points; @@ -156,6 +189,210 @@ struct CellEq { } }; +// Certain operations are applied to either a sequence of +// `internal::TransformedDriverSpec` used by the `StackDriverSpec` to represent +// layers, or to a sequence of `StackLayer` used by the open `StackDriver` to +// represent layers. +template +constexpr bool IsStackLayerLike = false; + +template <> +constexpr bool IsStackLayerLike = true; + +template <> +constexpr bool IsStackLayerLike = true; + +template +absl::Status ForEachLayer(size_t num_layers, Callback&& callback) { + for (size_t layer_i = 0; layer_i < num_layers; ++layer_i) { + absl::Status status = callback(layer_i); + if (!status.ok()) { + return tensorstore::MaybeAnnotateStatus( + status, absl::StrFormat("Layer %d", layer_i)); + } + } + return absl::OkStatus(); +} + +using internal::GetEffectiveDomain; + +Result> GetEffectiveDomain(const StackLayer& layer) { + return layer.is_open() + ? layer.transform.domain() + : internal::GetEffectiveDomain(layer.GetTransformedDriverSpec()); +} + +template +Result>> GetEffectiveDomainsForLayers( + span layers) { + static_assert(IsStackLayerLike); + std::vector> domains; + domains.reserve(layers.size()); + DimensionIndex rank; + auto status = ForEachLayer(layers.size(), [&](size_t layer_i) { + TENSORSTORE_ASSIGN_OR_RETURN( + auto effective_domain, + internal_stack::GetEffectiveDomain(layers[layer_i])); + if (!effective_domain.valid()) { + return absl::InvalidArgumentError( + tensorstore::StrCat("Domain must be specified")); + } + domains.emplace_back(std::move(effective_domain)); + // validate rank. + if (layer_i == 0) { + rank = domains.back().rank(); + } else if (domains.back().rank() != rank) { + return absl::InvalidArgumentError(tensorstore::StrCat( + "Layer domain ", domains.back(), " of rank ", domains.back().rank(), + " does not match layer 0 rank of ", rank)); + } + return absl::OkStatus(); + }); + if (!status.ok()) return status; + return domains; +} + +Result> GetCombinedDomain( + const Schema& schema, span> layer_domains) { + // Each layer is expected to have an effective domain so that each layer + // can be queried when resolving chunks. + // The overall domain is Hull(layer.domain...) + IndexDomain<> domain; + auto status = ForEachLayer(layer_domains.size(), [&](size_t layer_i) { + TENSORSTORE_ASSIGN_OR_RETURN( + domain, HullIndexDomains(domain, layer_domains[layer_i])); + return absl::OkStatus(); + }); + if (!status.ok()) return status; + TENSORSTORE_ASSIGN_OR_RETURN( + domain, ConstrainIndexDomain(schema.domain(), std::move(domain))); + // stack disallows resize, so mark all dimensions as explicit. + return WithImplicitDimensions(std::move(domain), false, false); +} + +using internal::GetEffectiveDimensionUnits; + +Result GetEffectiveDimensionUnits( + const StackLayer& layer) { + if (layer.is_open()) { + TENSORSTORE_ASSIGN_OR_RETURN(auto units, layer.driver->GetDimensionUnits()); + return tensorstore::TransformOutputDimensionUnits(layer.transform, + std::move(units)); + } else { + return internal::GetEffectiveDimensionUnits( + layer.GetTransformedDriverSpec()); + } +} + +template +Result GetDimensionUnits(const Schema& schema, + span layers) { + static_assert(IsStackLayerLike); + // Retrieve dimension units from schema. These take precedence over computed + // dimensions, so disallow further assignment. + DimensionUnitsVector dimension_units(schema.dimension_units()); + DimensionSet allow_assignment(true); + for (size_t i = 0; i < dimension_units.size(); i++) { + if (dimension_units[i].has_value()) { + allow_assignment[i] = false; + } + } + // Merge dimension units from the layers. If there are conflicting values, + // clear and disallow further assignment. + auto status = ForEachLayer(layers.size(), [&](size_t layer_i) { + const auto& d = layers[layer_i]; + TENSORSTORE_ASSIGN_OR_RETURN(auto du, + internal_stack::GetEffectiveDimensionUnits(d)); + if (du.size() > dimension_units.size()) { + dimension_units.resize(du.size()); + } + for (size_t i = 0; i < du.size(); i++) { + if (!allow_assignment[i]) continue; + if (!du[i].has_value()) continue; + if (!dimension_units[i].has_value()) { + dimension_units[i] = du[i]; + } else if (dimension_units[i].value() != du[i].value()) { + // mismatch; clear and disallow future assignment. + dimension_units[i] = std::nullopt; + allow_assignment[i] = false; + } + } + return absl::OkStatus(); + }); + if (!status.ok()) return status; + return dimension_units; +} + +using internal::TransformAndApplyOptions; + +absl::Status TransformAndApplyOptions(StackLayer& layer, + SpecOptions&& options) { + assert(!layer.is_open()); + return internal::TransformAndApplyOptions(layer.driver_spec, layer.transform, + std::move(options)); +} + +template +absl::Status ApplyLayerOptions(span layers, Schema& schema, + const SpecOptions& options) { + if (&schema != &options) { + TENSORSTORE_RETURN_IF_ERROR(schema.Set(options.rank())); + TENSORSTORE_RETURN_IF_ERROR(schema.Set(options.dtype())); + TENSORSTORE_RETURN_IF_ERROR(schema.Set(options.domain())); + TENSORSTORE_RETURN_IF_ERROR(schema.Set(options.dimension_units())); + } + if (options.codec().valid()) { + return absl::InvalidArgumentError( + "codec option not supported by \"stack\" driver"); + } + if (options.fill_value().valid()) { + return absl::InvalidArgumentError( + "fill value option not supported by \"stack\" driver"); + } + if (options.kvstore.valid()) { + return absl::InvalidArgumentError( + "kvstore option not supported by \"stack\" driver"); + } + if (options.chunk_layout().HasHardConstraints()) { + return absl::InvalidArgumentError( + "chunk layout option not supported by \"stack\" driver"); + } + return ForEachLayer(layers.size(), [&](size_t layer_i) { + auto& layer = layers[layer_i]; + if constexpr (std::is_same_v) { + if (layer.is_open()) { + if (options.open_mode != OpenMode{} && + !(options.open_mode & OpenMode::open)) { + return absl::InvalidArgumentError(tensorstore::StrCat( + "Open mode of ", options.open_mode, + " is not compatible with already-open layer")); + } + if (options.recheck_cached_data.specified() || + options.recheck_cached_metadata.specified()) { + return absl::InvalidArgumentError( + "Cannot specify cache rechecking options with already-open " + "layer"); + } + return absl::OkStatus(); + } + } + // Filter the options to only those that we wish to pass on to the + // layers, otherwise we may be passing on nonsensical settings for a + // layer. + SpecOptions layer_options; + layer_options.open_mode = options.open_mode; + layer_options.recheck_cached_data = options.recheck_cached_data; + layer_options.recheck_cached_metadata = options.recheck_cached_metadata; + layer_options.minimal_spec = options.minimal_spec; + TENSORSTORE_RETURN_IF_ERROR( + static_cast(layer_options).Set(schema.dtype())); + TENSORSTORE_RETURN_IF_ERROR( + static_cast(layer_options).Set(schema.rank())); + return internal_stack::TransformAndApplyOptions(layers[layer_i], + std::move(layer_options)); + }); +} + class StackDriverSpec : public internal::RegisteredDriverSpec { @@ -163,7 +400,7 @@ class StackDriverSpec constexpr static char id[] = "stack"; Context::Resource data_copy_concurrency; - std::vector layers; + std::vector layers; constexpr static auto ApplyMembers = [](auto&& x, auto f) { return f(internal::BaseCast(x), @@ -171,17 +408,19 @@ class StackDriverSpec }; absl::Status InitializeLayerRankAndDtype() { - if (layers.empty()) { - return absl::InvalidArgumentError("\"stack\" driver spec has no layers"); - } - /// Set the schema rank and dtype based on the layers. - for (auto& d : layers) { - TENSORSTORE_RETURN_IF_ERROR( - this->schema.Set(RankConstraint{internal::GetRank(d)})); + return ForEachLayer(layers.size(), [&](size_t layer_i) { + auto& layer = layers[layer_i]; + DimensionIndex layer_rank = internal::GetRank(layer); + if (schema.rank() != dynamic_rank && layer_rank != schema.rank()) { + return absl::InvalidArgumentError(tensorstore::StrCat( + "Rank of ", layer_rank, " does not match existing rank of ", + schema.rank())); + } + schema.Set(RankConstraint{layer_rank}).IgnoreError(); TENSORSTORE_RETURN_IF_ERROR( - this->schema.Set(d.driver_spec->schema.dtype())); - } - return absl::OkStatus(); + schema.Set(layer.driver_spec->schema.dtype())); + return absl::OkStatus(); + }); } constexpr static auto default_json_binder = jb::Sequence( @@ -196,28 +435,8 @@ class StackDriverSpec })); absl::Status ApplyOptions(SpecOptions&& options) override { - if (options.codec().valid()) { - return absl::InvalidArgumentError( - "\"codec\" not supported by \"stack\" driver"); - } - if (options.fill_value().valid()) { - return absl::InvalidArgumentError( - "\"fill_value\" not supported by \"stack\" driver"); - } - for (auto& d : layers) { - // Filter the options to only those that we wish to pass on to the - // layers, otherwise we may be passing on nonsensical settings for a - // layer. - SpecOptions o; - o.open_mode = options.open_mode; - o.recheck_cached_data = options.recheck_cached_data; - o.recheck_cached_metadata = options.recheck_cached_metadata; - TENSORSTORE_RETURN_IF_ERROR(static_cast(o).Set(schema.dtype())); - TENSORSTORE_RETURN_IF_ERROR(static_cast(o).Set(schema.rank())); - TENSORSTORE_RETURN_IF_ERROR( - internal::TransformAndApplyOptions(d, std::move(o))); - } - return schema.Set(static_cast(options)); + return internal_stack::ApplyLayerOptions( + layers, schema, std::move(options)); } OpenMode open_mode() const override { @@ -231,78 +450,16 @@ class StackDriverSpec return prev_mode; } - Result>> GetEffectiveDomainsForLayers() const { - assert(!layers.empty()); - std::vector> domains; - domains.reserve(layers.size()); - DimensionIndex rank; - for (size_t i = 0; i < layers.size(); i++) { - TENSORSTORE_ASSIGN_OR_RETURN(auto effective_domain, - internal::GetEffectiveDomain(layers[i])); - if (!effective_domain.valid()) { - return absl::InvalidArgumentError( - tensorstore::StrCat("layer[", i, "] domain is unspecified")); - } - domains.emplace_back(std::move(effective_domain)); - // validate rank. - if (i == 0) { - rank = domains.back().rank(); - } else if (domains.back().rank() != rank) { - return absl::InvalidArgumentError( - tensorstore::StrCat("layer[", i, "] rank mismatch")); - } - } - return domains; - } - Result> GetDomain() const override { - // Each layer is expected to have an effective domain so that each layer - // can be queried when resolving chunks. - // The overall domain is Hull(layer.domain...) - IndexDomain<> domain; - { - TENSORSTORE_ASSIGN_OR_RETURN(auto domains, - GetEffectiveDomainsForLayers()); - for (auto& d : domains) { - TENSORSTORE_ASSIGN_OR_RETURN(domain, HullIndexDomains(domain, d)); - } - } - TENSORSTORE_ASSIGN_OR_RETURN( - domain, ConstrainIndexDomain(schema.domain(), std::move(domain))); - // stack disallows resize, so mark all dimensions as explicit. - return WithImplicitDimensions(std::move(domain), false, false); + TENSORSTORE_ASSIGN_OR_RETURN(auto layer_domains, + internal_stack::GetEffectiveDomainsForLayers< + internal::TransformedDriverSpec>(layers)); + return internal_stack::GetCombinedDomain(schema, layer_domains); } Result GetDimensionUnits() const override { - /// Retrieve dimension units from schema. These take precedence - /// over computed dimensions, so disallow further assignment. - DimensionUnitsVector dimension_units(schema.dimension_units()); - DimensionSet allow_assignment(true); - for (size_t i = 0; i < dimension_units.size(); i++) { - if (dimension_units[i].has_value()) { - allow_assignment[i] = false; - } - } - /// Merge dimension units from the layers. If there are conflicting - /// values, clear and disallow further assignment. - for (const auto& d : layers) { - TENSORSTORE_ASSIGN_OR_RETURN(auto du, GetEffectiveDimensionUnits(d)); - if (du.size() > dimension_units.size()) { - dimension_units.resize(du.size()); - } - for (size_t i = 0; i < du.size(); i++) { - if (!allow_assignment[i]) continue; - if (!du[i].has_value()) continue; - if (!dimension_units[i].has_value()) { - dimension_units[i] = du[i]; - } else if (dimension_units[i].value() != du[i].value()) { - // mismatch; clear and disallow future assignment. - dimension_units[i] = std::nullopt; - allow_assignment[i] = false; - } - } - } - return dimension_units; + return internal_stack::GetDimensionUnits( + schema, layers); } Future Open( @@ -314,18 +471,11 @@ class StackDriver : public internal::RegisteredDriver { public: - explicit StackDriver(StackDriverSpec bound_spec) - : bound_spec_(std::move(bound_spec)) { - assert(bound_spec_.context_binding_state_ == ContextBindingState::bound); - dimension_units_ = bound_spec_.GetDimensionUnits().value_or( - DimensionUnitsVector(bound_spec_.schema.rank())); - } - - DataType dtype() override { return bound_spec_.schema.dtype(); } - DimensionIndex rank() override { return bound_spec_.schema.rank(); } + DataType dtype() override { return dtype_; } + DimensionIndex rank() override { return layer_domain_.rank(); } Executor data_copy_executor() override { - return bound_spec_.data_copy_concurrency->executor; + return data_copy_concurrency_->executor; } Result GetDimensionUnits() override { @@ -347,14 +497,17 @@ class StackDriver void Write(OpenTransactionPtr transaction, IndexTransform<> transform, WriteChunkReceiver receiver) override; - absl::Status InitializeGridIndices(const std::vector>& domains); + absl::Status InitializeGridIndices(span> domains); constexpr static auto ApplyMembers = [](auto&& x, auto f) { // Exclude `context_binding_state_` because it is handled specially. - return f(x.bound_spec_); + return f(x.dtype_, x.data_copy_concurrency_, x.layers_, x.dimension_units_, + x.layer_domain_); }; - StackDriverSpec bound_spec_; + DataType dtype_; + Context::Resource data_copy_concurrency_; + std::vector layers_; DimensionUnitsVector dimension_units_; IndexDomain<> layer_domain_; IrregularGrid grid_; @@ -362,27 +515,50 @@ class StackDriver absl::flat_hash_map grid_to_layer_; }; +Result MakeStackDriverHandle( + internal::ReadWritePtr driver, + span> layer_domains, Transaction transaction, + const Schema& schema) { + driver->dtype_ = schema.dtype(); + TENSORSTORE_ASSIGN_OR_RETURN( + driver->layer_domain_, + internal_stack::GetCombinedDomain(schema, layer_domains)); + TENSORSTORE_RETURN_IF_ERROR(driver->InitializeGridIndices(layer_domains)); + auto transform = IdentityTransform(driver->layer_domain_); + driver->dimension_units_ = + internal_stack::GetDimensionUnits(schema, driver->layers_) + .value_or(DimensionUnitsVector(transform.input_rank())); + return internal::DriverHandle{std::move(driver), std::move(transform), + std::move(transaction)}; +} + Future StackDriverSpec::Open( internal::OpenTransactionPtr transaction, ReadWriteMode read_write_mode) const { - if (transaction) return TransactionError(); + if (!schema.dtype().valid()) { + return absl::InvalidArgumentError("dtype must be specified"); + } if (read_write_mode == ReadWriteMode::dynamic) { read_write_mode = ReadWriteMode::read_write; } - if (!schema.dtype().valid()) { - return absl::InvalidArgumentError( - "Unable to infer \"dtype\" in \"stack\" driver"); + auto driver = internal::MakeReadWritePtr(read_write_mode); + driver->data_copy_concurrency_ = data_copy_concurrency; + const size_t num_layers = layers.size(); + driver->layers_.resize(num_layers); + for (size_t layer_i = 0; layer_i < num_layers; ++layer_i) { + auto& layer_spec = layers[layer_i]; + auto& layer = driver->layers_[layer_i]; + layer.transform = layer_spec.transform; + layer.driver_spec = layer_spec.driver_spec; } - - TENSORSTORE_ASSIGN_OR_RETURN(auto domains, GetEffectiveDomainsForLayers()); - - auto driver_ptr = - internal::MakeReadWritePtr(read_write_mode, *this); - TENSORSTORE_ASSIGN_OR_RETURN(driver_ptr->layer_domain_, GetDomain()); - TENSORSTORE_RETURN_IF_ERROR(driver_ptr->InitializeGridIndices(domains)); - - auto transform = tensorstore::IdentityTransform(driver_ptr->layer_domain_); - return internal::Driver::Handle{std::move(driver_ptr), std::move(transform)}; + TENSORSTORE_ASSIGN_OR_RETURN( + auto layer_domains, + internal_stack::GetEffectiveDomainsForLayers( + driver->layers_)); + return internal_stack::MakeStackDriverHandle( + std::move(driver), layer_domains, + internal::TransactionState::ToTransaction(std::move(transaction)), + schema); } /// The mechanism used here is to construct an irregular grid based on @@ -393,9 +569,10 @@ Future StackDriverSpec::Open( /// an R-tree variant using the layer MBRs to restrict the query space, /// then applying a similar gridding decomposition to the read transforms. absl::Status StackDriver::InitializeGridIndices( - const std::vector>& domains) { - assert(domains.size() == bound_spec_.layers.size()); - grid_ = IrregularGrid::Make(span(domains)); + span> domains) { + assert(domains.size() == layers_.size()); + grid_ = IrregularGrid::Make(layers_.empty() ? span(&layer_domain_, 1) + : span(domains)); Index start[kMaxRank]; Index shape[kMaxRank]; @@ -428,16 +605,35 @@ absl::Status StackDriver::InitializeGridIndices( Result StackDriver::GetBoundSpec( internal::OpenTransactionPtr transaction, IndexTransformView<> transform) { - if (transaction) return TransactionError(); auto driver_spec = internal::DriverSpec::Make(); - *driver_spec = bound_spec_; - + driver_spec->data_copy_concurrency = data_copy_concurrency_; + driver_spec->schema.Set(dtype_).IgnoreError(); + driver_spec->schema.Set(RankConstraint{rank()}).IgnoreError(); // When constructing the bound spec, set the dimension_units_ and // the current domain on the schema. driver_spec->schema.Set(Schema::DimensionUnits(dimension_units_)) .IgnoreError(); driver_spec->schema.Set(layer_domain_).IgnoreError(); + const size_t num_layers = layers_.size(); + driver_spec->layers.resize(num_layers); + auto status = ForEachLayer(num_layers, [&](size_t layer_i) { + auto& layer_spec = driver_spec->layers[layer_i]; + const auto& layer = layers_[layer_i]; + if (layer.driver_spec) { + layer_spec.transform = layer.transform; + layer_spec.driver_spec = layer.driver_spec; + } else { + TENSORSTORE_ASSIGN_OR_RETURN( + auto driver_spec, + GetTransformedDriverSpec(layer.GetDriverHandle(transaction), + SpecRequestOptions())); + layer_spec.transform = std::move(driver_spec.transform); + layer_spec.driver_spec = std::move(driver_spec.driver_spec); + } + return absl::OkStatus(); + }); + if (!status.ok()) return status; TransformedDriverSpec spec; spec.driver_spec = std::move(driver_spec); spec.transform = transform; @@ -447,7 +643,6 @@ Result StackDriver::GetBoundSpec( Future> StackDriver::ResolveBounds( OpenTransactionPtr transaction, IndexTransform<> transform, ResolveBoundsOptions options) { - if (transaction) return TransactionError(); // All layer bounds are required to be specified in the spec, and // may not be modified later, so here we propagate the composed bounds // to the index transform. @@ -460,6 +655,26 @@ Future> StackDriver::ResolveBounds( return TransformAccess::Make>(std::move(transform_ptr)); } +template +absl::Status ComposeAndDispatchOperation( + StateType& state, const internal::DriverHandle& driver_handle, + IndexTransform<> cell_transform) { + TENSORSTORE_RETURN_IF_ERROR(internal::ValidateSupportsModes( + driver_handle.driver.read_write_mode(), StateType::kMode)); + // The transform for the layer to use is: + // `Compose(outer_request, Compose(driver, cell))` + TENSORSTORE_ASSIGN_OR_RETURN( + auto a_transform, + ComposeTransforms(state.orig_transform, cell_transform)); + TENSORSTORE_ASSIGN_OR_RETURN( + auto b_transform, + ComposeTransforms(driver_handle.transform, std::move(a_transform))); + + state.Dispatch(driver_handle, std::move(b_transform), + std::move(cell_transform)); + return absl::OkStatus(); +} + /// Starts reads for each cell against a provided driver handle. template struct AfterOpenOp { @@ -472,16 +687,9 @@ struct AfterOpenOp { return f.result().status(); } // After opening the layer, issue reads to each of the the grid cells. - // The read transform is: Compose(outer_request, Compose(driver, cell)). for (auto& cell_transform : cells) { - TENSORSTORE_ASSIGN_OR_RETURN( - auto a_transform, - ComposeTransforms(state->orig_transform, cell_transform)); - TENSORSTORE_ASSIGN_OR_RETURN( - auto b_transform, - ComposeTransforms(f.result()->transform, std::move(a_transform))); - - state->Dispatch(*f.result(), std::move(b_transform), cell_transform); + TENSORSTORE_RETURN_IF_ERROR(ComposeAndDispatchOperation( + *state, f.value(), std::move(cell_transform))); } return absl::OkStatus(); } @@ -496,7 +704,7 @@ struct AfterOpenOp { } }; -// OpenLayerOp partitons the transform by layer, invoking UnmappedOp for each +// OpenLayerOp partitions the transform by layer, invoking UnmappedOp for each // grid cell which is not backed by a layer, and then opens each layer and // and initiates OpType (one of LayerReadOp/LayerWriteOp) for each layer's // cells. @@ -521,7 +729,19 @@ struct OpenLayerOp { IndexTransformView<> cell_transform) { auto it = self->grid_to_layer_.find(grid_cell_indices); if (it != self->grid_to_layer_.end()) { - layers_to_load[it->second].emplace_back(cell_transform); + const size_t layer_i = it->second; + const auto& layer = self->layers_[layer_i]; + if (layer.driver) { + // Layer is already open, dispatch operation directly. + TENSORSTORE_RETURN_IF_ERROR( + ComposeAndDispatchOperation( + *state, layer.GetDriverHandle(state->transaction), + std::move(cell_transform)), + tensorstore::MaybeAnnotateStatus( + _, absl::StrFormat("Layer %d", layer_i))); + } else { + layers_to_load[it->second].emplace_back(cell_transform); + } return absl::OkStatus(); } else { return unmapped(grid_cell_indices, cell_transform); @@ -535,14 +755,17 @@ struct OpenLayerOp { return; } - // Open each layer and invoke OpType for all corresponding cell transforms. + // Open each layer and invoke OpType for all corresponding cell + // transforms. for (auto& kv : layers_to_load) { + const size_t layer_i = kv.first; + const auto& layer = self->layers_[layer_i]; Link(WithExecutor( self->data_copy_executor(), - AfterOpenOp{state, kv.first, std::move(kv.second)}), + AfterOpenOp{state, layer_i, std::move(kv.second)}), state->promise, internal::OpenDriver(state->transaction, - self->bound_spec_.layers[kv.first], + layer.GetTransformedDriverSpec(), StateType::kMode)); } } @@ -563,7 +786,7 @@ struct ReadState : public internal::ChunkOperationState { IndexTransform<> orig_transform; // Initiate the read of an individual transform; dispatched by AfterOpenOp - void Dispatch(internal::Driver::Handle& h, + void Dispatch(const internal::Driver::Handle& h, IndexTransform<> composed_transform, IndexTransform<> cell_transform) { h.driver->Read(transaction, std::move(composed_transform), @@ -609,7 +832,7 @@ struct WriteState : public internal::ChunkOperationState { IndexTransform<> orig_transform; // Initiate the write of an individual transform; dispatched by AfterOpenOp - void Dispatch(internal::Driver::Handle& h, + void Dispatch(const internal::Driver::Handle& h, IndexTransform<> composed_transform, IndexTransform<> cell_transform) { h.driver->Write(transaction, std::move(composed_transform), @@ -640,7 +863,238 @@ void StackDriver::Write(OpenTransactionPtr transaction, OpenLayerOp{std::move(state)}); } +Result> MakeDriverFromLayerSpecs( + span layer_specs, StackOpenOptions& options, + DimensionIndex& rank) { + auto driver = + internal::MakeReadWritePtr(ReadWriteMode::read_write); + auto& dtype = driver->dtype_; + dtype = options.dtype(); + auto& transaction = options.transaction; + auto& context = options.context; + auto& read_write_mode = options.read_write_mode; + driver->layers_.resize(layer_specs.size()); + if (!context) context = Context::Default(); + Transaction common_transaction{no_transaction}; + ReadWriteMode common_read_write_mode = ReadWriteMode::dynamic; + DimensionIndex common_rank = dynamic_rank; + auto status = ForEachLayer(layer_specs.size(), [&](size_t layer_i) { + auto& layer = driver->layers_[layer_i]; + const auto& layer_spec = layer_specs[layer_i]; + const auto& layer_transaction = + layer_spec.is_open() ? layer_spec.transaction : transaction; + if (layer_i == 0) { + common_transaction = layer_transaction; + } else if (layer_transaction != common_transaction) { + return absl::InvalidArgumentError("Transaction mismatch"); + } + layer.transform = layer_spec.transform; + layer.driver_spec = layer_spec.driver_spec; + layer.driver = layer_spec.driver; + DataType layer_dtype; + if (layer_spec.is_open()) { + common_read_write_mode |= layer_spec.driver.read_write_mode(); + layer_dtype = layer_spec.driver->dtype(); + } else { + common_read_write_mode = ReadWriteMode::read_write; + TENSORSTORE_RETURN_IF_ERROR( + DriverSpecBindContext(layer.driver_spec, context)); + layer_dtype = layer_spec.driver_spec->schema.dtype(); + if (!layer.transform.valid()) { + TENSORSTORE_ASSIGN_OR_RETURN( + auto domain, + internal::GetEffectiveDomain(layer.GetTransformedDriverSpec())); + if (!domain.valid()) { + return absl::InvalidArgumentError( + tensorstore::StrCat("Domain must be specified")); + } + layer.transform = IdentityTransform(domain); + } + } + if (layer_dtype.valid()) { + if (!dtype.valid()) { + dtype = layer_dtype; + } else if (dtype != layer_dtype) { + return absl::InvalidArgumentError( + tensorstore::StrCat("Layer dtype of ", layer_dtype, + " does not match existing dtype of ", dtype)); + } + } + DimensionIndex layer_rank = layer.transform.input_rank(); + if (common_rank == dynamic_rank) { + common_rank = layer_rank; + } else if (common_rank != layer_rank) { + return absl::InvalidArgumentError(tensorstore::StrCat( + "Layer domain ", layer.transform.domain(), " of rank ", layer_rank, + " does not match layer 0 rank of ", common_rank)); + } + return absl::OkStatus(); + }); + if (!status.ok()) return status; + + if (common_read_write_mode == ReadWriteMode::dynamic) { + common_read_write_mode = ReadWriteMode::read_write; + } + + if (read_write_mode != ReadWriteMode::dynamic) { + TENSORSTORE_RETURN_IF_ERROR(internal::ValidateSupportsModes( + common_read_write_mode, read_write_mode)); + } else { + read_write_mode = common_read_write_mode; + } + + if (transaction != no_transaction) { + if (common_transaction != no_transaction && + common_transaction != transaction) { + return absl::InvalidArgumentError("Transaction mismatch"); + } + } else { + transaction = std::move(common_transaction); + } + + if (!dtype.valid()) { + return absl::InvalidArgumentError("dtype must be specified"); + } + rank = common_rank; + driver.set_read_write_mode(options.read_write_mode); + options.Set(driver->dtype_).IgnoreError(); + return driver; +} + +Result FinalizeStackHandle( + internal::ReadWritePtr driver, StackOpenOptions&& options) { + Schema& schema = options; + TENSORSTORE_RETURN_IF_ERROR(internal_stack::ApplyLayerOptions( + driver->layers_, schema, options)); + TENSORSTORE_ASSIGN_OR_RETURN( + driver->data_copy_concurrency_, + options.context.GetResource()); + TENSORSTORE_ASSIGN_OR_RETURN( + auto layer_domains, + internal_stack::GetEffectiveDomainsForLayers( + driver->layers_)); + return internal_stack::MakeStackDriverHandle( + std::move(driver), layer_domains, std::move(options.transaction), schema); +} + } // namespace + +Result Overlay(span layer_specs, + StackOpenOptions&& options) { + DimensionIndex rank; + TENSORSTORE_ASSIGN_OR_RETURN( + auto driver, + internal_stack::MakeDriverFromLayerSpecs(layer_specs, options, rank)); + TENSORSTORE_RETURN_IF_ERROR(options.Set(RankConstraint{rank})); + return internal_stack::FinalizeStackHandle(std::move(driver), + std::move(options)); +} + +Result Stack(span layer_specs, + DimensionIndex stack_dimension, + StackOpenOptions&& options) { + if (layer_specs.empty()) { + return absl::InvalidArgumentError( + "At least one layer must be specified for stack"); + } + DimensionIndex orig_rank; + TENSORSTORE_ASSIGN_OR_RETURN(auto driver, + internal_stack::MakeDriverFromLayerSpecs( + layer_specs, options, orig_rank)); + if (orig_rank == kMaxRank) { + return absl::InvalidArgumentError( + tensorstore::StrCat("stack would exceed maximum rank of ", kMaxRank)); + } + const DimensionIndex new_rank = orig_rank + 1; + TENSORSTORE_RETURN_IF_ERROR(options.Set(RankConstraint{new_rank})); + TENSORSTORE_ASSIGN_OR_RETURN( + stack_dimension, + tensorstore::NormalizeDimensionIndex(stack_dimension, new_rank)); + if (auto status = ForEachLayer( + layer_specs.size(), + [&](size_t layer_i) { + auto& transform = driver->layers_[layer_i].transform; + TENSORSTORE_ASSIGN_OR_RETURN( + transform, + std::move(transform) | + Dims(stack_dimension) + .AddNew() + .SizedInterval(static_cast(layer_i), 1)); + return absl::OkStatus(); + }); + !status.ok()) { + return status; + } + return internal_stack::FinalizeStackHandle(std::move(driver), + std::move(options)); +} + +Result Concat(span layer_specs, + DimensionIdentifier concat_dimension, + StackOpenOptions&& options) { + if (layer_specs.empty()) { + return absl::InvalidArgumentError( + "At least one layer must be specified for concat"); + } + DimensionIndex rank; + TENSORSTORE_ASSIGN_OR_RETURN( + auto driver, + internal_stack::MakeDriverFromLayerSpecs(layer_specs, options, rank)); + TENSORSTORE_RETURN_IF_ERROR(options.Set(RankConstraint{rank})); + DimensionIndex concat_dimension_index; + if (concat_dimension.label().data()) { + // concat_dimension specified by label must be resolved to an index. + std::string_view labels[kMaxRank]; + if (auto domain = options.domain(); domain.valid()) { + std::copy(domain.labels().begin(), domain.labels().end(), labels); + } + if (auto status = ForEachLayer( + driver->layers_.size(), + [&](size_t layer_i) { + auto layer_labels = + driver->layers_[layer_i].transform.domain().labels(); + for (DimensionIndex i = 0; i < rank; ++i) { + TENSORSTORE_ASSIGN_OR_RETURN( + labels[i], MergeDimensionLabels(labels[i], layer_labels[i]), + tensorstore::MaybeAnnotateStatus( + _, absl::StrFormat("Mismatch in dimension %d", i))); + } + return absl::OkStatus(); + }); + !status.ok()) { + return status; + } + TENSORSTORE_ASSIGN_OR_RETURN( + concat_dimension_index, + tensorstore::NormalizeDimensionLabel( + concat_dimension.label(), + span(&labels[0], rank))); + } else { + TENSORSTORE_ASSIGN_OR_RETURN( + concat_dimension_index, + tensorstore::NormalizeDimensionIndex(concat_dimension.index(), rank)); + } + Index offset; + if (auto status = ForEachLayer( + layer_specs.size(), + [&](size_t layer_i) { + auto& transform = driver->layers_[layer_i].transform; + if (layer_i != 0) { + TENSORSTORE_ASSIGN_OR_RETURN( + transform, + std::move(transform) | + Dims(concat_dimension_index).TranslateTo(offset)); + } + offset = transform.domain()[concat_dimension_index].exclusive_max(); + return absl::OkStatus(); + }); + !status.ok()) { + return status; + } + return internal_stack::FinalizeStackHandle(std::move(driver), + std::move(options)); +} + } // namespace internal_stack } // namespace tensorstore @@ -648,4 +1102,5 @@ namespace { const tensorstore::internal::DriverRegistration< tensorstore::internal_stack::StackDriverSpec> driver_registration; + } // namespace diff --git a/tensorstore/driver/stack/driver.h b/tensorstore/driver/stack/driver.h new file mode 100644 index 000000000..d05419562 --- /dev/null +++ b/tensorstore/driver/stack/driver.h @@ -0,0 +1,108 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENSORSTORE_DRIVER_STACK_DRIVER_H_ +#define TENSORSTORE_DRIVER_STACK_DRIVER_H_ + +#include +#include + +#include "tensorstore/driver/driver_handle.h" +#include "tensorstore/driver/driver_spec.h" +#include "tensorstore/index.h" +#include "tensorstore/index_space/dimension_identifier.h" +#include "tensorstore/index_space/index_transform.h" +#include "tensorstore/open_mode.h" +#include "tensorstore/open_options.h" +#include "tensorstore/spec.h" +#include "tensorstore/tensorstore.h" +#include "tensorstore/transaction.h" +#include "tensorstore/util/result.h" +#include "tensorstore/util/span.h" + +namespace tensorstore { +namespace internal_stack { + +using StackOpenOptions = TransactionalOpenOptions; + +// Specifies a stack layer as either a `TransformedDriverSpec` or a +// `DriverHandle`. +struct StackLayerSpec { + StackLayerSpec() = default; + explicit StackLayerSpec(const internal::TransformedDriverSpec& spec) + : transform(spec.transform), driver_spec(spec.driver_spec) {} + explicit StackLayerSpec(const internal::DriverHandle& handle) + : transform(handle.transform), + driver(handle.driver), + transaction(handle.transaction) {} + template + explicit StackLayerSpec(const TensorStore& store) + : StackLayerSpec(internal::TensorStoreAccess::handle(store)) {} + explicit StackLayerSpec(const Spec& spec) + : StackLayerSpec(internal_spec::SpecAccess::impl(spec)) {} + template + explicit StackLayerSpec(const std::variant& layer_spec) { + std::visit([&](auto& obj) { *this = StackLayerSpec(obj); }, layer_spec); + } + + bool is_open() const { return static_cast(driver); } + + internal::TransformedDriverSpec GetTransformedDriverSpec() const { + assert(!is_open()); + return internal::TransformedDriverSpec{driver_spec, transform}; + } + + // Index transform that applies to either `driver_spec` or `driver`. + IndexTransform<> transform; + + // Indicates a layer backed by a driver spec that will be opened on demand. + internal::DriverSpecPtr driver_spec; + + // Indicates a layer backed by an open driver. + internal::ReadWritePtr driver; + + // May be specified only if `driver` is non-null. + Transaction transaction{no_transaction}; +}; + +Result Overlay(span layer_specs, + StackOpenOptions&& options); +Result Stack(span layer_specs, + DimensionIndex stack_dimension, + StackOpenOptions&& options); +Result Concat(span layer_specs, + DimensionIdentifier concat_dimension, + StackOpenOptions&& options); + +template +struct StackResult { + using type = TensorStore<>; +}; +template +struct StackResult> { + using type = TensorStore; +}; + +template +struct OverlayResult { + using type = TensorStore<>; +}; +template +struct OverlayResult> { + using type = TensorStore; +}; +} // namespace internal_stack +} // namespace tensorstore + +#endif // TENSORSTORE_DRIVER_STACK_DRIVER_H_ diff --git a/tensorstore/driver/stack/driver_test.cc b/tensorstore/driver/stack/driver_test.cc index a0f5b5595..9adb53b63 100644 --- a/tensorstore/driver/stack/driver_test.cc +++ b/tensorstore/driver/stack/driver_test.cc @@ -13,28 +13,39 @@ // limitations under the License. #include +#include +#include #include #include +#include #include #include -#include "absl/strings/string_view.h" +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include #include "tensorstore/array.h" #include "tensorstore/array_testutil.h" +#include "tensorstore/box.h" #include "tensorstore/context.h" #include "tensorstore/data_type.h" #include "tensorstore/index.h" #include "tensorstore/index_space/dim_expression.h" +#include "tensorstore/index_space/index_domain.h" #include "tensorstore/internal/json_gtest.h" +#include "tensorstore/kvstore/spec.h" #include "tensorstore/open.h" +#include "tensorstore/open_mode.h" #include "tensorstore/progress.h" +#include "tensorstore/schema.h" +#include "tensorstore/stack.h" #include "tensorstore/strided_layout.h" #include "tensorstore/tensorstore.h" -#include "tensorstore/util/future.h" +#include "tensorstore/transaction.h" #include "tensorstore/util/result.h" #include "tensorstore/util/status_testutil.h" +#include "tensorstore/util/unit.h" namespace { @@ -43,7 +54,10 @@ using ::tensorstore::DimensionIndex; using ::tensorstore::Index; using ::tensorstore::MatchesArray; using ::tensorstore::MatchesJson; +using ::tensorstore::MatchesStatus; +using ::tensorstore::OpenMode; using ::tensorstore::ReadProgressFunction; +using ::tensorstore::ReadWriteMode; using ::tensorstore::span; ::nlohmann::json GetRank1Length4ArrayDriver(int inclusive_min, @@ -98,7 +112,6 @@ ::nlohmann::json GetRank1Length4N5Driver(int inclusive_min, } TEST(StackDriverTest, OpenAndResolveBounds) { - auto context = tensorstore::Context::Default(); ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", ::nlohmann::json::array_t({GetRank1Length4ArrayDriver(-3), @@ -107,8 +120,8 @@ TEST(StackDriverTest, OpenAndResolveBounds) { {"schema", {{"dimension_units", {"2px"}}}}, }; - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, + tensorstore::Open(json_spec).result()); EXPECT_EQ(1, store.rank()); EXPECT_EQ(tensorstore::dtype_v, store.dtype()); @@ -166,7 +179,6 @@ TEST(StackDriverTest, OpenAndResolveBounds) { } TEST(StackDriverTest, OpenWithDomain) { - auto context = tensorstore::Context::Default(); ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", ::nlohmann::json::array_t({GetRank1Length4ArrayDriver(-3), @@ -175,15 +187,13 @@ TEST(StackDriverTest, OpenWithDomain) { {"schema", {{"domain", {{"shape", {5}}}}}}, }; - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, + tensorstore::Open(json_spec).result()); ASSERT_EQ(tensorstore::Box<1>({0}, {5}), store.domain().box()); } TEST(StackDriverTest, Read) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", ::nlohmann::json::array_t({GetRank1Length4ArrayDriver(-3), @@ -191,8 +201,8 @@ TEST(StackDriverTest, Read) { GetRank1Length4ArrayDriver(3, 6)})}, }; - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, + tensorstore::Open(json_spec).result()); { TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto array, @@ -230,22 +240,19 @@ TEST(StackDriverTest, Read) { } TEST(StackDriverTest, ReadSparse) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", ::nlohmann::json::array_t({GetRank1Length4ArrayDriver(-3), GetRank1Length4ArrayDriver(2)})}, }; - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, + tensorstore::Open(json_spec).result()); // Cannot read everything EXPECT_THAT(tensorstore::Read(store).result(), - tensorstore::MatchesStatus( - absl::StatusCode::kInvalidArgument, - "Read cell origin=.1. missing layer mapping.*")); + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Read cell origin=.1. missing layer mapping.*")); // Can read the backed data. { @@ -258,9 +265,22 @@ TEST(StackDriverTest, ReadSparse) { } } -TEST(StackDriverTest, Rank0) { - auto context = tensorstore::Context::Default(); +TEST(StackDriverTest, NoLayers) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto spec, tensorstore::Spec::FromJson( + {{"driver", "stack"}, + {"layers", ::nlohmann::json::array_t()}, + {"dtype", "int32"}, + {"schema", {{"domain", {{"shape", {2, 3}}}}}}})); + EXPECT_EQ(OpenMode::unknown, spec.open_mode()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, + tensorstore::Open(spec).result()); + EXPECT_EQ(tensorstore::IndexDomain({2, 3}), store.domain()); + EXPECT_EQ(tensorstore::dtype_v, store.dtype()); + EXPECT_EQ(ReadWriteMode::read_write, store.read_write_mode()); +} +TEST(StackDriverTest, Rank0) { ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", @@ -278,8 +298,8 @@ TEST(StackDriverTest, Rank0) { }}}, }; - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, + tensorstore::Open(json_spec).result()); TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto array, tensorstore::Read<>(store).result()); @@ -288,8 +308,6 @@ TEST(StackDriverTest, Rank0) { } TEST(StackDriverTest, WriteCreate) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", ::nlohmann::json::array_t({GetRank1Length4N5Driver(-3), @@ -298,9 +316,8 @@ TEST(StackDriverTest, WriteCreate) { }; TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context, - tensorstore::OpenMode::open_or_create) - .result()); + auto store, + tensorstore::Open(json_spec, OpenMode::open_or_create).result()); TENSORSTORE_EXPECT_OK( tensorstore::Write( @@ -310,26 +327,22 @@ TEST(StackDriverTest, WriteCreate) { } TEST(StackDriverTest, WriteSparse) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", ::nlohmann::json::array_t({GetRank1Length4N5Driver(-3), GetRank1Length4N5Driver(2)})}, }; TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context, - tensorstore::OpenMode::open_or_create) - .result()); + auto store, + tensorstore::Open(json_spec, OpenMode::open_or_create).result()); // Cannot write everything EXPECT_THAT(tensorstore::Write(tensorstore::MakeOffsetArray( {-3}, {9, 8, 7, 6, 5, 4, 3, 2, 1}), store) .result(), - tensorstore::MatchesStatus( - absl::StatusCode::kInvalidArgument, - "Write cell origin=.1. missing layer mapping.*")); + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Write cell origin=.1. missing layer mapping.*")); // Can write the backed data. TENSORSTORE_ASSERT_OK( @@ -339,8 +352,6 @@ TEST(StackDriverTest, WriteSparse) { } TEST(StackDriverTest, ReadWriteNonExistingLayers) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", ::nlohmann::json::array_t({GetRank1Length4N5Driver(-3), @@ -348,26 +359,24 @@ TEST(StackDriverTest, ReadWriteNonExistingLayers) { GetRank1Length4N5Driver(3, 6)})}, }; - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, + tensorstore::Open(json_spec).result()); // Sublayers are opened on read; they do not exist. EXPECT_THAT(tensorstore::Read(store).result(), - tensorstore::MatchesStatus(absl::StatusCode::kNotFound, - ".*Error opening \"n5\" driver: .*")); + MatchesStatus(absl::StatusCode::kNotFound, + ".*Error opening \"n5\" driver: .*")); // Sublayers are opened on write; they do not exist. EXPECT_THAT(tensorstore::Write(tensorstore::MakeOffsetArray( {-3}, {9, 8, 7, 6, 5, 4, 3, 2, 1}), store) .result(), - tensorstore::MatchesStatus(absl::StatusCode::kNotFound, - ".*Error opening \"n5\" driver: .*")); + MatchesStatus(absl::StatusCode::kNotFound, + ".*Error opening \"n5\" driver: .*")); } TEST(StackDriverTest, Schema_MismatchedDtype) { - auto context = tensorstore::Context::Default(); - auto a = GetRank1Length4N5Driver(0); a["dtype"] = "int64"; @@ -380,22 +389,20 @@ TEST(StackDriverTest, Schema_MismatchedDtype) { }}, }; - EXPECT_THAT(tensorstore::Open(json_spec, context).result(), - tensorstore::MatchesStatus( - absl::StatusCode::kInvalidArgument, - ".*dtype .int32. does not match existing value .int64.*")); + EXPECT_THAT( + tensorstore::Open(json_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + ".*dtype .int32. does not match existing value .int64.*")); json_spec.erase("schema"); - EXPECT_THAT(tensorstore::Open(json_spec, context).result(), - tensorstore::MatchesStatus( - absl::StatusCode::kInvalidArgument, - ".*dtype .int32. does not match existing value .int64.*")); + EXPECT_THAT( + tensorstore::Open(json_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + ".*dtype .int32. does not match existing value .int64.*")); } TEST(StackDriverTest, Schema_MismatchedRank) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json a{ {"driver", "array"}, {"array", {{1, 3, 3}, {4, 6, 6}}}, @@ -412,23 +419,19 @@ TEST(StackDriverTest, Schema_MismatchedRank) { }; EXPECT_THAT( - tensorstore::Open(json_spec, context).result(), - tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, - "Rank specified by rank .2. does not match " - "existing rank specified by schema .1.")); + tensorstore::Open(json_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 0: Rank of 2 does not match existing rank of 1")); json_spec.erase("schema"); EXPECT_THAT( - tensorstore::Open(json_spec, context).result(), - tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, - "Rank specified by rank .1. does not match " - "existing rank specified by schema .2.")); + tensorstore::Open(json_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 1: Rank of 1 does not match existing rank of 2")); } TEST(StackDriverTest, Schema_HasCodec) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json a{ {"driver", "array"}, {"array", {{1, 3, 3}, {4, 6, 6}}}, @@ -444,15 +447,77 @@ TEST(StackDriverTest, Schema_HasCodec) { }}, }; - EXPECT_THAT(tensorstore::Open(json_spec, context).result(), - tensorstore::MatchesStatus( - absl::StatusCode::kInvalidArgument, - "\"codec\" not supported by \"stack\" driver")); + EXPECT_THAT(tensorstore::Open(json_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "codec option not supported by \"stack\" driver")); } -TEST(StackDriverTest, Schema_DimensionUnitsInSchema) { - auto context = tensorstore::Context::Default(); +TEST(StackDriverTest, Schema_HasChunkLayout) { + ::nlohmann::json a{ + {"driver", "array"}, + {"array", {{1, 3, 3}, {4, 6, 6}}}, + {"dtype", "int32"}, + }; + + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", ::nlohmann::json::array_t({a})}, + {"schema", + { + {"chunk_layout", {{"inner_order", {0, 1}}}}, + }}, + }; + + EXPECT_THAT( + tensorstore::Open(json_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "chunk layout option not supported by \"stack\" driver")); +} + +TEST(StackDriverTest, Schema_HasFillValue) { + ::nlohmann::json a{ + {"driver", "array"}, + {"array", {{1, 3, 3}, {4, 6, 6}}}, + {"dtype", "int32"}, + }; + + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", ::nlohmann::json::array_t({a})}, + {"schema", + { + {"fill_value", 42}, + }}, + }; + + EXPECT_THAT( + tensorstore::Open(json_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "fill value option not supported by \"stack\" driver")); +} +TEST(StackDriverTest, HasKvstore) { + ::nlohmann::json a{ + {"driver", "array"}, + {"array", {{1, 3, 3}, {4, 6, 6}}}, + {"dtype", "int32"}, + }; + + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", ::nlohmann::json::array_t({a})}, + }; + + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto kvstore_spec, tensorstore::kvstore::Spec::FromJson("memory://")); + + EXPECT_THAT( + tensorstore::Open(json_spec, kvstore_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "kvstore option not supported by \"stack\" driver")); +} + +TEST(StackDriverTest, Schema_DimensionUnitsInSchema) { ::nlohmann::json a{ {"driver", "array"}, {"array", {{1, 3, 3}, {4, 6, 6}}}, @@ -466,9 +531,8 @@ TEST(StackDriverTest, Schema_DimensionUnitsInSchema) { }; TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context, - tensorstore::OpenMode::open_or_create) - .result()); + auto store, + tensorstore::Open(json_spec, OpenMode::open_or_create).result()); EXPECT_THAT(store.dimension_units(), ::testing::Optional(::testing::ElementsAre( @@ -476,8 +540,6 @@ TEST(StackDriverTest, Schema_DimensionUnitsInSchema) { } TEST(StackDriverTest, Schema_DimensionUnitsInLayer) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json a{ {"driver", "array"}, {"array", {{1, 3, 3}, {4, 6, 6}}}, @@ -491,9 +553,8 @@ TEST(StackDriverTest, Schema_DimensionUnitsInLayer) { }; TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context, - tensorstore::OpenMode::open_or_create) - .result()); + auto store, + tensorstore::Open(json_spec, OpenMode::open_or_create).result()); EXPECT_THAT(store.dimension_units(), ::testing::Optional(::testing::ElementsAre( @@ -501,8 +562,6 @@ TEST(StackDriverTest, Schema_DimensionUnitsInLayer) { } TEST(StackDriverTest, Schema_MismatchedDimensionUnits) { - auto context = tensorstore::Context::Default(); - ::nlohmann::json a{ {"driver", "array"}, {"array", {{1, 3, 3}, {4, 6, 6}}}, @@ -517,18 +576,43 @@ TEST(StackDriverTest, Schema_MismatchedDimensionUnits) { }; TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto store, tensorstore::Open(json_spec, context, - tensorstore::OpenMode::open_or_create) - .result()); + auto store, + tensorstore::Open(json_spec, OpenMode::open_or_create).result()); EXPECT_THAT(store.dimension_units(), ::testing::Optional(::testing::ElementsAre( tensorstore::Unit("4nm"), tensorstore::Unit("2ft")))); } -TEST(StackDriverTest, SchemaDomain_MismatchedShape) { - auto context = tensorstore::Context::Default(); +TEST(StackDriverTest, Schema_MismatchedLayerDimensionUnits) { + ::nlohmann::json a{ + {"driver", "array"}, + {"array", {{1, 3, 3}, {4, 6, 6}}}, + {"dtype", "int32"}, + {"schema", {{"dimension_units", {"1ft", "2ft"}}}}, + }; + + ::nlohmann::json b{ + {"driver", "array"}, + {"array", {{1, 3, 3}, {4, 6, 6}}}, + {"dtype", "int32"}, + {"schema", {{"dimension_units", {"4nm", nullptr}}}}, + }; + + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", {a, b}}, + }; + + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto store, + tensorstore::Open(json_spec).result()); + + EXPECT_THAT(store.dimension_units(), + ::testing::Optional(::testing::ElementsAre( + std::nullopt, tensorstore::Unit("2ft")))); +} +TEST(StackDriverTest, SchemaDomain_MismatchedShape) { ::nlohmann::json json_spec{ {"driver", "stack"}, {"layers", ::nlohmann::json::array_t({GetRank1Length4N5Driver(0), @@ -536,9 +620,10 @@ TEST(StackDriverTest, SchemaDomain_MismatchedShape) { {"schema", {{"domain", {{"shape", {2, 2}}}}}}, }; - EXPECT_THAT(tensorstore::Open(json_spec, context).result(), - tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, - "Rank specified by.*")); + EXPECT_THAT( + tensorstore::Open(json_spec).result(), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 0: Rank of 1 does not match existing rank of 2")); } TEST(StackDriverTest, InvalidLayerMappingBug) { @@ -591,4 +676,737 @@ TEST(StackDriverTest, InvalidLayerMappingBug) { TENSORSTORE_ASSERT_OK(tensorstore::Read(store).result()); } +TEST(StackDriverTest, OpenMode_Open) { + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}, + }, + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}, + }, + }}, + }; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec, + tensorstore::Spec::FromJson(json_spec)); + EXPECT_EQ(OpenMode::open, spec.open_mode()); +} + +TEST(StackDriverTest, SpecDimensionUnitsMatch) { + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", {{"shape", {5}}}}, + {"dimension_units", {"4nm"}}}}, + {"open", true}, + {"create", true}, + }, + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", {{"shape", {5}}}}, + {"dimension_units", {"4nm"}}}}, + {"open", true}, + {"create", true}, + }, + }}, + }; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec, + tensorstore::Spec::FromJson(json_spec)); + EXPECT_THAT(spec.dimension_units(), + ::testing::Optional(::testing::ElementsAre("4nm"))); +} + +TEST(StackDriverTest, SpecDimensionUnitsMismatch) { + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", {{"shape", {5}}}}, + {"dimension_units", {"4nm"}}}}, + {"open", true}, + {"create", true}, + }, + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", {{"shape", {5}}}}, + {"dimension_units", {"8nm"}}}}, + {"open", true}, + {"create", true}, + }, + }}, + }; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec, + tensorstore::Spec::FromJson(json_spec)); + EXPECT_THAT(spec.dimension_units(), + ::testing::Optional(::testing::ElementsAre(std::nullopt))); +} + +TEST(StackDriverTest, SpecDimensionUnitsInSchema) { + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", {{"shape", {5}}}}, + {"dimension_units", {"4nm"}}}}, + {"open", true}, + {"create", true}, + }, + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", {{"shape", {5}}}}, + {"dimension_units", {"8nm"}}}}, + {"open", true}, + {"create", true}, + }, + }}, + {"schema", {{"dimension_units", {"7s"}}}}, + }; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec, + tensorstore::Spec::FromJson(json_spec)); + EXPECT_THAT(spec.dimension_units(), + ::testing::Optional(::testing::ElementsAre("7s"))); +} + +TEST(StackDriverTest, OpenMode_OpenCreate) { + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}, + {"open", true}, + {"create", true}, + }, + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}, + {"open", true}, + {"create", true}, + }, + }}, + }; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec, + tensorstore::Spec::FromJson(json_spec)); + EXPECT_EQ(OpenMode::open_or_create, spec.open_mode()); +} + +TEST(StackDriverTest, OpenMode_Mismatch) { + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}, + {"open", true}, + {"create", true}, + }, + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}, + {"open", true}, + }, + }}, + }; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec, + tensorstore::Spec::FromJson(json_spec)); + EXPECT_EQ(OpenMode::unknown, spec.open_mode()); +} + +TEST(StackDriverTest, DomainUnspecified) { + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}}}, + }, + }}, + }; + EXPECT_THAT(tensorstore::Spec::FromJson(json_spec), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 0: Domain must be specified")); +} + +TEST(StackDriverTest, RankMismatch) { + ::nlohmann::json json_spec{ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}, + }, + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, {"domain", {{"shape", {5, 2}}}}}}, + }, + }}, + }; + EXPECT_THAT( + tensorstore::Spec::FromJson(json_spec), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 1: Rank of 2 does not match existing rank of 1")); +} + +// Checks that Context is bound properly to unopened (`Spec`) layers. +TEST(OverlayTest, Context) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeArray({1, 2, 3, 4, 5}), base) + .result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec_shifted, base_spec | tensorstore::Dims(0).TranslateBy(5)); + // Create stack from two `Spec` objects. + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Overlay({base_spec, base_spec_shifted}, context)); + EXPECT_THAT(tensorstore::Read(store).result(), + ::testing::Optional(tensorstore::MakeArray( + {1, 2, 3, 4, 5, 1, 2, 3, 4, 5}))); + } + + // Create stack from `TensorStore` and `Spec`. + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay({base, base_spec_shifted}, context)); + EXPECT_THAT(tensorstore::Read(store).result(), + ::testing::Optional(tensorstore::MakeArray( + {1, 2, 3, 4, 5, 1, 2, 3, 4, 5}))); + } +} + +TEST(OverlayTest, Transaction) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeArray({1, 2, 3, 4, 5}), base) + .result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec_shifted, base_spec | tensorstore::Dims(0).TranslateBy(5)); + tensorstore::Transaction txn1(tensorstore::atomic_isolated); + tensorstore::Transaction txn2(tensorstore::atomic_isolated); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto base_shifted, + base | tensorstore::Dims(0).TranslateBy(5)); + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Overlay({base_spec, base_spec_shifted}, context, txn1)); + EXPECT_EQ(txn1, store.transaction()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Overlay({(base | txn1).value(), base_spec_shifted}, + context, txn1)); + EXPECT_EQ(txn1, store.transaction()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay({base, base_shifted}, context, txn1)); + EXPECT_EQ(txn1, store.transaction()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Overlay( + {(base | txn1).value(), (base_shifted | txn1).value()}, context)); + EXPECT_EQ(txn1, store.transaction()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay( + {(base | txn1).value(), (base_shifted | txn1).value()}, + context, txn1)); + EXPECT_EQ(txn1, store.transaction()); + } + + EXPECT_THAT( + tensorstore::Overlay({(base | txn1).value(), base_spec_shifted}, context), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 1: Transaction mismatch")); + + EXPECT_THAT( + tensorstore::Overlay({(base | txn1).value(), base_spec_shifted}, context), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 1: Transaction mismatch")); + + EXPECT_THAT(tensorstore::Overlay( + {(base | txn1).value(), (base_shifted | txn1).value()}, + context, txn2), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Transaction mismatch")); +} + +TEST(OverlayTest, NoLayers) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay({}, tensorstore::dtype_v, + tensorstore::Schema::Shape({2, 3}))); + EXPECT_EQ(tensorstore::IndexDomain({2, 3}), store.domain()); + EXPECT_EQ(tensorstore::dtype_v, store.dtype()); + EXPECT_EQ(ReadWriteMode::read_write, store.read_write_mode()); +} + +TEST(OverlayTest, ReadWriteMode) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeArray({1, 2, 3, 4, 5}), base) + .result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec_shifted, base_spec | tensorstore::Dims(0).TranslateBy(5)); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto base_shifted, + base | tensorstore::Dims(0).TranslateBy(5)); + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Overlay({base_spec, base_spec_shifted}, context)); + EXPECT_EQ(ReadWriteMode::read_write, store.read_write_mode()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay({base_spec, base_spec_shifted}, + context, ReadWriteMode::read)); + EXPECT_EQ(ReadWriteMode::read, store.read_write_mode()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay({base_spec, base_spec_shifted}, + context, ReadWriteMode::write)); + EXPECT_EQ(ReadWriteMode::write, store.read_write_mode()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay({base, base_shifted}, context)); + EXPECT_EQ(ReadWriteMode::read_write, store.read_write_mode()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay({base, base_shifted}, context, + ReadWriteMode::read)); + EXPECT_EQ(ReadWriteMode::read, store.read_write_mode()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Overlay( + {tensorstore::ModeCast(base, ReadWriteMode::read).value(), + base_shifted}, + context)); + EXPECT_EQ(ReadWriteMode::read_write, store.read_write_mode()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Overlay( + {tensorstore::ModeCast(base, ReadWriteMode::read).value(), + tensorstore::ModeCast(base_shifted, ReadWriteMode::write).value()}, + context)); + EXPECT_EQ(ReadWriteMode::read_write, store.read_write_mode()); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Overlay( + {tensorstore::ModeCast(base, ReadWriteMode::read).value(), + tensorstore::ModeCast(base_shifted, ReadWriteMode::read).value()}, + context)); + EXPECT_EQ(ReadWriteMode::read, store.read_write_mode()); + } + + EXPECT_THAT( + tensorstore::Overlay( + {tensorstore::ModeCast(base, ReadWriteMode::read).value(), + tensorstore::ModeCast(base_shifted, ReadWriteMode::read).value()}, + context, ReadWriteMode::write), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Write mode not supported")); +} + +TEST(OverlayTest, DomainUnspecified) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson({{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}}}})); + EXPECT_THAT(tensorstore::Overlay({base_spec}), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 0: Domain must be specified")); +} + +TEST(OverlayTest, DomainRankMismatch) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec1, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {2, 3}}}}}}})); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec2, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {2}}}}}}})); + EXPECT_THAT(tensorstore::Overlay({base_spec1, base_spec2}), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 1: Layer domain .* of rank 1 does not match " + "layer 0 rank of 2")); +} + +TEST(OverlayTest, Spec) { + using array_t = ::nlohmann::json::array_t; + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto base_shifted, + base | tensorstore::Dims(0).TranslateBy(5)); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Overlay({base, base_shifted}, context)); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec, + store.spec(tensorstore::MinimalSpec{true})); + EXPECT_THAT( + spec.ToJson(), + ::testing::Optional(MatchesJson({ + {"driver", "stack"}, + {"layers", + { + { + {"driver", "zarr3"}, + {"dtype", "int32"}, + {"transform", + {{"input_exclusive_max", array_t(1, array_t{5})}, + {"input_inclusive_min", array_t{0}}}}, + {"kvstore", {{"driver", "memory"}}}, + }, + { + {"driver", "zarr3"}, + {"dtype", "int32"}, + {"transform", + {{"input_exclusive_max", array_t(1, array_t{10})}, + {"input_inclusive_min", array_t{5}}, + {"output", {{{"input_dimension", 0}, {"offset", -5}}}}}}, + {"kvstore", {{"driver", "memory"}}}, + }, + }}, + {"dtype", "int32"}, + {"schema", + {{"domain", + {{"inclusive_min", array_t{0}}, {"exclusive_max", array_t{10}}}}}}, + {"transform", + {{"input_exclusive_max", array_t{10}}, + {"input_inclusive_min", array_t{0}}}}, + }))); +} + +TEST(OverlayTest, OpenMode) { + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto base_shifted, + base | tensorstore::Dims(0).TranslateBy(5)); + TENSORSTORE_ASSERT_OK( + tensorstore::Overlay({base, base_shifted}, context, OpenMode::open)); + TENSORSTORE_ASSERT_OK(tensorstore::Overlay({base, base_shifted}, context, + OpenMode::open_or_create)); + EXPECT_THAT( + tensorstore::Overlay({base, base_shifted}, context, OpenMode::create), + MatchesStatus( + absl::StatusCode::kInvalidArgument, + "Layer 0: Open mode of create is not compatible with already-open " + "layer")); +} + +TEST(OverlayTest, RecheckCached) { + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto base_shifted, + base | tensorstore::Dims(0).TranslateBy(5)); + TENSORSTORE_ASSERT_OK(tensorstore::Overlay({base, base_shifted}, context)); + EXPECT_THAT(tensorstore::Overlay({base, base_shifted}, context, + tensorstore::RecheckCachedData{false}), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 0: Cannot specify cache rechecking options " + "with already-open layer")); + EXPECT_THAT(tensorstore::Overlay({base, base_shifted}, context, + tensorstore::RecheckCachedMetadata{false}), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 0: Cannot specify cache rechecking options " + "with already-open layer")); +} + +TEST(OverlayTest, MissingDtype) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"domain", {{"shape", {5}}}}}}})); + EXPECT_THAT(tensorstore::Overlay({base_spec}), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "dtype must be specified")); +} + +TEST(OverlayTest, DtypeMismatch) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec1, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://a/"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec2, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://b/"}, + {"schema", {{"dtype", "int64"}, {"domain", {{"shape", {5}}}}}}})); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec2_shifted, + base_spec2 | tensorstore::Dims(0).TranslateBy(5)); + EXPECT_THAT(tensorstore::Overlay({base_spec1, base_spec2_shifted}), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Layer 1: Layer dtype of int64 does not match " + "existing dtype of int32")); +} + +TEST(StackTest, Basic) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeArray({1, 2, 3, 4, 5}), base) + .result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Stack({base_spec, base}, 0, context)); + EXPECT_EQ(tensorstore::IndexDomain({2, 5}), store.domain()); + EXPECT_THAT(tensorstore::Read(store).result(), + ::testing::Optional(tensorstore::MakeArray( + {{1, 2, 3, 4, 5}, {1, 2, 3, 4, 5}}))); +} + +TEST(StackTest, NoLayers) { + EXPECT_THAT(tensorstore::Stack({}, 0), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "At least one layer must be specified for stack")); +} + +TEST(StackTest, MaxRank) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", {{"shape", ::nlohmann::json::array_t(32, 1)}}}}}})); + EXPECT_THAT(tensorstore::Stack({base_spec, base_spec}, 0), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "stack would exceed maximum rank of 32")); +} + +TEST(ConcatTest, Basic) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson( + {{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", {{"dtype", "int32"}, {"domain", {{"shape", {5}}}}}}})); + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeArray({1, 2, 3, 4, 5}), base) + .result()); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Concat({base_spec, base}, 0, context)); + EXPECT_THAT(tensorstore::Read(store).result(), + ::testing::Optional(tensorstore::MakeArray( + {1, 2, 3, 4, 5, 1, 2, 3, 4, 5}))); +} + +TEST(ConcatTest, DimensionLabel) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec, + tensorstore::Spec::FromJson({{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", + { + {"shape", {1, 5}}, + {"labels", {"x", "y"}}, + }}}}})); + auto context = tensorstore::Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base, + tensorstore::Open(base_spec, OpenMode::create, context).result()); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeArray({{1, 2, 3, 4, 5}}), + base) + .result()); + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Concat({base_spec, base}, "x", context)); + EXPECT_THAT(tensorstore::Read(store).result(), + ::testing::Optional(tensorstore::MakeArray( + {{1, 2, 3, 4, 5}, {1, 2, 3, 4, 5}}))); + } + + { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, tensorstore::Concat({base_spec, base}, "y", context)); + EXPECT_THAT(tensorstore::Read(store).result(), + ::testing::Optional(tensorstore::MakeArray( + {{1, 2, 3, 4, 5, 1, 2, 3, 4, 5}}))); + } + + EXPECT_THAT(tensorstore::Concat({base_spec, base}, "z", context), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "Label \"z\" does not match one of .*")); +} + +TEST(ConcatTest, NoLayers) { + EXPECT_THAT(tensorstore::Concat({}, 0), + MatchesStatus(absl::StatusCode::kInvalidArgument, + "At least one layer must be specified for concat")); +} + +// When the concat_dimension is specified as a label, the dimension labels have +// to be resolved early using a different code path. This tests that code path. +TEST(ConcatTest, DimensionLabelMismatch) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec1, + tensorstore::Spec::FromJson({{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", + { + {"shape", {1, 5}}, + {"labels", {"x", "y"}}, + }}}}})); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto base_spec2, + tensorstore::Spec::FromJson({{"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"schema", + {{"dtype", "int32"}, + {"domain", + { + {"shape", {1, 5}}, + {"labels", {"x", "z"}}, + }}}}})); + EXPECT_THAT( + tensorstore::Concat({base_spec1, base_spec2}, "x"), + MatchesStatus( + absl::StatusCode::kInvalidArgument, + "Layer 1: Mismatch in dimension 1: Dimension labels do not match")); +} + } // namespace diff --git a/tensorstore/driver/stack/index.rst b/tensorstore/driver/stack/index.rst index cd5b61ea6..7e2908f61 100644 --- a/tensorstore/driver/stack/index.rst +++ b/tensorstore/driver/stack/index.rst @@ -1,32 +1,50 @@ +.. _stack-driver: + ``stack`` Driver ================ -The ``stack`` driver specifies a TensorStore backed by a sequence of layered -drivers, where each underlying layer describes it's bounds via a transform. +The ``stack`` driver specifies a TensorStore virtually overlays a sequence of +TensorStore *layers* within a common domain. + +- If the domains of the :json:schema:`~driver/stack.layers` may overlap, the + last layer that contains any given position within the domain takes + precedence. This last layer that contains the position is said to be the + *backing* layer for that position. -- The TensorStore for a stack driver has static bounds, which may be derived - from the underlying :json:schema:`~driver/stack.layers` +- By choosing appropriate transforms for the layers, this driver may be used to + virtually stack a sequence of TensorStores along a new dimension, or to + concatenate a sequence of TensorStores along an existing dimension. For these + use cases, the layer domains do not overlap. -- Reading is supported if the :json:schema:`~driver/stack.layers` TensorStore - supports reading. +Python API +---------- -- Writing is supported if the :json:schema:`~driver/stack.layers` TensorStore - supports writing. +Python APIs are provided for constructing a stack-driver-backed TensorStore +based on a sequence of :py:obj:`~tensorstore.Spec` or already-open +:py:obj:`~tensorstore.TensorStore` objects. -The layers may overlap, with the subsequent layers overriding values from -earlier layers in the stack. +- :py:obj:`tensorstore.overlay` virtually overlays a sequence of layers. +- :py:obj:`tensorstore.stack` stacks a sequence of layers along a new dimension. +- :py:obj:`tensorstore.concat` concatenates a sequence of layers along an + existing dimension. -The underlying ``stack`` TensorStore :json:schema:`~driver/stack.layers` are -opened on demand by each operation. +Supported operations +-------------------- +Opening the stack driver itself does not open any of the layers. Instead, +layers are opened on demand as needed to satisfy read and write operations. + +Reading or writing at a given position is supported, if any only if, the backing +layer supports reading or writing, respectively. Reading or writing at a +position that has no backing layer results in an error. .. json:schema:: driver/stack Example ------------ +------- -.. admonition:: Example of layer driver +.. admonition:: Example of stack driver :class: example >>> ts.open({ @@ -56,33 +74,62 @@ Example dtype("int32") - TensorStore Schema ------------------------------ +------------------ + +Data type +~~~~~~~~~ + +The :json:schema:`~Schema.dtype` must be specified for at least one layer, or in +the top-level :json:schema:`~driver/stack.schema`. All specified data types +must match. + +Domain +~~~~~~ + +The bounds of a stack driver-backed TensorStore, and of each layer, are fixed +when it is opened, based on the domains specified for each of the +:json:schema:`~driver/stack.layers`, and the :json:schema:`~Schema.domain` +specified for the stack driver, if any. + +Any :ref:`implicit/resizeable bounds` of layers become fixed +(explicit) bounds and will not be affected by subsequent resize operations. + +By default, the bounds of the stack-driver-backed TensorStore are determined by +computing the hull of the effective domains of the layers, but any finite or +explicit bound specified on the top-level stack-driver +:json:schema:`~Schema.domain` takes precedence and overrides the bound +determined from the hull. + +Note that it is valid for the domain to include positions that are not backed by +any layer, but any attempt to read or write such positions results in an error. + +The :ref:`dimension labels` are merged and must be compatible. + +Dimension units +~~~~~~~~~~~~~~~ + +For each dimension, the corresponding dimension unit is determined as follows: + +1. If a unit for the dimension is specified directly in the top-level stack-driver + :json:schema:`Schema.dimension_units`, the specified unit is assigned. -Only the following TensorStore schema constraints are useful to the ``stack`` -driver. They will be propagated and will cause errors when incompatible with -the layer stack. +2. Otherwise, if there is agreement among all the layers that specify a unit for + the dimension, the common unit is assigned. -- :json:schema:`Schema.rank` Sets the rank of the TensorStore and constrains - the rank of each of the :json:schema:`~driver/stack.layers` +3. Otherwise, the unit is unassigned (:json:`null`). -- :json:schema:`Schema.dtype` Sets the dtype of the TensorStore and constrains - the dtype of each of the :json:schema:`~driver/stack.layers` +Fill value +~~~~~~~~~~ -- :json:schema:`Schema.domain` Sets the domain of the TensorStore. The ``stack`` - driver computes an effective domain as the hull of each layer's effective - domain, as derived from the layer schema and transform. For each dimension - bound which is implicit and unbounded (as when the ``domain`` is omitted), - the bounds for that dimension are taken from the effective domain. +Fill values are not supported. -- :json:schema:`Schema.dimension_units` Sets the dimension units of the - TensorStore. The ``stack`` driver will set any unset dimension units from the - individual layers as long as they are in common. +Codec +~~~~~ +Codecs are not supported. -Scheama options that are not allowed: +Chunk layout +~~~~~~~~~~~~ -- :json:schema:`Schema.fill_value` -- :json:schema:`Schema.codec` -- :json:schema:`Schema.chunk_layout` +Chunk layouts are not supported. diff --git a/tensorstore/driver/stack/integration_test.cc b/tensorstore/driver/stack/integration_test.cc index d86167659..783d37ec6 100644 --- a/tensorstore/driver/stack/integration_test.cc +++ b/tensorstore/driver/stack/integration_test.cc @@ -13,15 +13,20 @@ // limitations under the License. #include +#include #include #include +#include #include +#include "absl/status/status.h" #include "absl/strings/cord.h" #include "absl/strings/string_view.h" #include +#include "tensorstore/array.h" #include "tensorstore/array_testutil.h" +#include "tensorstore/box.h" #include "tensorstore/context.h" #include "tensorstore/data_type.h" #include "tensorstore/index.h" @@ -29,9 +34,9 @@ #include "tensorstore/kvstore/kvstore.h" #include "tensorstore/kvstore/operations.h" #include "tensorstore/open.h" +#include "tensorstore/open_mode.h" #include "tensorstore/progress.h" #include "tensorstore/tensorstore.h" -#include "tensorstore/util/future.h" #include "tensorstore/util/result.h" #include "tensorstore/util/status.h" #include "tensorstore/util/status_testutil.h" @@ -366,9 +371,8 @@ TEST(IntegrationTest, NeuroglancerPrecomputed) { }, }; EXPECT_THAT(tensorstore::Open(spec, context).result(), - tensorstore::MatchesStatus( - absl::StatusCode::kInvalidArgument, - ".*Unable to infer \"dtype\" in \"stack\" driver.*")); + tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, + ".*dtype must be specified.*")); } // Missing transform results in an unbounded domain, which may be opened, diff --git a/tensorstore/driver/virtual_chunked/BUILD b/tensorstore/driver/virtual_chunked/BUILD index 3bea8a55a..54e892543 100644 --- a/tensorstore/driver/virtual_chunked/BUILD +++ b/tensorstore/driver/virtual_chunked/BUILD @@ -4,22 +4,41 @@ package(default_visibility = ["//visibility:public"]) licenses(["notice"]) -# This source file is compiled as part of the -# //tensorstore:virtual_chunked target in order to allow the -# `virtual_chunked.h` header to be part of the top-level package. -exports_files(["virtual_chunked.cc"]) - -# Simple forwarding target. -tensorstore_cc_library( - name = "virtual_chunked", - deps = ["//tensorstore:virtual_chunked"], -) - filegroup( name = "doc_sources", srcs = [], ) +tensorstore_cc_library( + name = "virtual_chunked", + srcs = ["virtual_chunked.cc"], + hdrs = ["//tensorstore:virtual_chunked.h"], + deps = [ + "//tensorstore", + "//tensorstore:array", + "//tensorstore:box", + "//tensorstore:context", + "//tensorstore:staleness_bound", + "//tensorstore:transaction", + "//tensorstore/driver", + "//tensorstore/driver:chunk_cache_driver", + "//tensorstore/index_space:index_transform", + "//tensorstore/internal:data_copy_concurrency_resource", + "//tensorstore/internal/cache:cache_pool_resource", + "//tensorstore/internal/cache:chunk_cache", + "//tensorstore/kvstore:generation", + "//tensorstore/serialization", + "//tensorstore/serialization:absl_time", + "//tensorstore/serialization:function", + "//tensorstore/util:executor", + "//tensorstore/util:option", + "//tensorstore/util/garbage_collection", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/status", + "@com_google_absl//absl/time", + ], +) + tensorstore_cc_test( name = "virtual_chunked_test", srcs = ["virtual_chunked_test.cc"], diff --git a/tensorstore/index_space/BUILD b/tensorstore/index_space/BUILD index f4c3277e3..2c799ef50 100644 --- a/tensorstore/index_space/BUILD +++ b/tensorstore/index_space/BUILD @@ -193,6 +193,7 @@ tensorstore_cc_library( "//tensorstore/util:division", "//tensorstore/util:quote_string", "//tensorstore/util:result", + "//tensorstore/util:span", "//tensorstore/util:status", "//tensorstore/util:str_cat", "@com_google_absl//absl/status", diff --git a/tensorstore/index_space/dimension_identifier.cc b/tensorstore/index_space/dimension_identifier.cc index 4b3710f9a..6f74b8134 100644 --- a/tensorstore/index_space/dimension_identifier.cc +++ b/tensorstore/index_space/dimension_identifier.cc @@ -14,13 +14,20 @@ #include "tensorstore/index_space/dimension_identifier.h" +#include +#include +#include #include // NOLINT +#include #include "absl/status/status.h" #include "absl/strings/str_join.h" +#include "tensorstore/index.h" +#include "tensorstore/index_space/dimension_index_buffer.h" #include "tensorstore/util/division.h" #include "tensorstore/util/quote_string.h" #include "tensorstore/util/result.h" +#include "tensorstore/util/span.h" #include "tensorstore/util/status.h" #include "tensorstore/util/str_cat.h" @@ -55,8 +62,10 @@ Result NormalizeDimensionExclusiveStopIndex( return index >= 0 ? index : index + rank; } -Result NormalizeDimensionLabel(std::string_view label, - span labels) { +namespace { +template +Result NormalizeDimensionLabelImpl(std::string_view label, + span labels) { if (label.empty()) { return absl::InvalidArgumentError( "Dimension cannot be specified by empty label"); @@ -74,6 +83,16 @@ Result NormalizeDimensionLabel(std::string_view label, } return dim; } +} // namespace + +Result NormalizeDimensionLabel(std::string_view label, + span labels) { + return NormalizeDimensionLabelImpl(label, labels); +} +Result NormalizeDimensionLabel( + std::string_view label, span labels) { + return NormalizeDimensionLabelImpl(label, labels); +} Result NormalizeDimensionIdentifier( DimensionIdentifier identifier, span labels) { diff --git a/tensorstore/index_space/dimension_identifier.h b/tensorstore/index_space/dimension_identifier.h index 0bb3c50a6..7f07afba3 100644 --- a/tensorstore/index_space/dimension_identifier.h +++ b/tensorstore/index_space/dimension_identifier.h @@ -141,6 +141,8 @@ Result NormalizeDimensionIndex(DimensionIndex index, /// \relates DimensionIdentifier Result NormalizeDimensionLabel(std::string_view label, span labels); +Result NormalizeDimensionLabel( + std::string_view label, span labels); /// Normalizes a dimension identifier to a dimension index in the range /// ``[0, rank)``. diff --git a/tensorstore/spec.h b/tensorstore/spec.h index 23eecbf87..0acfc3411 100644 --- a/tensorstore/spec.h +++ b/tensorstore/spec.h @@ -38,7 +38,7 @@ #include "tensorstore/open_options.h" #include "tensorstore/schema.h" #include "tensorstore/serialization/fwd.h" -#include "tensorstore/spec_impl.h" +#include "tensorstore/spec_impl.h" // IWYU pragma: export #include "tensorstore/util/garbage_collection/fwd.h" #include "tensorstore/util/option.h" #include "tensorstore/util/result.h" diff --git a/tensorstore/stack.h b/tensorstore/stack.h new file mode 100644 index 000000000..d6b7d2e49 --- /dev/null +++ b/tensorstore/stack.h @@ -0,0 +1,137 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENSORSTORE_STACK_H_ +#define TENSORSTORE_STACK_H_ + +#include +#include +#include +#include + +#include "tensorstore/driver/stack/driver.h" +#include "tensorstore/index.h" +#include "tensorstore/index_space/dimension_identifier.h" +#include "tensorstore/internal/type_traits.h" +#include "tensorstore/open_options.h" +#include "tensorstore/spec.h" +#include "tensorstore/tensorstore.h" +#include "tensorstore/util/option.h" +#include "tensorstore/util/result.h" +#include "tensorstore/util/span.h" + +namespace tensorstore { + +using StackOpenOptions = TransactionalOpenOptions; + +/// Virtually overlays a sequence of `TensorStore` or `Spec` layers within a +/// common domain. +/// +/// \param layers Sequence of layers to overlay. Later layers take precedence. +template >>> +Result>::type> +Overlay(const Layers& layers, StackOpenOptions&& options) { + std::vector layers_internal( + std::begin(layers), std::end(layers)); + TENSORSTORE_ASSIGN_OR_RETURN( + auto handle, + internal_stack::Overlay(layers_internal, std::move(options))); + return {std::in_place, + internal::TensorStoreAccess::Construct< + typename internal_stack::OverlayResult< + internal::remove_cvref_t>::type>( + std::move(handle))}; +} +template >>, + typename... Option> +std::enable_if_t< + IsCompatibleOptionSequence, + Result>::type>> +Overlay(const Layers& layers, Option&&... option) { + TENSORSTORE_INTERNAL_ASSIGN_OPTIONS_OR_RETURN(StackOpenOptions, options, + option) + return Overlay(layers, std::move(options)); +} + +/// Virtually stacks a sequence of `TensorStore` or `Spec` layers along a new +/// dimension. +/// +/// \param layers Sequence of layers to stack. +template >>> +Result>::type> +Stack(const Layers& layers, DimensionIndex stack_dimension, + StackOpenOptions&& options) { + std::vector layers_internal( + std::begin(layers), std::end(layers)); + TENSORSTORE_ASSIGN_OR_RETURN( + auto handle, internal_stack::Stack(layers_internal, stack_dimension, + std::move(options))); + return {std::in_place, + internal::TensorStoreAccess::Construct< + typename internal_stack::StackResult< + internal::remove_cvref_t>::type>( + std::move(handle))}; +} +template >>, + typename... Option> +std::enable_if_t< + IsCompatibleOptionSequence, + Result>::type>> +Stack(const Layers& layers, DimensionIndex stack_dimension, + Option&&... option) { + TENSORSTORE_INTERNAL_ASSIGN_OPTIONS_OR_RETURN(StackOpenOptions, options, + option) + return Stack(layers, stack_dimension, std::move(options)); +} + +/// Virtually concatenates a sequence of `TensorStore` or `Spec` layers along an +/// existing dimension. +/// +/// \param layers Sequence of layers to concatenate. +template >>> +Result>::type> +Concat(const Layers& layers, DimensionIdentifier concat_dimension, + StackOpenOptions&& options) { + std::vector layers_internal( + std::begin(layers), std::end(layers)); + TENSORSTORE_ASSIGN_OR_RETURN( + auto handle, internal_stack::Concat(layers_internal, concat_dimension, + std::move(options))); + return {std::in_place, + internal::TensorStoreAccess::Construct< + typename internal_stack::OverlayResult< + internal::remove_cvref_t>::type>( + std::move(handle))}; +} +template >>, + typename... Option> +std::enable_if_t< + IsCompatibleOptionSequence, + Result>::type>> +Concat(const Layers& layers, DimensionIdentifier concat_dimension, + Option&&... option) { + TENSORSTORE_INTERNAL_ASSIGN_OPTIONS_OR_RETURN(StackOpenOptions, options, + option) + return Concat(layers, concat_dimension, std::move(options)); +} + +} // namespace tensorstore + +#endif // TENSORSTORE_STACK_H_ diff --git a/tensorstore/tensorstore.h b/tensorstore/tensorstore.h index 365657be6..20c5407f6 100644 --- a/tensorstore/tensorstore.h +++ b/tensorstore/tensorstore.h @@ -437,6 +437,20 @@ ModeCast(SourceRef&& source) { std::forward(source)); } +/// Changes the `ReadWriteMode` of `store` to the specified `new_mode`. +/// +/// If `new_mode == ReadWriteMode::dynamic`, the existing mode is unchanged. +/// +/// \relates TensorStore +template +Result> ModeCast( + TensorStore store, ReadWriteMode new_mode) { + Result> result{std::move(store)}; + TENSORSTORE_RETURN_IF_ERROR(internal::SetReadWriteMode( + internal::TensorStoreAccess::handle(*result), new_mode)); + return result; +} + /// Read-only `TensorStore` alias. /// /// \relates TensorStore