Add a tmpdir for allowing to specify a dir for temps during full indexes creation

Author: FrancescAlted

src/blosc2/indexing.py

@@ -272,6 +272,18 @@ def _tmpdir_for_array(array: blosc2.NDArray) -> str | None:
     return None
 
 
+def _resolve_full_index_tmpdir(array: blosc2.NDArray, tmpdir: str | None) -> str | None:
+    """Resolve the workspace for temporary files used by OOC full index builds.
+
+    An explicit ``tmpdir`` wins. Otherwise, persistent arrays default to the
+    directory that stores the array so temporary sidecars stay on the same
+    filesystem. In-memory arrays fall back to the system temporary directory.
+    """
+    if tmpdir is not None:
+        return tmpdir
+    return _tmpdir_for_array(array)
+
+
 def _load_store(array: blosc2.NDArray) -> dict:
     if _is_persistent_array(array):
         key = _array_key(array)
@@ -3006,8 +3018,16 @@ def create_index(
     persistent: bool | None = None,
     in_mem: bool = False,
     name: str | None = None,
+    tmpdir: str | None = None,
     **kwargs,
 ) -> dict:
+    """Create an index descriptor for a 1-D array or structured field.
+
+    Parameters are equivalent to :meth:`blosc2.NDArray.create_index`.
+    ``tmpdir`` controls where temporary files for out-of-core ``kind="full"``
+    builds are created. If ``None``, persistent arrays default to their own
+    directory and in-memory arrays use the system temporary directory.
+    """
     cparams = _normalize_index_cparams(kwargs.pop("cparams", None))
     del kwargs
     dtype = _validate_index_target(array, field)
@@ -3034,7 +3054,7 @@ def create_index(
         full = None
         if kind == "full":
             with tempfile.TemporaryDirectory(
-                prefix="blosc2-index-ooc-", dir=_tmpdir_for_array(array)
+                prefix="blosc2-index-ooc-", dir=_resolve_full_index_tmpdir(array, tmpdir)
             ) as tmpdir:
                 full = _build_full_descriptor_ooc(
                     array, target, token, kind, dtype, persistent, Path(tmpdir), cparams
@@ -3201,8 +3221,15 @@ def create_expr_index(
     return _copy_descriptor(descriptor)
 
 
-def create_csindex(array: blosc2.NDArray, field: str | None = None, **kwargs) -> dict:
-    return create_index(array, field=field, kind="full", **kwargs)
+def create_csindex(
+    array: blosc2.NDArray, field: str | None = None, tmpdir: str | None = None, **kwargs
+) -> dict:
+    """Create a full sorted index descriptor.
+
+    This is shorthand for :func:`create_index` with ``kind="full"``.
+    ``tmpdir`` has the same meaning as in :func:`create_index`.
+    """
+    return create_index(array, field=field, kind="full", tmpdir=tmpdir, **kwargs)
 
 
 def _resolve_index_token(store: dict, field: str | None, name: str | None) -> str:

src/blosc2/ndarray.py

@@ -4763,6 +4763,7 @@ def create_index(
         persistent: bool | None = None,
         in_mem: bool = False,
         name: str | None = None,
+        tmpdir: str | None = None,
         **kwargs: Any,
     ) -> dict:
         """Create an index for a 1-D array or structured field.
@@ -4791,6 +4792,12 @@ def create_index(
             Optional logical label stored in the descriptor. Index identity is
             still driven by the target field, so creating another index on the
             same field replaces the previous one.
+        tmpdir : str or None, optional
+            Directory to use for temporary files during out-of-core
+            ``kind="full"`` builds. If ``None``, persistent arrays use the same
+            directory as the array being indexed so temporaries stay on the
+            same filesystem. In-memory arrays fall back to the system
+            temporary directory.
         kwargs : dict, optional
             Keyword arguments forwarded to the index builder. At the moment the
             supported option is ``cparams``. Pass ``cparams`` to control the
@@ -4822,19 +4829,22 @@ def create_index(
             persistent=persistent,
             in_mem=in_mem,
             name=name,
+            tmpdir=tmpdir,
             **kwargs,
         )
 
-    def create_csindex(self, field: str | None = None, **kwargs: Any) -> dict:
+    def create_csindex(self, field: str | None = None, tmpdir: str | None = None, **kwargs: Any) -> dict:
         """Create a fully sorted index for a 1-D array or structured field.
 
         This is a convenience wrapper for ``create_index(kind="full")`` and is
         the required index tier for direct ordered reuse in
         ``sort(order=...)``, ``indices(order=...)``, and ``itersorted(...)``.
+        ``tmpdir`` controls where temporary files for out-of-core full builds
+        are created, following the same rules as :meth:`create_index`.
         """
         from . import indexing
 
-        return indexing.create_csindex(self, field=field, **kwargs)
+        return indexing.create_csindex(self, field=field, tmpdir=tmpdir, **kwargs)
 
     def create_expr_index(
         self,

tests/ndarray/test_indexing.py

@@ -1123,6 +1123,50 @@ def test_forced_ooc_full_index_merge_preserves_sorted_sidecars(monkeypatch, tmp_
     np.testing.assert_array_equal(values_sidecar[:], data[positions_sidecar[:]])
 
 
+def test_create_index_full_ooc_defaults_tmpdir_to_array_directory(monkeypatch, tmp_path):
+    path = tmp_path / "default_tmpdir_full.b2nd"
+    data = np.arange(4096, dtype=np.int64)
+    arr = blosc2.asarray(data, urlpath=path, mode="w", chunks=(256,), blocks=(64,))
+
+    recorded = {}
+    real_temporary_directory = indexing.tempfile.TemporaryDirectory
+
+    def tracking_temporary_directory(*args, **kwargs):
+        recorded["dir"] = kwargs.get("dir")
+        return real_temporary_directory(*args, **kwargs)
+
+    monkeypatch.setattr(indexing.tempfile, "TemporaryDirectory", tracking_temporary_directory)
+
+    descriptor = arr.create_index(kind="full")
+
+    assert descriptor["ooc"] is True
+    assert recorded["dir"] == str(path.parent.resolve())
+
+
+def test_create_csindex_full_ooc_uses_explicit_tmpdir(monkeypatch, tmp_path):
+    path = tmp_path / "explicit_tmpdir_full.b2nd"
+    custom_tmpdir = tmp_path / "custom-index-tmp"
+    custom_tmpdir.mkdir()
+    dtype = np.dtype([("a", np.int64), ("payload", np.int32)])
+    data = np.zeros(4096, dtype=dtype)
+    data["a"] = np.arange(data.shape[0], dtype=np.int64)
+    arr = blosc2.asarray(data, urlpath=path, mode="w", chunks=(256,), blocks=(64,))
+
+    recorded = {}
+    real_temporary_directory = indexing.tempfile.TemporaryDirectory
+
+    def tracking_temporary_directory(*args, **kwargs):
+        recorded["dir"] = kwargs.get("dir")
+        return real_temporary_directory(*args, **kwargs)
+
+    monkeypatch.setattr(indexing.tempfile, "TemporaryDirectory", tracking_temporary_directory)
+
+    descriptor = arr.create_csindex("a", tmpdir=str(custom_tmpdir))
+
+    assert descriptor["ooc"] is True
+    assert recorded["dir"] == str(custom_tmpdir)
+
+
 @pytest.mark.parametrize("persistent", [False, True])
 def test_compact_full_index_rebuilds_navigation_without_whole_loading(monkeypatch, tmp_path, persistent):
     dtype = np.dtype([("a", np.int64), ("b", np.int64)])