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
