From 17b83d491ef6bbff66ebf319c2dd1d52635a303b Mon Sep 17 00:00:00 2001 From: Takahiro Tsutsumi Date: Tue, 18 Nov 2025 09:11:36 -0700 Subject: [PATCH 1/2] Fix a bug that was not handling field_name in a list in partitioning ps --- .../graph_tools/coordinate_utils.py | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/src/graphviper/graph_tools/coordinate_utils.py b/src/graphviper/graph_tools/coordinate_utils.py index 3eea799..62d44e2 100644 --- a/src/graphviper/graph_tools/coordinate_utils.py +++ b/src/graphviper/graph_tools/coordinate_utils.py @@ -582,13 +582,28 @@ def _partition_ps_by_non_dimensions(ps, ps_partition_keys): partition_info = xds.xr_ms.get_partition_info() for key in ps_partition_keys: val_for_xds = partition_info[key] + # value could be a list + list_for_xds = None + if isinstance(val_for_xds, list): + list_for_xds = val_for_xds # OK I think I can punt: the key should probably be an integer but that doesn't feel very Pythonic # But I *can* reasonably demand it is hashable - if not isinstance(val_for_xds, Hashable): - raise ValueError( - f"Can't split by {key}; value {val_for_xds} is not suitable for splitting" + if list_for_xds is not None: + for v in list_for_xds: + if not isinstance(v, Hashable): + raise ValueError( + f"Can't split by {key}; value {v} is not suitable for splitting" + ) + ps_split_map.setdefault(key, {}).setdefault(v, []).append(name) + else: + if not isinstance(val_for_xds, Hashable): + raise ValueError( + f"Can't split by {key}; value {val_for_xds} is not suitable for splitting" + ) + ps_split_map.setdefault(key, {}).setdefault(val_for_xds, []).append( + name ) - ps_split_map.setdefault(key, {}).setdefault(val_for_xds, []).append(name) + d = {} # We loop over the cartersian product of the keys for multi_index in itertools.product( From e28aa930732df7fd2c0e1ae72f3b90f87111d74e Mon Sep 17 00:00:00 2001 From: Takahiro Tsutsumi Date: Tue, 18 Nov 2025 15:31:59 -0700 Subject: [PATCH 2/2] Add "field_name" to ps_partition to test the selection and file cleanup --- tests/test_graph_tools.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/tests/test_graph_tools.py b/tests/test_graph_tools.py index 0cff0aa..55f467c 100644 --- a/tests/test_graph_tools.py +++ b/tests/test_graph_tools.py @@ -1,3 +1,20 @@ +from importlib.metadata import files + + +def delete_files(filepaths=None): + import os + import shutil + + if filepaths is None: + return + for file in filepaths: + print(f"Removing {file}...") + if os.path.isdir(file): + shutil.rmtree(file) + elif os.path.isfile(file): + os.remove(file) + + def test_map_reduce(): from toolviper.utils.data import download from graphviper.graph_tools.map import map @@ -101,6 +118,9 @@ def my_sum(graph_inputs, input_params): dask_graph = generate_dask_workflow(graph_reduce) assert dask.compute(dask_graph)[0] == 178177980857.54022 + delete_files( + filepaths=[ms_name, ps_store], + ) def test_ps_partition(): @@ -135,11 +155,11 @@ def test_ps_partition(): node_task_data_mapping = interpolate_data_coords_onto_parallel_coords( parallel_coords=parallel_coords, input_data=ps, - ps_partition=["spectral_window_name"], + ps_partition=["spectral_window_name", "field_name"], ) # print(node_task_data_mapping) - assert len(node_task_data_mapping.keys()) == 2 + assert len(node_task_data_mapping.keys()) == 4 # We check that for each data selection the spw_id is unique: spw_split_success = all( [ @@ -157,11 +177,14 @@ def test_ps_partition(): ) assert spw_split_success + delete_files( + filepaths=[msv2name, zarrPath], + ) + if __name__ == "__main__": test_map_reduce() test_ps_partition() - """ chunk_indx 0 (0, 0) chunk_indx 1 (0, 1)