feat: Support for masked arrays (fixes zarr-developers#3744)

Author: Test User

changes/3744.feature.md

@@ -0,0 +1 @@
+Add support for numpy masked arrays (numpy.ma.MaskedArray) in zarr.array(). When a masked array is provided, it is automatically converted to a filled array, with a warning that the mask is not preserved. Users who need to preserve mask information should use separate arrays or structured dtypes for storing both data and mask information.

src/zarr/abc/store.py

@@ -187,6 +187,26 @@ def _check_writable(self) -> None:
         if self.read_only:
             raise ValueError("store was opened in read-only mode and does not support writing")
 
+    def _ensure_buffer(self, value: Buffer | bytes) -> Buffer:
+        """Convert bytes to Buffer if needed.
+
+        Parameters
+        ----------
+        value : Buffer or bytes
+            The value to ensure is a Buffer.
+
+        Returns
+        -------
+        Buffer
+            The input value if it's already a Buffer, or a new Buffer created from bytes.
+        """
+        # avoid circular import
+        from zarr.core.buffer import Buffer
+
+        if isinstance(value, bytes):
+            return Buffer.from_bytes(value)
+        return value
+
     @abstractmethod
     def __eq__(self, value: object) -> bool:
         """Equality comparison."""
@@ -465,24 +485,28 @@ def supports_writes(self) -> bool:
         ...
 
     @abstractmethod
-    async def set(self, key: str, value: Buffer) -> None:
+    async def set(self, key: str, value: Buffer | bytes) -> None:
         """Store a (key, value) pair.
 
         Parameters
         ----------
         key : str
-        value : Buffer
+        value : Buffer or bytes
+            The value to store. If bytes are provided, they will be converted
+            to a Buffer internally.
         """
         ...
 
-    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
+    async def set_if_not_exists(self, key: str, value: Buffer | bytes) -> None:
         """
         Store a key to ``value`` if the key is not already present.
 
         Parameters
         ----------
         key : str
-        value : Buffer
+        value : Buffer or bytes
+            The value to store. If bytes are provided, they will be converted
+            to a Buffer internally.
         """
         # Note for implementers: the default implementation provided here
         # is not safe for concurrent writers. There's a race condition between
@@ -491,7 +515,7 @@ async def set_if_not_exists(self, key: str, value: Buffer) -> None:
         if not await self.exists(key):
             await self.set(key, value)
 
-    async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None:
+    async def _set_many(self, values: Iterable[tuple[str, Buffer | bytes]]) -> None:
         """
         Insert multiple (key, value) pairs into storage.
         """

src/zarr/api/asynchronous.py

@@ -621,6 +621,16 @@ async def array(data: npt.ArrayLike | AnyArray, **kwargs: Any) -> AnyAsyncArray:
     if isinstance(data, Array):
         return await from_array(data=data, **kwargs)
 
+    # Handle masked arrays by converting to filled array
+    if isinstance(data, np.ma.MaskedArray):
+        warnings.warn(
+            "Masked arrays are not fully supported in Zarr. The mask will not be preserved. "
+            "Consider using zarr's structured dtype or a separate array for the mask if you need to preserve it.",
+            UserWarning,
+            stacklevel=2,
+        )
+        data = cast(np.ndarray, data.filled())
+
     # ensure data is array-like
     if not hasattr(data, "shape") or not hasattr(data, "dtype"):
         data = np.asanyarray(data)

src/zarr/experimental/cache_store.py

@@ -357,18 +357,21 @@ async def get(
         else:
             return await self._get_try_cache(key, prototype, byte_range)
 
-    async def set(self, key: str, value: Buffer) -> None:
+    async def set(self, key: str, value: Buffer | bytes) -> None:
         """
         Store data in the underlying store and optionally in cache.
 
         Parameters
         ----------
         key : str
             The key to store under
-        value : Buffer
-            The data to store
+        value : Buffer or bytes
+            The data to store. If bytes are provided, they will be converted
+            to a Buffer internally.
         """
         await super().set(key, value)
+        # Ensure value is a Buffer for caching
+        value = self._ensure_buffer(value)
         # Invalidate all cached byte-range entries (source data changed)
         async with self._state.lock:
             self._invalidate_range_entries(key)

src/zarr/storage/_common.py

@@ -173,14 +173,14 @@ async def get(
             prototype = default_buffer_prototype()
         return await self.store.get(self.path, prototype=prototype, byte_range=byte_range)
 
-    async def set(self, value: Buffer) -> None:
+    async def set(self, value: Buffer | bytes) -> None:
         """
         Write bytes to the store.
 
         Parameters
         ----------
-        value : Buffer
-            The buffer to write.
+        value : Buffer or bytes
+            The buffer or bytes to write.
         """
         await self.store.set(self.path, value)
 
@@ -201,14 +201,14 @@ async def delete_dir(self) -> None:
         """
         await self.store.delete_dir(self.path)
 
-    async def set_if_not_exists(self, default: Buffer) -> None:
+    async def set_if_not_exists(self, default: Buffer | bytes) -> None:
         """
         Store a key to ``value`` if the key is not already present.
 
         Parameters
         ----------
-        default : Buffer
-            The buffer to store if the key is not already present.
+        default : Buffer or bytes
+            The buffer or bytes to store if the key is not already present.
         """
         await self.store.set_if_not_exists(self.path, default)
 

src/zarr/storage/_fsspec.py

@@ -318,17 +318,14 @@ async def get(
     async def set(
         self,
         key: str,
-        value: Buffer,
+        value: Buffer | bytes,
         byte_range: tuple[int, int] | None = None,
     ) -> None:
         # docstring inherited
         if not self._is_open:
             await self._open()
         self._check_writable()
-        if not isinstance(value, Buffer):
-            raise TypeError(
-                f"FsspecStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
-            )
+        value = self._ensure_buffer(value)
         path = _dereference_path(self.path, key)
         # write data
         if byte_range:

src/zarr/storage/_local.py

@@ -269,26 +269,23 @@ async def get_partial_values(
             args.append((_get, path, prototype, byte_range))
         return await concurrent_map(args, asyncio.to_thread, limit=None)  # TODO: fix limit
 
-    async def set(self, key: str, value: Buffer) -> None:
+    async def set(self, key: str, value: Buffer | bytes) -> None:
         # docstring inherited
         return await self._set(key, value)
 
-    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
+    async def set_if_not_exists(self, key: str, value: Buffer | bytes) -> None:
         # docstring inherited
         try:
             return await self._set(key, value, exclusive=True)
         except FileExistsError:
             pass
 
-    async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
+    async def _set(self, key: str, value: Buffer | bytes, exclusive: bool = False) -> None:
         if not self._is_open:
             await self._open()
         self._check_writable()
         assert isinstance(key, str)
-        if not isinstance(value, Buffer):
-            raise TypeError(
-                f"LocalStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
-            )
+        value = self._ensure_buffer(value)
         path = self.root / key
         await asyncio.to_thread(_put, path, value, exclusive=exclusive)
 

src/zarr/storage/_logging.py

@@ -190,12 +190,12 @@ async def exists(self, key: str) -> bool:
         with self.log(key):
             return await self._store.exists(key)
 
-    async def set(self, key: str, value: Buffer) -> None:
+    async def set(self, key: str, value: Buffer | bytes) -> None:
         # docstring inherited
         with self.log(key):
             return await self._store.set(key=key, value=value)
 
-    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
+    async def set_if_not_exists(self, key: str, value: Buffer | bytes) -> None:
         # docstring inherited
         with self.log(key):
             return await self._store.set_if_not_exists(key=key, value=value)

src/zarr/storage/_memory.py

@@ -156,26 +156,24 @@ async def exists(self, key: str) -> bool:
         # docstring inherited
         return key in self._store_dict
 
-    async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None:
+    async def set(self, key: str, value: Buffer | bytes, byte_range: tuple[int, int] | None = None) -> None:
         # docstring inherited
         self._check_writable()
         await self._ensure_open()
         assert isinstance(key, str)
-        if not isinstance(value, Buffer):
-            raise TypeError(
-                f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
-            )
+        value = self._ensure_buffer(value)
         if byte_range is not None:
             buf = self._store_dict[key]
             buf[byte_range[0] : byte_range[1]] = value
             self._store_dict[key] = buf
         else:
             self._store_dict[key] = value
 
-    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
+    async def set_if_not_exists(self, key: str, value: Buffer | bytes) -> None:
         # docstring inherited
         self._check_writable()
         await self._ensure_open()
+        value = self._ensure_buffer(value)
         self._store_dict.setdefault(key, value)
 
     async def delete(self, key: str) -> None:
@@ -506,14 +504,11 @@ def from_dict(cls, store_dict: MutableMapping[str, Buffer]) -> Self:
         gpu_store_dict = {k: gpu.Buffer.from_buffer(v) for k, v in store_dict.items()}
         return cls(gpu_store_dict)
 
-    async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None:
+    async def set(self, key: str, value: Buffer | bytes, byte_range: tuple[int, int] | None = None) -> None:
         # docstring inherited
         self._check_writable()
         assert isinstance(key, str)
-        if not isinstance(value, Buffer):
-            raise TypeError(
-                f"GpuMemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
-            )
+        value = self._ensure_buffer(value)
         # Convert to gpu.Buffer
         gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value)
         await super().set(key, gpu_value, byte_range=byte_range)

src/zarr/storage/_obstore.py

@@ -166,20 +166,21 @@ def supports_writes(self) -> bool:
         # docstring inherited
         return True
 
-    async def set(self, key: str, value: Buffer) -> None:
+    async def set(self, key: str, value: Buffer | bytes) -> None:
         # docstring inherited
         import obstore as obs
 
         self._check_writable()
-
+        value = self._ensure_buffer(value)
         buf = value.as_buffer_like()
         await obs.put_async(self.store, key, buf)
 
-    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
+    async def set_if_not_exists(self, key: str, value: Buffer | bytes) -> None:
         # docstring inherited
         import obstore as obs
 
         self._check_writable()
+        value = self._ensure_buffer(value)
         buf = value.as_buffer_like()
         with contextlib.suppress(obs.exceptions.AlreadyExistsError):
             await obs.put_async(self.store, key, buf, mode="create")