restore PUT_Value with query param, modified rsp to return indices

Author: jreadey

hsds/chunk_crawl.py

@@ -26,7 +26,7 @@
 from h5json.hdf5dtype import createDataType
 from h5json.array_util import jsonToArray, getNumpyValue
 from h5json.array_util import getNumElements, arrayToBytes, bytesToArray
-from h5json.shape_util import getShapeDims
+from h5json.shape_util import getShapeDims, getRank
 from h5json.dset_util import getChunkDims
 from h5json.time_util import getNow
 
@@ -220,6 +220,7 @@ async def read_chunk_hyperslab(
         raise HTTPInternalServerError()
     type_json = dset_json["type"]
     dset_dt = createDataType(type_json)
+    dset_rank = getRank(dset_json)
 
     chunk_shape = None  # expected return array shape
     chunk_sel = None  # for hyperslab
@@ -255,6 +256,10 @@ async def read_chunk_hyperslab(
 
     if query is None and query_update is None:
         query_dtype = None
+    elif query_update is not None:
+        # PUT_Chunk's query-update handling returns the global dataset
+        # indices of matching elements, as (n, rank) coordinate tuples
+        query_dtype = np.dtype("i8")
     else:
         # GET_Chunk's query handling (h5json.query_util.arrayQuery) returns
         # the matching values themselves, typed as select_dtype
@@ -381,9 +386,15 @@ async def read_chunk_hyperslab(
         log.debug(f"data: {len(array_data)} bytes")
         if query is not None or query_update is not None:
             # TBD: this needs to be fixed up for variable length dtypes
-            nrows = len(array_data) // query_dtype.itemsize
+            if query_update is not None:
+                # indices are returned as (n, rank)
+                nrows = len(array_data) // (query_dtype.itemsize * dset_rank)
+                rsp_shape = (nrows, dset_rank)
+            else:
+                nrows = len(array_data) // query_dtype.itemsize
+                rsp_shape = (nrows,)
             try:
-                chunk_arr = bytesToArray(array_data, query_dtype, (nrows,))
+                chunk_arr = bytesToArray(array_data, query_dtype, rsp_shape)
             except ValueError as ve:
                 log.warn(f"bytesToArray ValueError: {ve}")
                 raise HTTPBadRequest()

hsds/chunk_dn.py

@@ -14,10 +14,11 @@
 # handles regauests to read/write chunk data
 #
 
+import json
 import numpy as np
 import traceback
 from aiohttp.web_exceptions import HTTPBadRequest, HTTPInternalServerError
-from aiohttp.web_exceptions import HTTPNotFound, HTTPServiceUnavailable, HTTPNotImplemented
+from aiohttp.web_exceptions import HTTPNotFound, HTTPServiceUnavailable
 from aiohttp.web import json_response, StreamResponse
 
 from h5json.hdf5dtype import createDataType, getSubType
@@ -53,10 +54,20 @@ async def PUT_Chunk(request):
     bucket = None
     input_arr = None
     element_count = None
+    limit = 0
 
     if "query" in params:
         query = params["query"]
         log.info(f"PUT_Chunk query: {query}")
+    if "Limit" in params:
+        param_limit = params["Limit"]
+        try:
+            limit = int(param_limit)
+        except ValueError:
+            msg = f"invalid Limit param: {param_limit}"
+            log.warn(msg)
+            raise HTTPBadRequest(reason=msg)
+        log.debug(f"PUT_Chunk limit: {limit}")
     chunk_id = request.match_info.get("id")
     if not chunk_id:
         msg = "Missing chunk id"
@@ -184,13 +195,74 @@ async def PUT_Chunk(request):
             raise HTTPNotFound()
 
     if query:
-        # TBD: query+update support was dropped when chunkUtil.chunkQuery was
-        # removed in favor of h5json.query_util.arrayQuery. arrayQuery returns
-        # match coordinates rather than rows, so this needs to be rewritten to
-        # update chunk_arr at those coordinates and rebuild a response array.
-        msg = "PUT_Chunk with query is not currently supported"
-        log.error(msg)
-        raise HTTPNotImplemented(reason=msg)
+        try:
+            indices = arrayQuery(query, chunk_arr, selection=selection, limit=limit)
+        except (TypeError, ValueError) as e:
+            msg = f"query: {query} is not valid, got exception: {e}"
+            log.warn(msg)
+            raise HTTPBadRequest(reason=msg)
+
+        log.debug(f"PUT_Chunk - query matched {len(indices)} elements")
+
+        try:
+            update_value = await request.json()
+        except json.JSONDecodeError:
+            msg = "Unable to load JSON body for query update"
+            log.warn(msg)
+            raise HTTPBadRequest(reason=msg)
+
+        rank = len(chunk_arr.shape)
+        fancy_index = tuple(indices[:, i] for i in range(rank))
+
+        if len(indices) > 0:
+            # query_update is only allowed when the value is one element -
+            # that element gets broadcast across all matching positions
+            if select_dt.names:
+                # compound type - value is a JSON object of field name to
+                # value; only the given fields are updated, others are
+                # left as-is
+                if not isinstance(update_value, dict):
+                    msg = "expected a JSON object for compound type query update"
+                    log.warn(msg)
+                    raise HTTPBadRequest(reason=msg)
+                for field_name, field_value in update_value.items():
+                    if field_name not in select_dt.names:
+                        msg = f"field: {field_name} not found in dataset type"
+                        log.warn(msg)
+                        raise HTTPBadRequest(reason=msg)
+                    chunk_arr[field_name][fancy_index] = field_value
+            else:
+                # simple type - value is the (scalar) element itself
+                if isinstance(update_value, dict) and "value" in update_value:
+                    update_value = update_value["value"]
+                chunk_arr[fancy_index] = update_value
+            is_dirty = True
+
+        # return the global dataset indices of the matching elements -
+        # the chunk's offset within the dataset (per dimension) is its
+        # grid index times the chunk dims along that dimension
+        chunk_index = getChunkIndex(chunk_id)
+        offset = np.array([chunk_index[i] * dims[i] for i in range(rank)], dtype=indices.dtype)
+        global_indices = indices + offset
+
+        read_resp = arrayToBytes(global_indices)
+        try:
+            resp = StreamResponse()
+            resp.headers["Content-Type"] = "application/octet-stream"
+            resp.content_length = len(read_resp)
+            await resp.prepare(request)
+            await resp.write(read_resp)
+        except Exception as e:
+            log.error(f"Exception during binary data write: {e}")
+            raise HTTPInternalServerError()
+        finally:
+            await resp.write_eof()
+
+        if is_dirty or config.get("write_zero_chunks", default=False):
+            save_chunk(app, chunk_id, dset_json, chunk_arr, bucket=bucket)
+
+        log.response(request, resp=resp)
+        return resp
     else:
         # regular chunk update
         # check that the content_length is what we expect

hsds/chunk_sn.py

@@ -277,7 +277,7 @@ def _getPoints(body, rank=1):
     return points
 
 
-def _getQuery(params, dtype, rank=1, body=None):
+def _getQuery(params, dtype, body=None):
     """ get query parameter and validate if set """
 
     kw = "query"
@@ -293,11 +293,6 @@ def _getQuery(params, dtype, rank=1, body=None):
             msg = "Query string can not be used with append parameter"
             log.warn(msg)
             raise HTTPBadRequest(reason=msg)
-        # validate the query string
-        if rank > 1:
-            msg = "Query string is not supported for multidimensional datasets"
-            log.warn(msg)
-            raise HTTPBadRequest(reason=msg)
 
         # following will throw HTTPBadRequest if query is malformed
         validateQuery(query, dtype)
@@ -383,12 +378,13 @@ async def _getRequestData(request, http_streaming=True):
 
 
 async def arrayResponse(arr, request, dset_json):
-    """ return the array as binary or json response based on accept type """
+    """ return the query-update match indices as a binary or json response
+    based on accept type """
     response_type = getAcceptType(request)
 
     if response_type == "binary":
         output_data = arr.tobytes()
-        msg = f"PUT_Value query - returning {len(output_data)} bytes binary data"
+        msg = f"PUT_Value query - returning {len(output_data)} bytes binary indices"
         log.debug(msg)
 
         # write response
@@ -405,16 +401,16 @@ async def arrayResponse(arr, request, dset_json):
         except Exception as e:
             log.error(f"Exception during binary data write: {e}")
     else:
-        log.debug("PUT Value query - returning JSON data")
+        log.debug("PUT Value query - returning JSON indices")
         rsp_json = {}
         data = arr.tolist()
-        log.debug(f"got rsp data {len(data)} points")
+        log.debug(f"got rsp data {len(data)} indices")
         try:
-            json_query_data = bytesArrayToList(data)
+            indices_data = bytesArrayToList(data)
         except ValueError as err:
             msg = f"Cannot decode provided bytes to list: {err}"
             raise HTTPBadRequest(reason=msg)
-        rsp_json["value"] = json_query_data
+        rsp_json["indices"] = indices_data
         rsp_json["hrefs"] = get_hrefs(request, dset_json)
 
         resp = await jsonResponse(request, rsp_json)
@@ -518,7 +514,7 @@ async def PUT_Value(request):
     # if there's no selection parameter, this will return entire dataspace
     selection = _getSelect(params, dset_json, body=body)
 
-    query = _getQuery(params, dset_dtype, rank=rank, body=body)
+    query = _getQuery(params, dset_dtype, body=body)
 
     element_count = _getElementCount(params, body=body)
 
@@ -808,7 +804,7 @@ async def GET_Value(request):
         ignore_nan = False
     log.debug(f"ignore nan: {ignore_nan}")
 
-    query = _getQuery(params, dset_dtype, rank=rank)
+    query = _getQuery(params, dset_dtype)
 
     response_type = getAcceptType(request)
 

hsds/dset_lib.py

@@ -520,6 +520,12 @@ async def doReadSelection(
         select_dtype = dset_dtype
     if query is None:
         query_dtype = None
+    elif query_update is not None:
+        # PUT_Chunk's query-update handling returns the global dataset
+        # indices of matching elements, as (n, rank) coordinate tuples
+        log.debug(f"query: {query} limit: {limit} query_update: {query_update}")
+        query_dtype = np.dtype("i8")
+        query_rank = getRank(dset_json)
     else:
         # GET_Chunk's query handling (h5json.query_util.arrayQuery) returns
         # the matching values themselves, typed as select_dtype
@@ -600,7 +606,10 @@ async def doReadSelection(
             nrows = limit
         else:
             nrows = crawler._hits
-        arr = np.empty((nrows,), dtype=query_dtype)
+        if query_update is not None:
+            arr = np.empty((nrows, query_rank), dtype=query_dtype)
+        else:
+            arr = np.empty((nrows,), dtype=query_dtype)
         start = 0
         for chunkid in chunk_ids:
             if chunkid not in chunk_map: