Add GatherToZarr and save pipeline node extension data to disk

src/spikeinterface/core/analyzer_extension_core.py

@@ -1552,13 +1552,23 @@ def _set_params(self, **kwargs):
     def _run(self, verbose=False, **job_kwargs):
         from spikeinterface.core.node_pipeline import run_node_pipeline
 
-        # TODO: should we save directly to npy in binary_folder format / or to zarr?
-        # if self.sorting_analyzer.format == "binary_folder":
-        #     gather_mode = "npy"
-        #     extension_folder = self.sorting_analyzer.folder / "extenstions" / self.extension_name
-        #     gather_kwargs = {"folder": extension_folder}
-        gather_mode = "memory"
-        gather_kwargs = {}
+        # gather results directly to the final on-disk location (one npy file / zarr dataset per
+        # nodepipeline variable) to avoid an extra in-memory copy. This is only done when we are
+        # actually saving to a disk format (see AnalyzerExtension.run()); otherwise gather in memory.
+        gather_to_disk = self._save_to_disk and self.format in ("binary_folder", "zarr")
+        if gather_to_disk:
+            extension_folder = self.sorting_analyzer.folder / "extensions" / self.extension_name
+            names = self.nodepipeline_variables
+            if self.format == "binary_folder":
+                gather_mode = "npy"
+                dest = [extension_folder / f"{name}.npy" for name in names]
+            else:
+                gather_mode = "zarr"
+                dest = [extension_folder / name for name in names]
+        else:
+            gather_mode = "memory"
+            dest = None
+            names = None
 
         job_kwargs = fix_job_kwargs(job_kwargs)
         nodes = self.get_pipeline_nodes()
@@ -1568,7 +1578,8 @@ def _run(self, verbose=False, **job_kwargs):
             job_kwargs=job_kwargs,
             job_name=self.extension_name,
             gather_mode=gather_mode,
-            gather_kwargs=gather_kwargs,
+            dest=dest,
+            names=names,
             verbose=False,
         )
         if isinstance(data, tuple):
@@ -1619,6 +1630,11 @@ def _get_data(self, outputs="numpy", concatenated=False, return_data_name=None,
                 ), f"return_data_name {return_data_name} not in nodepipeline_variables {self.nodepipeline_variables}"
 
         all_data = self.data[return_data_name]
+        # data gathered directly into a zarr store (e.g. by the node pipeline) is kept as a
+        # zarr.Array handle. On a non-lazy analyzer we materialize it to a numpy array (this mirrors
+        # the non-lazy load convention). A memmap is an np.ndarray subclass so it is left untouched.
+        if not self.sorting_analyzer._lazy and not isinstance(all_data, np.ndarray):
+            all_data = np.asarray(all_data)
         keep_mask = None
         if periods is not None:
             keep_mask = select_sorting_periods_mask(