added Hdf5db.queryDataset method

Author: jreadey

src/h5json/h5reader.py

@@ -73,14 +73,26 @@ def getAttribute(self, obj_id, name, includeData=True):
         pass
 
     @abstractmethod
-    def getDatasetValues(self, obj_id, sel=None, dtype=None, query=None):
+    def getDatasetValues(self, obj_id, sel=None, dtype=None):
         """
         Get values from dataset identified by obj_id.
         If a slices list or tuple is provided, it should have the same
         number of elements as the rank of the dataset.
         """
         pass
 
+    def queryDataset(self, obj_id, query, sel=None):
+        """
+        Query the given dataset using the selection and query expression
+
+        Return a numpy array of indices for the elements that match the query.
+        Readers are not required to implement this — by default it raises
+        NotImplementedError, and Hdf5db falls back to querying the dataset values
+        it fetches via getDatasetValues. Override this only if the storage backend
+        has a more efficient way to evaluate the query (e.g. pushing it down to storage).
+        """
+        raise NotImplementedError("queryDataset not implemented for " + type(self).__name__)
+
     @abstractmethod
     def open(self):
         """ Open data source for reading """

src/h5json/hdf5db.py

@@ -31,20 +31,11 @@ def _query_rel_to_abs(x_sel, rel_indices, rank):
     """Map arrayQuery relative indices (within a sub-array) to absolute dataset indices.
 
     x_sel: SimpleSelection whose sub-array was queried
-    rel_indices: arrayQuery result — (N,) for rank=1, (N, rank) for rank>1
+    rel_indices: arrayQuery result, ndarray of shape (N, rank)
     """
     slices = x_sel.slices
-    if rank == 1:
-        s = slices[0]
-        if isinstance(s, slice):
-            start = s.start if s.start is not None else 0
-            step = s.step if s.step is not None else 1
-            return (rel_indices * step + start).astype(np.dtype('u8'))
-        if isinstance(s, list):
-            return np.array(s, dtype='u8')[rel_indices.astype(int)]
-        return np.full(len(rel_indices), int(s), dtype='u8')
     if len(rel_indices) == 0:
-        return rel_indices.astype('u8')
+        return np.zeros((0, rank), dtype='u8')
     abs_result = np.zeros((len(rel_indices), rank), dtype='u8')
     for d in range(rank):
         s = slices[d]
@@ -696,7 +687,7 @@ def deleteAttribute(self, obj_id, name):
 
         self.make_dirty(obj_id)
 
-    def getDatasetValues(self, dset_id, sel, query=None):
+    def getDatasetValues(self, dset_id, sel):
         """
         Get values from dataset identified by obj_id.
         If a slices list or tuple is provided, it should have the same
@@ -796,8 +787,6 @@ def init_arr(rdtype, cpl):
                 raise ValueError("Only SELECT_ALL selections are supported for scalar datasets")
             if sel.shape != ():
                 raise ValueError("Selection shape does not match dataset shape")
-            if query:
-                raise ValueError("Query is not supported for scalar datasets")
             if updates:
                 # for scalars the update has to be the requested value
                 (update_sel, arr) = updates[-1]
@@ -816,55 +805,6 @@ def init_arr(rdtype, cpl):
         arr = None
         fetch = True
 
-        if query:
-            full_shape = sel.shape
-            rank = len(full_shape)
-            result_mask = np.zeros(full_shape, dtype=bool)
-
-            # Delegate query to the reader when it has relevant data
-            query_fetch = not (isinstance(self._reader, H5NullReader) or dset_id in self._new_objects)
-            if query_fetch:
-                for (update_sel, _) in updates:
-                    if selections.contained(sel, update_sel):
-                        query_fetch = False
-                        break
-
-            if query_fetch:
-                reader_result = self.reader.getDatasetValues(dset_id, sel, dtype=dtype, query=query)
-                if reader_result is not None and len(reader_result) > 0:
-                    if rank == 1:
-                        result_mask[reader_result.astype(int)] = True
-                    else:
-                        result_mask[tuple(reader_result[:, d].astype(int) for d in range(rank))] = True
-
-            for (update_sel, update_val) in updates:
-                x_sel = selections.intersect(sel, update_sel)
-                if x_sel.nselect == 0:
-                    continue
-
-                # Invalidate reader results overwritten by this update
-                inter_mask = np.zeros(full_shape, dtype=bool)
-                inter_mask[x_sel.slices] = True
-                result_mask &= ~inter_mask
-
-                # Query the updated values at the intersection
-                local_sel = selections.translate(update_sel, x_sel)
-                x_vals = update_val[local_sel.slices]
-                x_rel = arrayQuery(query, x_vals)
-
-                if len(x_rel) > 0:
-                    abs_result = _query_rel_to_abs(x_sel, x_rel, rank)
-                    if rank == 1:
-                        result_mask[abs_result.astype(int)] = True
-                    else:
-                        result_mask[tuple(abs_result[:, d].astype(int) for d in range(rank))] = True
-
-            indices = np.argwhere(result_mask)
-            if rank == 1:
-                return indices.reshape(-1).astype(np.dtype('u8'))
-            else:
-                return indices.astype(np.dtype('u8'))
-
         # determine if we need to get data from the reader
         if isinstance(self._reader, H5NullReader) or dset_id in self._new_objects:
             fetch = False
@@ -981,6 +921,101 @@ def init_arr(rdtype, cpl):
 
         return arr
 
+    def queryDataset(self, dset_id, query, sel=None, limit=0):
+        """
+        Query the given dataset using the selection and query expression
+        If sel is provided, only the elements in the selection will be queried,
+        otherwise the entire dataset will be queried.
+        If limit is provided, only the first limit number of elements that match the query will be returned.
+
+        Return a numpy array of indices for the elements that match the query
+        """
+
+        def queryReader(dset_id, query, sel=None, limit=0):
+            result = None
+            try:
+                result = self.reader.queryDataset(dset_id, query, sel=sel, limit=limit)
+            except NotImplementedError:
+                print("This reader doesn't support queryDataset")
+
+            if result is not None:
+                return result
+
+            # query the dataset by fetching the data and applying the query locally
+            arr = self.getDatasetValues(dset_id, sel)
+
+            result = arrayQuery(query, arr, limit=limit)
+            result = _query_rel_to_abs(sel, result, len(sel.shape))
+
+            return result
+        #
+        # start of queryDataset
+        #
+        if not isinstance(query, str):
+            raise TypeError("Expected query string")
+
+        if sel is not None and not isinstance(sel, selections.Selection):
+            raise TypeError("Expected Selection class")
+        if not isinstance(limit, int) or limit < 0:
+            raise TypeError("Expected non-negative integer for limit")
+
+        dset_json = self.getObjectById(dset_id)
+        shape_json = dset_json["shape"]
+
+        shape_class = getShapeClass(shape_json)
+        if shape_class == "H5S_NULL":
+            raise ValueError("querying null space dataset not supported")
+        dims = getShapeDims(shape_json)
+        if sel is None:
+            sel = selections.select(dims, ...)
+
+        updates = self._getDatasetUpdates(dset_id)
+
+        if sel.shape != dims:
+            raise TypeError("Selection shape does not match dataset shape")
+
+        full_shape = sel.shape
+        rank = len(full_shape)
+
+        # Delegate query to the reader when it has relevant data
+        query_fetch = not (isinstance(self._reader, H5NullReader) or dset_id in self._new_objects)
+        if query_fetch:
+            for (update_sel, _) in updates:
+                if selections.contained(sel, update_sel):
+                    query_fetch = False
+                    break
+
+        result_mask = np.zeros(full_shape, dtype=bool)
+        if query_fetch:
+            fetched = queryReader(dset_id, query, sel=sel, limit=limit)
+            if len(fetched) > 0:
+                result_mask[tuple(fetched[:, d].astype(int) for d in range(rank))] = True
+
+        for (update_sel, update_val) in updates:
+            x_sel = selections.intersect(sel, update_sel)
+            if x_sel.nselect == 0:
+                continue
+
+            # Invalidate reader results overwritten by this update
+            inter_mask = np.zeros(full_shape, dtype=bool)
+            inter_mask[x_sel.slices] = True
+            result_mask &= ~inter_mask
+
+            # Query the updated values at the intersection
+            local_sel = selections.translate(update_sel, x_sel)
+            x_vals = update_val[local_sel.slices]
+            x_rel = arrayQuery(query, x_vals)
+
+            if len(x_rel) > 0:
+                abs_result = _query_rel_to_abs(x_sel, x_rel, rank)
+                result_mask[tuple(abs_result[:, d].astype(int) for d in range(rank))] = True
+
+        indices = np.argwhere(result_mask)
+        if limit > 0 and len(indices) > limit:
+            indices = indices[:limit]
+
+        return indices
+
     def setDatasetValues(self, dset_id, sel, arr):
         """
         Write the given ndarray to the dataset using the selection

src/h5json/query_util.py

@@ -377,11 +377,7 @@ def arrayQuery(
 
     if limit is not 0, only up to limit indices will be returned
 
-    The return value will be an ndarray.  If data_arr is one dimensionsal, the return value will be
-    an array of shape (count,) where count is the number of matches of the query.  Each element will be
-    an index to the matching element.  If data_arr is multi-dimensional, the array will have shape (count, rank),
-    where count is the number of matching elements and rank is the rank of the data_arr.  In this case, element [i,j]
-    of the returned array will be the nth coordinate of the ith match of the query.
+    The return value will be an ndarray.  The array shape (count, rank) where rank is the number of array dimensions.
 
     Example queries:
         "_ > 1.0" # match any array element with a value greater than 1.0
@@ -417,8 +413,4 @@ def arrayQuery(
     indices = np.argwhere(mask)
     if limit > 0:
         indices = indices[:limit]
-
-    if rank == 1:
-        return indices.reshape(-1).astype(np.dtype('u8'))
-    else:
-        return indices.astype(np.dtype('u8'))
+    return indices

src/h5json/selections.py

@@ -197,13 +197,13 @@ def from_query_result(shape, indices):
     """Create a PointSelection from an arrayQuery result.
 
     shape: full dataset shape tuple
-    indices: ndarray of shape (N,) for 1D or (N, rank) for nD, as returned by arrayQuery
+    indices: ndarray of shape (N, rank), as returned by arrayQuery
     """
     rank = len(shape)
     if len(indices) == 0:
         return _empty_paired_sel(shape)
     if rank == 1:
-        return select(shape, indices.astype(int).tolist())
+        return select(shape, indices[:, 0].astype(int).tolist())
     coords = tuple(indices[:, d].astype(int).tolist() for d in range(rank))
     return select(shape, coords)
 

test/unit/array_query_test.py

@@ -69,7 +69,7 @@ def testArrayQueryNoneCompound(self):
         query = "_ > 1.0"
         result = arrayQuery(query, data_arr)
         self.assertTrue(isinstance(result, np.ndarray))
-        self.assertEqual(result.dtype, np.dtype("u8"))
+        self.assertEqual(result.dtype, np.dtype("int64"))
         self.assertEqual(len(result.shape), 2)
         self.assertEqual(result.shape[0], len(expected))
         self.assertEqual(result.shape[1], 2)
@@ -83,18 +83,18 @@ def testArrayQuery1D(self):
         query = "symbol == b'AAPL'"
         result = arrayQuery(query, data_arr)
         self.assertTrue(isinstance(result, np.ndarray))
-        self.assertEqual(result.dtype, np.dtype("u8"))
-        self.assertEqual(result.shape, (4,))
+        self.assertEqual(result.dtype, np.dtype("int64"))
+        self.assertEqual(result.shape, (4, 1))
         expected_indexes = (1, 4, 7, 10)  # rows above with AAPL as symbol
         for i in range(4):
             item = result[i]
-            self.assertTrue(item in expected_indexes)
+            self.assertTrue(item[0] in expected_indexes)
 
         # read just one row back
         result = arrayQuery(query, data_arr, limit=1)
         self.assertTrue(isinstance(result, np.ndarray))
         self.assertEqual(len(result), 1)
-        index = result[0]
+        index = result[0][0]
         self.assertEqual(index, 1)
 
         # query with selection, no limit