normalize io.gcp.DicomSearch

sdks/python/apache_beam/yaml/standard_io.yaml

@@ -118,6 +118,7 @@
     'ReadFromMongoDB': 'apache_beam.yaml.yaml_io.read_from_mongodb'
     'WriteToMongoDB': 'apache_beam.yaml.yaml_io.write_to_mongodb'
     'ReadFromDelta': 'apache_beam.yaml.yaml_io.read_from_delta'
+    'DicomSearch': 'apache_beam.yaml.yaml_io.dicom_search'
 
 # General File Formats
 # Declared as a renaming transform to avoid exposing all

sdks/python/apache_beam/yaml/yaml_io.py

@@ -909,3 +909,109 @@ def match_all(
           path=str(x.path), size_in_bytes=int(x.size_in_bytes),
           last_updated_in_seconds=float(x.last_updated_in_seconds)
           if x.last_updated_in_seconds is not None else None))
+
+
+def _dicom_search_result_to_row(result):
+  if not result.get('success'):
+    raise RuntimeError(
+        'DicomSearch failed with status: %s' % (result.get('status'), ))
+  return beam.Row(
+      result=json.dumps(result.get('result', [])),
+      status=str(result.get('status')),
+      input=json.dumps(result.get('input', {})))
+
+
+def _dicom_search_to_output_row(result):
+  return beam.Row(
+      result=json.dumps(result.get('result', [])),
+      status=str(result.get('status')),
+      input=json.dumps(result.get('input', {})))
+
+
+def _dicom_search_to_error_row(result):
+  inp = result.get('input') or {}
+  if isinstance(inp, Mapping):
+    element = beam.Row(**dict(inp))
+  else:
+    element = inp
+  return beam.Row(
+      element=element,
+      msg='DicomSearch failed with status: %s' % (result.get('status'), ),
+      stack='')
+
+
+@beam.ptransform_fn
+def dicom_search(
+    pcoll, *, buffer_size: int = 8, max_workers: int = 5, error_handling=None):
+  """Searches a Google Cloud Healthcare DICOM store using QIDO-RS.
+
+  This transform takes an input PCollection of Rows describing QIDO search
+  requests and returns Rows with the search results encoded as JSON.
+
+  Each input Row must include:
+
+    - project_id (str): GCP project containing the DICOM store.
+    - region (str): Region where the DICOM store resides.
+    - dataset_id (str): Dataset containing the DICOM store.
+    - dicom_store_id (str): DICOM store id.
+    - search_type (str): One of ``studies``, ``series``, or ``instances``.
+    - params (map of str to str, optional): QIDO search filters.
+
+  Successful outputs are Rows with:
+
+    - result (str): JSON-encoded list of matching DICOM resources.
+    - status (str): HTTP status from the DICOM API.
+    - input (str): JSON-encoded copy of the search request.
+
+  Failed searches raise unless ``error_handling`` is set, in which case they
+  are routed to the configured error output.
+
+  Args:
+    buffer_size: Number of requests to buffer before flushing.
+    max_workers: Maximum number of threads used to issue requests.
+    error_handling: If specified, should be a mapping giving an output into
+      which to emit failed searches, as described at
+      https://beam.apache.org/documentation/sdks/yaml-errors/
+  """
+  try:
+    from apache_beam.io.gcp.healthcare.dicomio import DicomSearch
+  except ImportError as exn:
+    raise ValueError(
+        "GCP dependencies are not installed. Cannot use DicomSearch. "
+        "Please install using 'pip install apache-beam[gcp]'.") from exn
+
+  def row_to_dict(value):
+    if value is None:
+      return None
+    if hasattr(value, '_asdict'):
+      return {k: row_to_dict(v) for k, v in value._asdict().items()}
+    elif hasattr(value, 'as_dict'):
+      return {k: row_to_dict(v) for k, v in value.as_dict().items()}
+    elif isinstance(value, (list, tuple)):
+      return [row_to_dict(v) for v in value]
+    elif isinstance(value, Mapping):
+      return {k: row_to_dict(v) for k, v in value.items()}
+    else:
+      return value
+
+  if error_handling:
+    error_handling = yaml_utils.SafeLineLoader.strip_metadata(error_handling)
+
+  results = (
+      pcoll
+      | beam.Map(row_to_dict)
+      | DicomSearch(buffer_size=buffer_size, max_workers=max_workers))
+
+  if error_handling and error_handling.get('output'):
+    return {
+        'good': (
+            results
+            | beam.Filter(lambda r: r.get('success'))
+            | beam.Map(_dicom_search_to_output_row)),
+        error_handling['output']: (
+            results
+            | beam.Filter(lambda r: not r.get('success'))
+            | beam.Map(_dicom_search_to_error_row)),
+    }
+
+  return results | beam.Map(_dicom_search_result_to_row)

sdks/python/apache_beam/yaml/yaml_io_test.py

@@ -807,6 +807,153 @@ def test_query_and_table_both_raises(self):
           schema={'fields': []})
 
 
+class FakeDicomSearch(beam.PTransform):
+  def __init__(
+      self, buffer_size=8, max_workers=5, client=None, credential=None):
+    self.buffer_size = buffer_size
+    self.max_workers = max_workers
+
+  def expand(self, pcoll):
+    def do_search(element):
+      required = [
+          'project_id',
+          'region',
+          'dataset_id',
+          'dicom_store_id',
+          'search_type',
+      ]
+      for key in required:
+        if key not in element:
+          return {
+              'result': [],
+              'status': 'Must have %s in the dict.' % key,
+              'input': element,
+              'success': False,
+          }
+      if element['search_type'] not in ('instances', 'studies', 'series'):
+        return {
+            'result': [],
+            'status': (
+                'Search type can only be "studies", '
+                '"instances" or "series"'),
+            'input': element,
+            'success': False,
+        }
+      if element.get('project_id') == 'bad_project':
+        return {
+            'result': [],
+            'status': 500,
+            'input': element,
+            'success': False,
+        }
+      params = element.get('params') or {}
+      result = [{'PatientName': 'Alice', 'params': params}]
+      return {
+          'result': result,
+          'status': 200,
+          'input': element,
+          'success': True,
+      }
+
+    return pcoll | beam.Map(do_search)
+
+
+def _patch_dicom_search():
+  """Install a fake dicomio module so tests do not require GCP extras."""
+  import sys
+  import types
+
+  dicomio_mod = types.ModuleType('apache_beam.io.gcp.healthcare.dicomio')
+  dicomio_mod.DicomSearch = FakeDicomSearch
+  return mock.patch.dict(
+      sys.modules, {'apache_beam.io.gcp.healthcare.dicomio': dicomio_mod})
+
+
+class YamlDicomSearchTest(unittest.TestCase):
+  def test_dicom_search_success(self):
+    with _patch_dicom_search():
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        result = (
+            p
+            | beam.Create([
+                beam.Row(
+                    project_id='proj',
+                    region='us-central1',
+                    dataset_id='dataset',
+                    dicom_store_id='store',
+                    search_type='instances',
+                    params={'PatientName': 'Alice'})
+            ])
+            | YamlTransform(
+                '''
+                type: DicomSearch
+                '''))
+        assert_that(
+            result
+            | beam.Map(
+                lambda row:
+                (row.status, json.loads(row.result)[0]['PatientName'])),
+            equal_to([('200', 'Alice')]))
+
+  def test_dicom_search_with_error_handling(self):
+    with _patch_dicom_search():
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        result = (
+            p
+            | beam.Create([
+                beam.Row(
+                    project_id='proj',
+                    region='us-central1',
+                    dataset_id='dataset',
+                    dicom_store_id='store',
+                    search_type='instances'),
+                beam.Row(
+                    project_id='bad_project',
+                    region='us-central1',
+                    dataset_id='dataset',
+                    dicom_store_id='store',
+                    search_type='instances'),
+            ])
+            | YamlTransform(
+                '''
+                type: DicomSearch
+                config:
+                  error_handling:
+                    output: errors
+                '''))
+        assert_that(
+            result['good'] | beam.Map(lambda row: row.status),
+            equal_to(['200']),
+            label='CheckGood')
+        assert_that(
+            result['errors'] | beam.Map(lambda error: error.msg),
+            equal_to(['DicomSearch failed with status: 500']),
+            label='CheckErrors')
+
+  def test_dicom_search_without_error_handling_raises(self):
+    with self.assertRaises(Exception):
+      with _patch_dicom_search():
+        with beam.Pipeline(
+            options=beam.options.pipeline_options.PipelineOptions(
+                pickle_library='cloudpickle')) as p:
+          _ = (
+              p
+              | beam.Create([
+                  beam.Row(
+                      project_id='bad_project',
+                      region='us-central1',
+                      dataset_id='dataset',
+                      dicom_store_id='store',
+                      search_type='instances')
+              ])
+              | YamlTransform(
+                  '''
+                  type: DicomSearch
+                  '''))
+
+
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)
   unittest.main()