diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index 796429b5bdfc..58080cff4051 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/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 diff --git a/sdks/python/apache_beam/yaml/yaml_io.py b/sdks/python/apache_beam/yaml/yaml_io.py index b3ef18f96086..3b2a13ebf7f7 100644 --- a/sdks/python/apache_beam/yaml/yaml_io.py +++ b/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) diff --git a/sdks/python/apache_beam/yaml/yaml_io_test.py b/sdks/python/apache_beam/yaml/yaml_io_test.py index c3df0328f22b..43955233d10a 100644 --- a/sdks/python/apache_beam/yaml/yaml_io_test.py +++ b/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()