Enabled field filtering for Proxy sources

Author: Luke Shaw

caterva2/client.py

@@ -313,13 +313,13 @@ def get_download_url(self):
         """
         return api_utils.get_download_url(self.path, self.urlbase)
 
-    def __getitem__(self, key):
+    def __getitem__(self, item):
         """
         Retrieves a slice of the dataset.
 
         Parameters
         ----------
-        key : int, slice, tuple of ints and slices, or None
+        item : int, slice, tuple of ints and slices, or None
             Specifies the slice to fetch.
 
         Returns
@@ -340,10 +340,17 @@ def __getitem__(self, key):
         >>> ds[0:10]
         array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
         """
-        if isinstance(key, str):  # used a filter or field to index so want blosc2 array as result
-            return self.client.get_slice(self.path, key, as_blosc2=True)
+        if isinstance(item, str):  # used a filter or field to index so want blosc2 array as result
+            fields = np.dtype(eval(self.dtype)).fields
+            if fields is None:
+                raise ValueError("The array is not structured (its dtype does not have fields)")
+            if item in fields:
+                # A shortcut to access fields
+                return self.client.get_slice(self.path, as_blosc2=True, field=item)  # arg key is None
+            else:  # used a filter (possibly lazyexpr)
+                return self.client.get_slice(self.path, item, as_blosc2=True)
         else:
-            return self.slice(key, as_blosc2=False)
+            return self.slice(item, as_blosc2=False)
 
     def slice(
         self, key: int | slice | Sequence[slice], as_blosc2: bool = True
@@ -874,7 +881,7 @@ def fetch(self, path, slice_=None):
         # Does the same as get_slice but forces return of np array
         return self.get_slice(path, key=slice_, as_blosc2=False)
 
-    def get_slice(self, path, key=None, as_blosc2=True):
+    def get_slice(self, path, key=None, as_blosc2=True, field=None):
         """Get a slice of a File/Dataset.
 
         Parameters
@@ -888,6 +895,8 @@ def get_slice(self, path, key=None, as_blosc2=True):
             If True (default), the result will be returned as a Blosc2 object
             (either a `SChunk` or `NDArray`).  If False, it will be returned
             as a NumPy array (equivalent to `self[key]`).
+        field: str
+            Shortcut to access a field in a structured array. If provided, `key` is ignored.
 
         Returns
         -------
@@ -904,7 +913,11 @@ def get_slice(self, path, key=None, as_blosc2=True):
               dtype=[('a', '<f4'), ('b', '<f8')])
         """
         urlbase, path = _format_paths(self.urlbase, path)
-        if isinstance(key, str):  # A filter or field has been passed
+        if field:  # blosc2 doesn't support indexing of multiple fields
+            return api_utils.fetch_data(
+                path, urlbase, {"field": field}, auth_cookie=self.cookie, as_blosc2=as_blosc2
+            )
+        if isinstance(key, str):  # A filter has been passed
             return api_utils.fetch_data(
                 path, urlbase, {"filter": key}, auth_cookie=self.cookie, as_blosc2=as_blosc2
             )

caterva2/services/sub.py

@@ -24,6 +24,7 @@
 import tarfile
 import typing
 import zipfile
+from argparse import ArgumentError
 
 # Requirements
 import blosc2
@@ -747,6 +748,7 @@ async def fetch_data(
     slice_: str | None = None,
     user: db.User = Depends(optional_user),
     filter: str | None = None,
+    field: str | None = None,
 ):
     """
     Fetch a dataset.
@@ -759,6 +761,8 @@ async def fetch_data(
         The slice to fetch.
     filter : str
         The filter to apply to the dataset.
+    field : str
+        The desired field of dataset. If provided, filter is ignored.
 
     Returns
     -------
@@ -774,14 +778,18 @@ async def fetch_data(
     abspath, dataprep = abspath_and_dataprep(path, slice_, user=user)
     # This is still needed and will only update the necessary chunks
     await dataprep()
+
     if filter:
+        if field:
+            raise ArgumentError("Cannot handle both field and filter parameters at the same time")
         filter = filter.strip()
         container, _ = get_filtered_array(abspath, path, filter, sortby=None)
     else:
         container = open_b2(abspath, path)
 
     if isinstance(container, blosc2.Proxy):
         container = container._cache
+    container = container[field] if field else container
 
     if isinstance(container, blosc2.NDArray | blosc2.LazyExpr | hdf5.HDF5Proxy | blosc2.NDField):
         array = container
@@ -808,7 +816,7 @@ async def fetch_data(
             for sl, sh in zip(slice_, shape, strict=False)
         )
 
-    if whole and not isinstance(array, blosc2.LazyExpr | hdf5.HDF5Proxy | blosc2.NDField):
+    if whole and (not isinstance(array, blosc2.LazyExpr | hdf5.HDF5Proxy | blosc2.NDField)) and (not filter):
         # Send the data in the file straight to the client,
         # avoiding slicing and re-compression.
         return FileResponse(abspath, filename=abspath.name, media_type="application/octet-stream")
@@ -1869,13 +1877,15 @@ def get_filtered_array(abspath, path, filter, sortby):
 
     # Filter rows only for NDArray with fields
     if filter:
+        arr = arr._cache if isinstance(arr, blosc2.Proxy) else arr
         # Check whether filter is the name of a field
         if filter in arr.fields:
             if arr.dtype.fields[filter][0] == bool:  # noqa: E721
                 # If boolean, give the filter a boolean expression
                 filter = f"{filter} == True"
             else:
-                return arr[filter], idx  # just return the blosc2 NDfield associated to field
+                raise IndexError("Filter should be a boolean expression")
+
         # Let's create a LazyExpr with the filter
         larr = arr[filter]
         # TODO: do some benchmarking to see if this is worth it

caterva2/tests/test_api.py

@@ -698,7 +698,6 @@ def test_lazyexpr_fields(auth_client):
     if not auth_client:
         pytest.skip("authentication support needed")
 
-    opnm = "ds"
     oppt = f"{TEST_CATERVA2_ROOT}/ds-1d-fields.b2nd"
     auth_client.subscribe(TEST_CATERVA2_ROOT)
 
@@ -708,9 +707,9 @@ def test_lazyexpr_fields(auth_client):
     np.testing.assert_allclose(field[:], arr[:]["a"])
 
     # Test a lazyexpr
-    servered = arr["a < 500 & b >= .1"]
-    downloaded = arr.slice(None)["a < 500 & b >= .1"]
-    np.testing.assert_allclose(servered[:], downloaded[:])
+    servered = arr["(a < 500) & (b >= .1)"][:]
+    downloaded = arr.slice(None)["(a < 500) & (b >= .1)"][:]
+    [np.testing.assert_array_equal(servered[f], downloaded[f]) for f in downloaded.dtype.fields]
 
 
 def test_expr_from_expr(auth_client):