From d7332572718f423c0d1f7891584edd527c19dabe Mon Sep 17 00:00:00 2001 From: JSKenyon Date: Thu, 21 Sep 2023 10:45:29 +0200 Subject: [PATCH] Add rechunking in fragments code and ROWID to columns in backup app. --- quartical/apps/backup.py | 4 ++-- quartical/data_handling/ms_handler.py | 7 +++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/quartical/apps/backup.py b/quartical/apps/backup.py index 1a15ea8c..594b7d4b 100644 --- a/quartical/apps/backup.py +++ b/quartical/apps/backup.py @@ -67,7 +67,7 @@ def backup(): # This call exists purely to get the relevant shape and dtype info. data_xds_list = xds_from_storage_ms( args.ms_path, - columns=column_name, + columns=(column_name, "ROWID"), index_cols=("TIME",), group_cols=("FIELD_ID", "DATA_DESC_ID", "SCAN_NUMBER"), ) @@ -81,7 +81,7 @@ def backup(): # Repeat of above call but now with correct chunking information. data_xds_list = xds_from_storage_ms( args.ms_path, - columns=column_name, + columns=(column_name, "ROWID"), index_cols=("TIME",), group_cols=("FIELD_ID", "DATA_DESC_ID", "SCAN_NUMBER"), chunks=chunks diff --git a/quartical/data_handling/ms_handler.py b/quartical/data_handling/ms_handler.py index 4768d62e..ca652fc8 100644 --- a/quartical/data_handling/ms_handler.py +++ b/quartical/data_handling/ms_handler.py @@ -8,6 +8,7 @@ xds_from_table_fragment, xds_to_table_fragment ) +from daskms.experimental.utils import rechunk_by_size from dask.graph_manipulation import clone from loguru import logger from quartical.weights.weights import initialize_weights @@ -300,6 +301,12 @@ def write_xds_list(xds_list, ref_xds_list, ms_path, output_opts): warnings.simplefilter("ignore") if output_opts.fragment_path: + + xds_list = [ + rechunk_by_size(xds, {'corr'}, only_when_needed=True) + for xds in xds_list + ] + write_xds_list = xds_to_table_fragment( xds_list, output_opts.fragment_path,