From fa1acb693f5ba7e8d16c735016b137ad9950164a Mon Sep 17 00:00:00 2001 From: aibrahiim Date: Thu, 6 Aug 2026 14:29:13 +0300 Subject: [PATCH 1/2] normalize io.gcp.DicomSearch --- sdks/python/apache_beam/yaml/standard_io.yaml | 1 + sdks/python/apache_beam/yaml/yaml_io.py | 72 +++++++++ sdks/python/apache_beam/yaml/yaml_io_test.py | 151 ++++++++++++++++++ 3 files changed, 224 insertions(+) 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..d3c02b444d65 100644 --- a/sdks/python/apache_beam/yaml/yaml_io.py +++ b/sdks/python/apache_beam/yaml/yaml_io.py @@ -909,3 +909,75 @@ 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', {}))) + + +@beam.ptransform_fn +@yaml_errors.maybe_with_exception_handling_transform_fn +def dicom_search( + pcoll, + *, + buffer_size: int = 8, + max_workers: int = 5): + """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 and can be routed with ``error_handling``. + + Args: + buffer_size: Number of requests to buffer before flushing. + max_workers: Maximum number of threads used to issue requests. + """ + 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 + + return ( + pcoll + | beam.Map(row_to_dict) + | DicomSearch(buffer_size=buffer_size, max_workers=max_workers) + | 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..ba74e180338f 100644 --- a/sdks/python/apache_beam/yaml/yaml_io_test.py +++ b/sdks/python/apache_beam/yaml/yaml_io_test.py @@ -807,6 +807,157 @@ 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() From 9fccc20fb4f28966198dd884fdc535420bd7e448 Mon Sep 17 00:00:00 2001 From: aibrahiim Date: Thu, 6 Aug 2026 15:48:52 +0300 Subject: [PATCH 2/2] fix dicomSearch error handling and formatting --- sdks/python/apache_beam/yaml/yaml_io.py | 56 ++++++++++++++++---- sdks/python/apache_beam/yaml/yaml_io_test.py | 18 +++---- 2 files changed, 52 insertions(+), 22 deletions(-) diff --git a/sdks/python/apache_beam/yaml/yaml_io.py b/sdks/python/apache_beam/yaml/yaml_io.py index d3c02b444d65..3b2a13ebf7f7 100644 --- a/sdks/python/apache_beam/yaml/yaml_io.py +++ b/sdks/python/apache_beam/yaml/yaml_io.py @@ -921,13 +921,28 @@ def _dicom_search_result_to_row(result): 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 -@yaml_errors.maybe_with_exception_handling_transform_fn def dicom_search( - pcoll, - *, - buffer_size: int = 8, - max_workers: int = 5): + 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 @@ -948,19 +963,22 @@ def dicom_search( - status (str): HTTP status from the DICOM API. - input (str): JSON-encoded copy of the search request. - Failed searches raise and can be routed with ``error_handling``. + 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 + "Please install using 'pip install apache-beam[gcp]'.") from exn def row_to_dict(value): if value is None: @@ -976,8 +994,24 @@ def row_to_dict(value): else: return value - return ( + 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) - | beam.Map(_dicom_search_result_to_row)) + | 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 ba74e180338f..43955233d10a 100644 --- a/sdks/python/apache_beam/yaml/yaml_io_test.py +++ b/sdks/python/apache_beam/yaml/yaml_io_test.py @@ -872,9 +872,8 @@ def _patch_dicom_search(): 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: + with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: result = ( p | beam.Create([ @@ -884,9 +883,7 @@ def test_dicom_search_success(self): dataset_id='dataset', dicom_store_id='store', search_type='instances', - params={ - 'PatientName': 'Alice' - }) + params={'PatientName': 'Alice'}) ]) | YamlTransform( ''' @@ -895,15 +892,14 @@ def test_dicom_search_success(self): assert_that( result | beam.Map( - lambda row: ( - row.status, json.loads(row.result)[0]['PatientName'])), + 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: + with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: result = ( p | beam.Create([