Fetch Argo data massively

CITATION.cff

@@ -40,6 +40,4 @@ preferred-citation:
   volume: 5
   issue: 53
   page: 2425
-
-# This CITATION.cff file was generated with cffinit.
-# Visit https://bit.ly/cffinit to generate yours today!
+swhid: "swh:1:dir:3d568702945207a22ac35e038256e4be88d20f17;origin=https://github.com/euroargodev/argopy;visit=swh:1:snp:b9840a337629714e13028db36df96ad323093e30;anchor=swh:1:rev:a4d92889fdd2ade2735569ac7968bf424abb4fc9"

argopy/stores/implementations/local.py

@@ -108,7 +108,7 @@ def open_dataset(
         path: str
             The local path of the netcdf file to open
 
-        errors:
+        errors: str, "raise", "ignore", "silent", default = "raise"
 
         lazy: bool, default=False
             Define if we should try to open the netcdf dataset lazily or not
@@ -152,6 +152,12 @@ def load_in_memory(path, errors="raise", xr_opts={}):
                     raise e
                 elif errors == "ignore":
                     log.error("FileNotFoundError raised from: %s" % path)
+            except Exception as e:
+                if errors == "raise":
+                    raise e
+                elif errors == "ignore":
+                    log.error(f"{str(e)} raised from: {path}")
+
             return None, None
 
         def load_lazily(path, errors="raise", xr_opts={}, akoverwrite: bool = False):

argopy/stores/index/implementations/pyarrow/index.py

@@ -19,6 +19,7 @@
 from argopy.options import OPTIONS
 from argopy.errors import DataNotFound, InvalidDatasetStructure
 from argopy.utils.checkers import check_index_cols
+from argopy.utils.chunking import chunk_wmos
 from argopy.utils.geo import conv_lon
 from argopy.utils.format import mono2multi
 from argopy.stores.index.spec import ArgoIndexStoreProto
@@ -406,7 +407,7 @@ def xmax(xtble):
             ]
         else:
             if not hasattr(self, "index"):
-                self.load()
+                self.load(nrows=self._nrows_index)
             return [
                 xmin(self.index),
                 xmax(self.index),
@@ -424,6 +425,8 @@ def read_files(self, index: bool = False, multi: bool = False) -> List[str]:
                 for f in self.search["file"]
             ]
         else:
+            if not hasattr(self, "index"):
+                self.load(nrows=self._nrows_index)
             flist = [
                 sep.join(["dac", f.as_py().replace("/", sep)])
                 for f in self.index["file"]
@@ -459,40 +462,52 @@ def records_per_wmo(self, index=False):
         """Return the number of records per unique WMOs in search results
 
         Fall back on full index if search not triggered
-
-        Notes
-        -----
-        Computation is parallelized over all WMOs with a ThreadPoolExecutor
         """
+        if hasattr(self, "search") and not index:
+            table = self.search
+        else:
+            if not hasattr(self, "index"):
+                self.load(nrows=self._nrows_index)
+            table = self.index
+        files = table["file"]
+        extracted = pa.compute.extract_regex(files, pattern=r"/(?P<wmo>\d+)/")
+        wmos = pa.compute.struct_field(extracted, "wmo").drop_null()
+        vc = pa.compute.value_counts(wmos).to_pylist()
+        count = {int(d["values"]): d["counts"] for d in vc}
+        return count
 
-        def count_wmo(self, wmo: int, index: bool):
-            if hasattr(self, "search") and not index:
-                search_filter = pa.compute.match_substring_regex(
-                    self.search["file"], pattern="/%i/" % wmo
-                )
-                return wmo, self.search.filter(search_filter).shape[0]
-            else:
-                if not hasattr(self, "index"):
-                    self.load(nrows=self._nrows_index)
-                search_filter = pa.compute.match_substring_regex(
-                    self.index["file"], pattern="/%i/" % wmo
-                )
-                return wmo, self.index.filter(search_filter).shape[0]
+    def chunks_wmo(self, index=False, *args, **kwargs):
+        """Split WMOs into chunks balanced on both float count and profile count.
 
-        ulist = self.read_wmo()
-        count = {}
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            # Submit tasks for each WMO
-            futures = {
-                executor.submit(count_wmo, self, wmo, index): wmo for wmo in ulist
-            }
+        Fall back on full index if search not triggered
 
-            # Process results as they complete
-            for future in concurrent.futures.as_completed(futures):
-                wmo, cnt = future.result()
-                count[wmo] = cnt
+        Parameters
+        ----------
+        index: bool, default=False
+            Use the full index or search result
+        counts : dict
+            {wmo: number_of_profiles}
+        n_chunks : int, optional
+            Number of chunks to return.
+        max_profiles : int, optional
+            Target maximum number of profiles per chunk. The number of chunks is
+            derived as ceil(total_profiles / max_profiles).
+        max_floats : int, optional
+            Maximum number of floats (WMOs) per chunk. The number of chunks is
+            derived as ceil(n_floats / max_floats), and the cap is strictly enforced.
+        floats_weight : float in [0, 1]
+            Relative importance of balancing the number of floats (1.0) versus
+            the number of profiles (0.0) across chunks.
+
+        Exactly one of `n_chunks`, `max_profiles` or `max_floats` must be given.
 
-        return count
+        Returns
+        -------
+        list of list of int
+            Each inner list holds the (sorted) WMOs of one chunk.
+        """
+        counts = self.records_per_wmo(index=index)
+        return chunk_wmos(counts, *args, **kwargs)
 
     def to_indexfile(self, file):
         """Save search results on file, following the Argo standard index formats

argopy/stores/index/implementations/pyarrow/search_engine.py

@@ -29,6 +29,32 @@
 log = logging.getLogger("argopy.stores.index.pa")
 
 
+
+def filter_by_wmo(table, wmos):
+    """Return the rows of `table` whose "file" path belongs to one of `wmos`.
+
+    Parameters
+    ----------
+    table : pyarrow.Table
+        Index or search table with a "file" column (e.g. "dac/<wmo>/profiles/...").
+    wmos : iterable of int or str
+        WMOs to keep.
+
+    Returns
+    -------
+    pyarrow.Table
+    """
+    # Compare as strings to avoid casting the whole column to int
+    value_set = pa.array([str(int(w)) for w in wmos], type=pa.string())
+
+    extracted = pa.compute.extract_regex(table["file"], pattern=r"/(?P<wmo>\d+)/")
+    wmo_col = pa.compute.struct_field(extracted, "wmo")  # null where no match
+
+    mask = pa.compute.is_in(wmo_col, value_set=value_set)
+    # return table.filter(mask)
+    return mask
+
+
 @lru_cache(maxsize=25_000)
 def compute_wmo(wmo: int, obj):
     return pa.compute.match_substring_regex(obj.index["file"], pattern="/%i/" % wmo)
@@ -75,8 +101,7 @@ def namer(WMOs):
             return {"WMO": WMOs}
 
         def composer(obj, WMOs):
-            filt = pmap(obj, compute_wmo, WMOs)
-            return obj._reduce_a_filter_list(filt, op="or")
+            return filter_by_wmo(obj.index, WMOs)
 
         WMOs = checker(WMOs)
         self._obj.load(nrows=self._obj._nrows_index)

argopy/utils/arco/__init__.py

@@ -0,0 +1,14 @@
+"""
+This submodule is a naive approach to massively load Argo data as a large collection of points, or profiles interpolated on standard pressure levels.
+
+This is an experimental submodule based on previous work from the EA-RISE project.
+This submodule is adapted from the original OSnet codebase developped by Sean Tokunaga and Guillaume Maze, but never published before.
+https://github.com/euroargodev/argopy-osnet
+
+API may change anytime, not all files may be necessary in here and shall be removed/refactored without notice.
+"""
+from argopy.utils.arco.fetcher import MassFetcher
+from argopy.utils.arco.flat import FlatData
+
+
+__all__ = ('MassFetcher', 'FlatData')