add update_value option to queryDataset

Author: jreadey

src/h5json/h5writer.py

@@ -70,6 +70,18 @@ def append(self):
     def no_data(self):
         return self._no_data
 
+    def queryDataset(self, obj_id, query, sel=None, limit=0, update_value=None):
+        """
+        Replace any elements that match the
+
+        Return a numpy array of indices for the elements that match the query.
+        Writers 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 storage handle, return root_id"""

src/h5json/hdf5db.py

@@ -35,8 +35,8 @@ def _query_rel_to_abs(x_sel, rel_indices, rank):
     """
     slices = x_sel.slices
     if len(rel_indices) == 0:
-        return np.zeros((0, rank), dtype='u8')
-    abs_result = np.zeros((len(rel_indices), rank), dtype='u8')
+        return np.zeros((0, rank), dtype='int64')
+    abs_result = np.zeros((len(rel_indices), rank), dtype='int64')
     for d in range(rank):
         s = slices[d]
         if isinstance(s, slice):
@@ -1110,13 +1110,16 @@ def getChunkIterator(self, dset_id, sel=None):
         """
         return ChunkIterator(self, dset_id, sel=sel)
 
-    def queryDataset(self, dset_id, query, sel=None, limit=0):
+    def queryDataset(self, dset_id, query, sel=None, limit=0, update_value=None):
         """
         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.
-
+        If update_value is provied, elements matching the query (up to limit elements if limit is non-zero)
+        will be updated to the given value. For a compound dtype, update_value may be a dict mapping
+        field names to the value to set for that field - only those fields are modified, and the rest
+        of each matching element is left unchanged.
         Return a numpy array of indices for the elements that match the query
         """
 
@@ -1155,10 +1158,75 @@ def queryReader(dset_id, query, sel=None, limit=0):
                 if limit > 0 and nhits >= limit:
                     break
 
-            result = np.concatenate(hits, axis=0) if hits else np.zeros((0, rank), dtype='u8')
+            result = np.concatenate(hits, axis=0) if hits else np.zeros((0, rank), dtype='int64')
             if limit > 0 and len(result) > limit:
                 result = result[:limit]
             return result
+
+        def queryWriter(dset_id, query, sel=None, limit=0, update_value=None):
+            result = None
+            try:
+                result = self.writer.queryDataset(dset_id, query, sel=sel, limit=limit, update_value=update_value)
+            except NotImplementedError:
+                # This reader doesn't support queryDataset
+                pass
+
+            if result is None:
+                rank = len(sel.shape)
+                try:
+                    chunk_iter = ChunkIterator(self, dset_id, sel=sel)
+                except ValueError:
+                    # ChunkIterator doesn't support this selection (e.g. a fancy/point
+                    # selection, or a scalar dataset) - fall back to querying the
+                    # entire selection at once
+                    arr = self.getDatasetValues(dset_id, sel)
+                    result = arrayQuery(query, arr, limit=limit)
+                    result = _query_rel_to_abs(sel, result, rank)
+                else:
+                    # query the dataset chunk by chunk so the whole selection is
+                    # never loaded into memory at once
+                    hits = []
+                    nhits = 0
+                    for chunk_arr in chunk_iter:
+                        chunk_rel = arrayQuery(query, chunk_arr)
+                        if len(chunk_rel) == 0:
+                            continue
+                        hits.append(_query_rel_to_abs(chunk_iter.sel, chunk_rel, rank))
+                        nhits += len(chunk_rel)
+                        if limit > 0 and nhits >= limit:
+                            break
+
+                    result = np.concatenate(hits, axis=0) if hits else np.zeros((0, rank), dtype='int64')
+                    if limit > 0 and len(result) > limit:
+                        result = result[:limit]
+
+            if update_value is not None and len(result) > 0:
+                # update the values at the matching indices
+                dtype = self.getDtype(self.getObjectById(dset_id))
+                if isinstance(update_value, dict):
+                    # a dict maps field names to the value to set for that field -
+                    # only those fields are modified, the rest of each record is
+                    # left as-is
+                    if len(dtype) == 0:
+                        raise TypeError("update_value dict is only supported for compound dtypes")
+                    fields = [f for f in dtype.names if f in update_value]
+                    if not fields:
+                        raise ValueError(
+                            f"None of the requested fields {list(update_value.keys())} found in dtype")
+                    if len(fields) == 1:
+                        value = np.asarray(update_value[fields[0]], dtype=dtype.fields[fields[0]][0])
+                    else:
+                        value_dtype = np.dtype([(f, dtype.fields[f][0]) for f in fields])
+                        value = np.zeros((), dtype=value_dtype)
+                        for f in fields:
+                            value[f] = update_value[f]
+                    update_sel = selections.select(sel.shape, result, fields=fields)
+                else:
+                    value = np.asarray(update_value, dtype=dtype)
+                    update_sel = selections.select(sel.shape, result)
+                self.setDatasetValues(dset_id, update_sel, value)
+            return result
+
         #
         # start of queryDataset
         #
@@ -1180,6 +1248,12 @@ def queryReader(dset_id, query, sel=None, limit=0):
         if sel is None:
             sel = selections.select(dims, ...)
 
+        if update_value is not None:
+            # do flush so we can be sure to do an atomic operation if the writer supports it
+            self.flush()
+            results = queryWriter(dset_id, query, sel=sel, limit=limit, update_value=update_value)
+            return results
+
         updates = self._getDatasetUpdates(dset_id)
 
         if sel.shape != dims:
@@ -1213,13 +1287,25 @@ def queryReader(dset_id, query, sel=None, limit=0):
             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
+            if update_sel.select_type == selections.H5S_SEL_POINTS:
+                # update_val is 1-D, indexed by position in update_sel (not
+                # sliceable via translate(), which only handles hyperslabs)
+                x_points = list(selections._iter_points(x_sel))
+                if not x_points:
+                    continue
+                upd_pt_to_idx = {pt: j for j, pt in enumerate(selections._iter_points(update_sel))}
+                x_vals = update_val[[upd_pt_to_idx[pt] for pt in x_points]]
+                x_rel = arrayQuery(query, x_vals)
+                if len(x_rel) > 0:
+                    abs_coords = np.array([x_points[i] for i in x_rel[:, 0]], dtype='u8')
+                    result_mask[tuple(abs_coords[:, d] for d in range(rank))] = True
+            else:
+                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:
@@ -1276,7 +1362,11 @@ def setDatasetValues(self, dset_id, sel, arr):
             raise TypeError(f"arr.dtype {src_dt} doesn't match expected dtype {expected_dt}")
 
         if sel.select_type == selections.H5S_SEL_POINTS:
-            if sel.nselect != arr.shape[0]:
+            if arr.shape == ():
+                # broadcast the scalar to match the number of selected points, so
+                # the stored update value can be indexed like any other point update
+                arr = np.full(sel.mshape, arr[()], dtype=arr.dtype)
+            elif sel.nselect != arr.shape[0]:
                 raise TypeError("Selection shape does not match number of points")
         elif sel.select_type == selections.H5S_SEL_FANCY:
             if arr.shape != sel.mshape:
@@ -1289,7 +1379,7 @@ def setDatasetValues(self, dset_id, sel, arr):
                 raise TypeError("Selection shape does not match dataset shape")
             # Allow scalar arrays when writing a field-restricted selection
             # (the scalar will be broadcast to all selected positions).
-            if arr.shape != () or sel.fields is None:
+            if arr.shape != ():
                 if 0 < arr.ndim < len(dims):
                     # arr has fewer dims than the dataset rank (e.g. a 1-D array
                     # written to a slice of a 3-D dataset).  Validate against

test/unit/hdf5db_test.py

@@ -709,6 +709,19 @@ def test2DDataset(self):
                 else:
                     self.assertEqual(x, i * 10 + j)
 
+        # point selection write with broadcasting
+        arr = np.array(42, dtype=dtype)
+        db.setDatasetValues(dset_id, sel, arr)
+        arr = db.getDatasetValues(dset_id, sel_all)
+        for i in range(nrows):
+            for j in range(ncols):
+                x = arr[i, j]
+                if i == j and i < 4:
+                    # these are the elements were set to 42 with the point write
+                    self.assertEqual(x, 42)
+                else:
+                    self.assertEqual(x, i * 10 + j)
+
         # test select all write
         arr = np.zeros(shape, dtype=dtype)
         arr[...] = 42
@@ -1073,6 +1086,18 @@ def testQuerySimpleType(self):
         query = "_ > 10"
         indices = db.queryDataset(dset_id, query)
         self.assertEqual(indices.shape, (56, 2))
+
+        indices = db.queryDataset(dset_id, query, update_value=0)
+        self.assertEqual(indices.shape, (56, 2))
+
+        indices = db.queryDataset(dset_id, query)
+        self.assertEqual(indices.shape, (0, 2))
+
+        # query update with limit
+        query = "_ == 0"
+        indices = db.queryDataset(dset_id, query, limit=5, update_value=-99)
+        self.assertEqual(indices.shape, (5, 2))
+
         db.close()
 
     def testQueryDataset1D(self):
@@ -1138,6 +1163,16 @@ def testQueryDataset1D(self):
         except ValueError:
             pass
 
+        # query with update_value
+        indices = db.queryDataset(dset_id, query, update_value={"open": -999, "close": 999})
+        self.assertEqual(indices.shape, (4, 1))
+        sel = selections.select(shape, indices)
+        values = db.getDatasetValues(dset_id, sel)
+        for i in range(len(values)):
+            self.assertEqual(values[i]["open"], -999)
+            self.assertEqual(values[i]["close"], 999)
+            self.assertEqual(values[i]["symbol"], b"AAPL")
+
         db.close()
 
     def testQueryDataset2D(self):
@@ -1171,6 +1206,16 @@ def testQueryDataset2D(self):
         for i, row in enumerate(indices):
             self.assertEqual(tuple(int(x) for x in row), expected_indexes[i])
 
+        # query with update_value
+        indices = db.queryDataset(dset_id, query, update_value={"open": -999, "close": 999})
+        self.assertEqual(indices.shape, (4, 2))
+        sel = selections.select(shape, indices)
+        values = db.getDatasetValues(dset_id, sel)
+        for i in range(len(values)):
+            self.assertEqual(values[i]["open"], -999)
+            self.assertEqual(values[i]["close"], 999)
+            self.assertEqual(values[i]["symbol"], b"AAPL")
+
         db.close()
 
     def testChunkIterator1D(self):