🚀 (zarr) make save_to_disk write path fsspec-aware (#485)

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- (storage/zarr) make the Zarr backend write path fsspec-aware: `save_to_disk` can now write to any fsspec target (`memory://`, `s3://`, `gs://`, …) in addition to local paths, in both sequential and parallel (`num_proc > 1`) modes. Remote targets require the matching fsspec backend to be installed (e.g. `s3fs` for `s3://`); a missing backend surfaces fsspec's own install hint. No new hard dependency is added (`fsspec` is already pulled in by `zarr`). Scope is limited to the zarr write path; the read path is tracked separately.
 - (storage/writer) add optional `sample_callback` to `save_to_disk`, invoked once per sample right after it is written to disk as `sample_callback(split_name, index, sample_path)`. This lets callers process samples one by one once written instead of waiting for the whole dataset. Currently supported for the `cgns` backend, including parallel writing (`num_proc > 1`), where the callback runs inside the worker processes and must be picklable and process-safe.
 
 ### Fixed

docs/source/tutorials/storage.md

@@ -10,7 +10,28 @@ End‑to‑end workflows for creating, saving, and loading PLAID datasets with t
 
 - **`sample_constructor`** is a simple function that takes a single identifier (of any type) and returns a PLAID `Sample`. The identifier can be an integer, a file path, a string, a tuple — anything that makes sense for your data.
 - **`ids`** is a dictionary mapping split names to **sliceable sequences** of identifiers — anything with `__getitem__` and `__len__` (list, tuple, numpy array, …). PLAID handles iteration, generator creation, and parallel sharding internally.
-- **`save_to_disk`** writes a dataset locally; **`push_to_hub`** uploads it to Hugging Face Hub.
+- **`save_to_disk`** writes a dataset to `output_folder`; **`push_to_hub`** uploads it to Hugging Face Hub.
+
+!!! info "Remote (fsspec) write targets — zarr backend"
+    With the **zarr** backend, `output_folder` may be a plain local path **or any
+    [fsspec](https://filesystem-spec.readthedocs.io/) URL** (`memory://`,
+    `s3://bucket/prefix`, `gs://…`, …), in both sequential and parallel
+    (`num_proc > 1`) modes:
+
+    ```python
+    save_to_disk(
+        output_folder="s3://my-bucket/datasets/shapenetcar",
+        sample_constructor=sample_constructor,
+        ids=ids,
+        backend="zarr",
+        num_proc=N_PROC,
+    )
+    ```
+
+    Remote targets require the matching fsspec backend to be installed
+    (`s3fs` for `s3://`, `gcsfs` for `gs://`, …); a missing backend raises an
+    `ImportError` carrying fsspec's own install hint. The other backends
+    (`cgns`, `hf_datasets`) currently expect a local `output_folder`.
 - **`init_from_disk`** / **`download_from_hub`** / **`init_streaming_from_hub`** load datasets back into PLAID.
 - Backend converters turn raw backend samples into PLAID `Sample` objects.
 

src/plaid/storage/zarr/writer.py

@@ -17,6 +17,7 @@
 from pathlib import Path
 from typing import Any, Callable, Generator, Optional, Union
 
+import fsspec
 import numpy as np
 import yaml
 import zarr
@@ -30,6 +31,69 @@
 from ...infos import Infos
 
 
+def _is_local_target(target: Union[str, Path]) -> bool:
+    """Return whether ``target`` resolves to the local filesystem.
+
+    A plain path (absolute or relative) or a ``file://`` URL is local; any other
+    fsspec protocol (``memory://``, ``s3://``, ``gs://``, ``http://``, ...) is
+    remote.
+
+    Args:
+        target (Union[str, Path]): Output folder path or fsspec URL.
+
+    Returns:
+        bool: ``True`` if the target lives on the local filesystem.
+
+    Raises:
+        ImportError: If ``target`` uses a protocol whose fsspec handler is not
+            installed (e.g. ``s3://`` without ``s3fs``). The original fsspec
+            message is preserved so the user knows which extra to install.
+    """
+    fs, _ = fsspec.core.url_to_fs(str(target))
+    protocols = fs.protocol
+    if isinstance(protocols, str):
+        protocols = (protocols,)
+    return any(p in ("file", "local") for p in protocols)
+
+
+def _join_target(base: Union[str, Path], *parts: str) -> Union[Path, str]:
+    """Join path components, preserving fsspec URLs.
+
+    ``pathlib.Path`` collapses the ``//`` in a URL (``memory://root`` becomes
+    ``memory:/root``), silently corrupting the target. This helper keeps local
+    targets as ``Path`` (unchanged behavior) and joins remote URLs with plain
+    ``/`` so the protocol separator survives.
+
+    Args:
+        base (Union[str, Path]): Base path or fsspec URL.
+        *parts (str): Additional path components to append.
+
+    Returns:
+        Union[Path, str]: A ``Path`` for local targets, a ``str`` URL otherwise.
+    """
+    if _is_local_target(base):
+        return Path(base).joinpath(*parts)
+    return "/".join([str(base).rstrip("/"), *parts])
+
+
+def _open_split_group(split_target: str, mode: str) -> Any:
+    """Open (or create) the Zarr group for a split, local or remote.
+
+    ``zarr.open_group`` natively resolves an fsspec URL to a ``FsspecStore``, so
+    the same call works for local paths and remote URLs. This wrapper exists so
+    the parallel worker (which runs in a separate process and cannot share an
+    open handle) can reopen the exact same store from a plain string.
+
+    Args:
+        split_target (str): Local path or fsspec URL of the split root.
+        mode (str): Zarr open mode (``"w"``, ``"a"``, ...).
+
+    Returns:
+        Any: An open Zarr group.
+    """
+    return zarr.open_group(split_target, mode=mode)
+
+
 def _auto_chunks(shape: tuple[int, ...], target_n: int) -> tuple[int, ...]:
     """Computes automatic chunk sizes for Zarr arrays based on shape and target size.
 
@@ -134,9 +198,12 @@ def _zarr_worker_batch_job(args) -> int:  # pragma: no cover
     """
     split_root_path, gen_func, var_features_keys, batch, start_index = args
 
-    # split_root = zarr.open_group(split_root_path, mode="a")
-    store = zarr.storage.LocalStore(split_root_path)
-    split_root = zarr.group(store=store)
+    # Reopen the split group from a plain string so it works for both local
+    # paths and fsspec URLs (memory://, s3://, ...). ``zarr.open_group`` builds
+    # a ``LocalStore`` or ``FsspecStore`` from the target automatically; the
+    # previous hardcoded ``zarr.storage.LocalStore`` silently wrote a remote URL
+    # to a literal local directory instead of the intended remote target.
+    split_root = _open_split_group(split_root_path, mode="a")
 
     sample_counter = start_index
     written = 0
@@ -185,15 +252,21 @@ def generate_datasetdict_to_disk(
         None: This function does not return a value; it writes the dataset directly
             to disk.
     """
-    output_folder = Path(output_folder) / "data"
-    output_folder.mkdir(exist_ok=True, parents=True)
+    # Build the ``data`` subfolder in a way that works for local paths and
+    # fsspec URLs alike. For local targets we keep the original ``Path`` + mkdir
+    # behavior; for remote targets we join with ``/`` (``Path`` would collapse
+    # the ``://`` separator) and skip the local ``mkdir`` — fsspec stores create
+    # keys lazily on write and have no notion of an empty directory.
+    data_target = _join_target(output_folder, "data")
+    if _is_local_target(output_folder):
+        Path(data_target).mkdir(exist_ok=True, parents=True)
 
     var_features_keys = list(variable_schema.keys())
     gen_kwargs_ = gen_kwargs or {sn: {} for sn in generators.keys()}
 
     for split_name, gen_func in generators.items():
-        split_root_path = str(output_folder / split_name)
-        _ = zarr.open_group(split_root_path, mode="w")  # create/overwrite
+        split_root_path = str(_join_target(data_target, split_name))
+        _ = _open_split_group(split_root_path, mode="w")  # create/overwrite
 
         batch_ids_list = gen_kwargs_.get(split_name, {}).get("shards_ids", [])
         total_samples = (

tests/storage/test_zarr_fsspec_write.py

@@ -0,0 +1,210 @@
+"""Tests for the fsspec-aware Zarr write path (issue #485).
+
+These tests exercise ``plaid.storage.zarr.writer.generate_datasetdict_to_disk``
+against a remote (non-local) fsspec target using the in-memory filesystem
+(``memory://``). They verify that:
+
+- local and remote targets are correctly discriminated;
+- URL joining preserves the ``://`` protocol separator (``pathlib.Path`` would
+  collapse it);
+- samples are written at the intended remote location in both sequential and
+  parallel (``num_proc > 1``) modes;
+- no literal ``memory:`` directory is created on the local filesystem (the
+  previous ``LocalStore`` implementation wrote remote URLs to a local folder).
+"""
+
+import os
+
+import fsspec
+import numpy as np
+import pytest
+import zarr
+
+from plaid.containers.sample import Sample
+from plaid.storage.zarr.writer import (
+    _is_local_target,
+    _join_target,
+    _open_split_group,
+    generate_datasetdict_to_disk,
+)
+
+# ---------------------------------------------------------------------------
+# Module-level (picklable) helpers required by the parallel writer.
+# ---------------------------------------------------------------------------
+
+_GLOBAL_VALUES = {0: 10.0, 1: 20.0, 2: 30.0, 3: 40.0}
+
+
+def _make_sample(value: float) -> Sample:
+    """Build a minimal sample carrying a single global scalar."""
+    sample = Sample()
+    sample.add_global("myglobal", float(value))
+    return sample
+
+
+class _MemorySampleGenerator:
+    """Picklable generator matching the ``gen_func(shards_ids)`` contract.
+
+    Defined at module level so ``multiprocessing`` can pickle it when
+    ``num_proc > 1``.
+    """
+
+    def __call__(self, shards_ids=None):
+        if shards_ids is None:
+            shards_ids = [[]]
+        for shard in shards_ids:
+            for sample_id in shard:
+                yield _make_sample(_GLOBAL_VALUES[sample_id])
+
+
+_VARIABLE_SCHEMA = {
+    "Global/myglobal": {"dtype": "float64", "ndim": 0},
+    "Global/myglobal_times": {"dtype": "float64", "ndim": 0},
+    "Global": {"dtype": "float64", "ndim": 1},
+    "Global_times": {"dtype": "float64", "ndim": 1},
+}
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for the target helpers.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("target", "expected_local"),
+    [
+        ("/abs/local/path", True),
+        ("relative/local/path", True),
+        ("file:///abs/local/path", True),
+        ("memory://root", False),
+    ],
+)
+def test_is_local_target(target, expected_local):
+    assert _is_local_target(target) is expected_local
+
+
+def test_is_local_target_missing_backend_raises_clear_error():
+    """An unavailable fsspec protocol surfaces the install hint from fsspec."""
+    with pytest.raises(ImportError, match="s3fs"):
+        _is_local_target("s3://bucket/key")
+
+
+def test_join_target_local_returns_path(tmp_path):
+    joined = _join_target(tmp_path, "data", "train")
+    assert joined == tmp_path / "data" / "train"
+
+
+def test_join_target_remote_preserves_protocol():
+    """``pathlib.Path`` would collapse ``memory://root`` to ``memory:/root``."""
+    joined = _join_target("memory://root", "data", "train")
+    assert joined == "memory://root/data/train"
+    # a trailing slash on the base must not produce a doubled separator
+    assert _join_target("memory://root/", "data") == "memory://root/data"
+
+
+# ---------------------------------------------------------------------------
+# End-to-end write to a remote (memory://) target.
+# ---------------------------------------------------------------------------
+
+
+def _clear_memory_fs():
+    """Reset the shared in-memory filesystem between test runs."""
+    fs = fsspec.filesystem("memory")
+    fs.store.clear()
+    if hasattr(fs, "pseudo_dirs"):
+        fs.pseudo_dirs.clear()
+        fs.pseudo_dirs.append("")
+
+
+def _read_split_groups(output_url: str, split_name: str):
+    data_target = _join_target(_join_target(output_url, "data"), split_name)
+    return zarr.open_group(str(data_target), mode="r")
+
+
+@pytest.mark.parametrize("num_proc", [1, 2])
+def test_generate_datasetdict_to_disk_parallel_still_works_locally(num_proc, tmp_path):
+    """The store-selection change must not regress the local parallel path.
+
+    ``memory://`` is per-process, so a ``num_proc > 1`` write cannot be read
+    back from the parent process; the parallel reopen is therefore validated on
+    a shared local target instead.
+    """
+    output_folder = tmp_path / f"dataset_np{num_proc}"
+
+    if num_proc == 1:
+        shards_ids = [[0, 1]]
+    else:
+        # one shard per worker to exercise the parallel reopen path
+        shards_ids = [[0], [1]]
+
+    generate_datasetdict_to_disk(
+        output_folder=output_folder,
+        generators={"train": _MemorySampleGenerator()},
+        variable_schema=_VARIABLE_SCHEMA,
+        gen_kwargs={"train": {"shards_ids": shards_ids}},
+        num_proc=num_proc,
+        verbose=False,
+    )
+
+    group = zarr.open_group(str(output_folder / "data" / "train"), mode="r")
+    assert sorted(group.group_keys()) == [
+        "sample_000000000",
+        "sample_000000001",
+    ]
+
+
+def test_generate_datasetdict_to_disk_writes_to_memory_fs():
+    """Sequential write to a remote (memory://) fsspec target.
+
+    A parallel (``num_proc > 1``) equivalent is intentionally omitted here: the
+    ``memory://`` filesystem is per-process, so samples written by worker
+    processes are not visible from the parent. The parallel store-selection path
+    is covered by ``test_generate_datasetdict_to_disk_parallel_still_works_locally``
+    and ``test_open_split_group_roundtrip_on_memory_fs``.
+    """
+    _clear_memory_fs()
+    output_url = "memory://dataset_seq"
+
+    generate_datasetdict_to_disk(
+        output_folder=output_url,
+        generators={"train": _MemorySampleGenerator()},
+        variable_schema=_VARIABLE_SCHEMA,
+        gen_kwargs={"train": {"shards_ids": [[0, 1]]}},
+        num_proc=1,
+        verbose=False,
+    )
+
+    # samples landed at the intended remote location
+    group = _read_split_groups(output_url, "train")
+    sample_groups = sorted(group.group_keys())
+    assert sample_groups == ["sample_000000000", "sample_000000001"]
+
+    # each sample carries the flattened global feature and is readable back
+    for sample_name in sample_groups:
+        array_keys = list(group[sample_name].array_keys())
+        assert "Global__myglobal" in array_keys
+
+    # the remote write must NOT create a literal local "memory:" directory
+    assert not os.path.exists("memory:")
+    assert not os.path.exists("memory:/")
+
+
+def test_open_split_group_roundtrip_on_memory_fs():
+    """Create (mode 'w') then reopen (mode 'a') a split group on memory://.
+
+    This mirrors the sequential-create / parallel-worker-reopen handshake used
+    by ``generate_datasetdict_to_disk``.
+    """
+    _clear_memory_fs()
+    target = str(_join_target(_join_target("memory://reopen", "data"), "train"))
+
+    created = _open_split_group(target, mode="w")
+    created.create_group("sample_000000000").create_array("x", data=np.arange(3))
+
+    reopened = _open_split_group(target, mode="a")
+    reopened.create_group("sample_000000001").create_array("x", data=np.arange(3) + 10)
+
+    assert sorted(reopened.group_keys()) == [
+        "sample_000000000",
+        "sample_000000001",
+    ]