From ca1726601393cb6d3555b59af8a59b674cb54da4 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 6 Feb 2024 00:41:34 -0500 Subject: [PATCH 1/4] only wrap BackendArrays in CopyOnWriteArray --- xarray/backends/api.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index 670a0ec6d68..8216753eb11 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -236,7 +236,10 @@ def _protect_dataset_variables_inplace(dataset, cache): for name, variable in dataset.variables.items(): if name not in dataset._indexes: # no need to protect IndexVariable objects - data = indexing.CopyOnWriteArray(variable._data) + if isinstance(variable._data, backends.BackendArray): + data = indexing.CopyOnWriteArray(variable._data) + else: + data = variable.data if cache: data = indexing.MemoryCachedArray(data) variable.data = data From 99a37be0e992a05ceca56474bda429fc1036336d Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 6 Feb 2024 00:41:34 -0500 Subject: [PATCH 2/4] only wrap BackendArrays in CopyOnWriteArray --- xarray/backends/api.py | 9 ++++++++- xarray/core/indexing.py | 7 +++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index 670a0ec6d68..467535e6a4d 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -236,9 +236,16 @@ def _protect_dataset_variables_inplace(dataset, cache): for name, variable in dataset.variables.items(): if name not in dataset._indexes: # no need to protect IndexVariable objects - data = indexing.CopyOnWriteArray(variable._data) + + if isinstance(variable._data, backends.BackendArray): + # only need to protect arrays that were lazy-loaded from disk + data = indexing.CopyOnWriteArray(variable._data) + else: + data = variable.data + if cache: data = indexing.MemoryCachedArray(data) + variable.data = data diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 6e6ce01a41f..a1958c0f637 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -652,6 +652,13 @@ def _wrap_numpy_scalars(array): class CopyOnWriteArray(ExplicitlyIndexedNDArrayMixin): + """ + Ensures that an attempt to write to this array only writes to a copy + of whatever might still be stored on disk. + + Used to wrap any array that inherits from BackendArray. + """ + __slots__ = ("array", "_copied") def __init__(self, array): From c80fefd02d0bf2f1490c707f0ac4876d84dd74ff Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 6 Feb 2024 00:59:31 -0500 Subject: [PATCH 3/4] document how CopyOnWrite happens for BackendArray subclasses --- doc/internals/how-to-add-new-backend.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/doc/internals/how-to-add-new-backend.rst b/doc/internals/how-to-add-new-backend.rst index 4352dd3df5b..447d0477462 100644 --- a/doc/internals/how-to-add-new-backend.rst +++ b/doc/internals/how-to-add-new-backend.rst @@ -421,6 +421,9 @@ That implies that all the reference to open files should be dropped. For opening files, we therefore suggest to use the helper class provided by Xarray :py:class:`~xarray.backends.CachingFileManager`. +Note that any ``BackendArray`` subclass will also be protected such that any attempt +to write to this array only writes to a copy of whatever might still be stored on disk. + .. _RST indexing: Indexing examples From 22ec10f07b432f52d9dd6b3ffc431ffc66824a2f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 6 Feb 2024 06:06:20 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- doc/internals/how-to-add-new-backend.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/internals/how-to-add-new-backend.rst b/doc/internals/how-to-add-new-backend.rst index 447d0477462..9a4b7c47323 100644 --- a/doc/internals/how-to-add-new-backend.rst +++ b/doc/internals/how-to-add-new-backend.rst @@ -421,7 +421,7 @@ That implies that all the reference to open files should be dropped. For opening files, we therefore suggest to use the helper class provided by Xarray :py:class:`~xarray.backends.CachingFileManager`. -Note that any ``BackendArray`` subclass will also be protected such that any attempt +Note that any ``BackendArray`` subclass will also be protected such that any attempt to write to this array only writes to a copy of whatever might still be stored on disk. .. _RST indexing: