def test_multireadwrite(ms, group_cols, index_cols): xds = xds_from_ms(ms, group_cols=group_cols, index_cols=index_cols) nds = [ds.copy() for ds in xds] writes = [xds_to_table(sds, ms, [k for k in sds.data_vars.keys() if k != "ROWID"]) for sds in nds] da.compute(writes)
def test_ms_update(ms, group_cols, index_cols, select_cols): # Zero everything to be sure with TableProxy(pt.table, ms, readonly=False, lockoptions='auto', ack=False) as T: nrows = T.nrows().result() T.putcol("STATE_ID", np.full(nrows, 0, dtype=np.int32)).result() data = np.zeros_like(T.getcol("DATA").result()) data_dtype = data.dtype T.putcol("DATA", data).result() xds = xds_from_ms(ms, columns=select_cols, group_cols=group_cols, index_cols=index_cols, chunks={"row": 2}) written_states = [] written_data = [] writes = [] # Write out STATE_ID and DATA for i, ds in enumerate(xds): dims = ds.dims chunks = ds.chunks state = da.arange(i, i + dims["row"], chunks=chunks["row"]) state = state.astype(np.int32) written_states.append(state) data = da.arange(i, i + dims["row"] * dims["chan"] * dims["corr"]) data = data.reshape(dims["row"], dims["chan"], dims["corr"]) data = data.rechunk((chunks["row"], chunks["chan"], chunks["corr"])) data = data.astype(data_dtype) written_data.append(data) nds = ds.assign(STATE_ID=(("row", ), state), DATA=(("row", "chan", "corr"), data)) write = xds_to_table(nds, ms, ["STATE_ID", "DATA"]) writes.append(write) # Do all writes in parallel dask.compute(writes) xds = xds_from_ms(ms, columns=select_cols, group_cols=group_cols, index_cols=index_cols, chunks={"row": 2}) # Check that state and data have been correctly written it = enumerate(zip(xds, written_states, written_data)) for i, (ds, state, data) in it: assert_array_equal(ds.STATE_ID.data, state) assert_array_equal(ds.DATA.data, data)
def _proc_map_fn(args): try: ms, i = args xds = xds_from_ms(ms, columns=["STATE_ID"], group_cols=["FIELD_ID"]) xds[i] = xds[i].assign(STATE_ID=(("row", ), xds[i].STATE_ID.data + i)) write = xds_to_table(xds[i], ms, ["STATE_ID"]) write.compute(scheduler='sync') except Exception as e: print(str(e)) return True
def test_write_array_names(ms, tmp_path): _, short_name, _ = table_path_split(ms) datasets = xds_from_ms(ms) out_table = str(tmp_path / short_name) writes = xds_to_table(datasets, out_table, "ALL") for ds in writes: for k, v in ds.data_vars.items(): prefix = "".join(("write~", k)) assert key_split(v.data.name) == prefix
def _proc_map_fn(args): import dask.threaded as dt # No dask pools are spun up with dt.pools_lock: assert dt.default_pool is None assert len(dt.pools) == 0 try: ms, i = args xds = xds_from_ms(ms, columns=["STATE_ID"], group_cols=["FIELD_ID"]) xds[i] = xds[i].assign(STATE_ID=(("row",), xds[i].STATE_ID.data + i)) write = xds_to_table(xds[i], ms, ["STATE_ID"]) dask.compute(write) except Exception as e: print(str(e)) return True