Limited write support

[sjperkins]

Thanks for contributing to xarray-ms.

We would appreciate it if you could add:

• Test Cases covering your PR.
• Documentation.
• A Changelog entry in doc/source/changelog.rst.

---

📚 Documentation preview 📚: https://xarray-ms--170.org.readthedocs.build/en/170/

[landmanbester]

Hi @sjperkins — @landmanbester pointed me at this branch while I was planning the move of
pfb degrid onto MSv4 + Ray (ratt-ru/pfb-imaging#278). I prototyped against it and it does
exactly what we need: region writes on the (time, frequency) grid remove the entire
row-mapping and column-creation layer I had been about to write by hand. Thank you.

Four findings from that prototyping, all measured. The first is the one that matters to us.

1. The arcae < 0.5.0 pin looks stale — the write path works on arcae 0.5.4

pyproject.toml on this branch pins arcae>=0.4.0a8, < 0.5.0. That is a problem for
pfb-imaging: we need arcae >= 0.5.2 because ska-sa/arcae#211 and #212 are what let arcae
and python-casacore coexist in one process, and our whole test suite is built on that
(pytest tests/ writes the test MS with python-casacore and reads it with arcae). Pinning
below 0.5.0 would take that away.

So I lifted the pin to arcae>=0.5.3 and tried it. The write path works unmodified on
arcae 0.5.4, provided the table is opened with a single instance:

dt = xarray.open_datatree(ms, engine="xarray-ms:msv2",
                          partition_schema=[...], ninstances=1)

Without ninstances=1, sync_msv2() fails at addcols:

pyarrow.lib.ArrowNotImplementedError: Write support when number of table instances 8
is greater than one

which makes sense: arcae 0.5.x reports safe_multithreaded_writes() == False, whereas the
0.4.0a line this branch was developed against was multithread-writeable, so the branch has
no reason to ask for a single-instance handle.

With ninstances=1 I ran six (time, frequency) region writes over a real MS and every
region landed exactly where it should. So the constraint is real but narrow.

Two things that would help consumers here:

• Could the write path acquire (or require) a single-instance handle itself, or at least
  translate that arcae error into something that names ninstances? As it stands the
  failure is quite far from its cause.
• ninstances=1 for the whole DataTree also costs read parallelism, which we would rather
  keep. Is a read-many / write-one split feasible, or is one handle per table the design?

2. sync_msv2() does not flush, so new columns are invisible to other handles

After dt.sync_msv2(write_map={"MODEL": "PFB_MODEL_DATA"}) the new column cannot be seen by
any other table handle until the DataTree is closed:

observer	sees the new column?
fresh arcae.table(ms) in the same process	no
separate process	no
separate process, after dt.close()	yes

For a single-process script this is invisible. For any distributed consumer it is a trap: the
natural shape is driver creates the columns, then fans work out to worker processes, and
those workers open their own handles and will not find the column. We can work around it by
closing the DataTree before dispatch, but that is an undocumented ordering constraint that
every consumer has to rediscover. A flush at the end of sync_msv2 would remove it.

3. Concurrent multi-process region writes are correct (nice!)

Four separate processes, each opening its own DataTree with ninstances=1 and writing a
disjoint time region of the same MS with to_msv2(compute=True, region=...): all four
regions came out correct. That is the scenario Ray puts us in, and it means we need no
writer actor and no serialisation layer at all. Worth having as a test on the 0.5.x line
alongside the existing dask test_distributed_write.

4. MainMSv2Array.__setitem__ does not filter row_map == -1 — but arcae saves us

rows = self._structure_factory.instance[self._partition].row_map[key[:2]]
row_key = (rows.ravel(),) + key[2:]
...putcol(self._column, value, index=row_key)

Absent (time, baseline) cells carry -1, which goes straight into putcol as a row index.
I checked what arcae does with that, since a numpy-style interpretation would silently
scribble on the last row of the table:

putcol("MODEL_DATA", payload, index=([-1, 5, -1, 9],))
  -> accepted, no error
  -> rows actually written: [5 9]
  -> row 5 got payload[1], row 9 got payload[3]   (correctly aligned)
  -> last row of the table untouched

So it is safe: arcae skips negative indices and keeps the payload aligned. Given how bad the
alternative would have been, and that partitions with missing baselines are the normal case
for MeerKAT, this feels like it deserves an explicit test rather than resting on arcae's
behaviour being what it is.

5. Minor: the test simulator is stale against arcae 0.5.x

tests/test_write.py cannot run on 0.5.x because of the fixture, not the write path:

xarray_ms/testing/simulator.py:233: in simulate_ms
    with Table.ms_from_descriptor(output_ms, "MAIN", 1, table_desc) as T:
E   TypeError: Argument 'table_desc' has incorrect type (expected dict, got int)

ms_from_descriptor's signature changed between the two arcae lines. Mentioning it because
it is the first thing anyone rebasing this branch will hit, and it makes the branch look more
broken on 0.5.x than it actually is.

---

Is there anything we can usefully do to help get this landed? pfb-imaging's degrid work is
gated on it — the rest of the plan (a pure-numpy degrid kernel in
landmanbester/pfb-model-spec#22 and a beam helper in
landmanbester/meerkat-beams#27) is independent and can proceed in parallel, but the MSv4
front-end itself waits on this PR. Happy to contribute tests or the ninstances plumbing if
that is useful rather than in the way.

[sjperkins]

The arcae < 0.5.0 pin looks stale — the write path works on arcae 0.5.4

This is by design. The write support branch is versioned in the following range0.4.0 <= versions < 0.5.0 so that downstream consumers only get the read support versions (i.e. >= 0.5.0).

pyproject.toml on this branch pins arcae>=0.4.0a8, < 0.5.0. That is a problem for
pfb-imaging: we need arcae >= 0.5.2 because ska-sa/arcae#211 and #212 are what let arcae
and python-casacore coexist in one process, and our whole test suite is built on that
(pytest tests/ writes the test MS with python-casacore and reads it with arcae). Pinning
below 0.5.0 would take that away.

I think you believe this because of some internal context that you've stored somewhere but the contents of #211 and #212 should have been merged into this branch. If pfb is explicitly depending on write support (much like tricolour does, it should pin 0.4.0 <= xarray-ms < 0.5.0. The other concerns should fall away once this version range is adopted.

After dt.sync_msv2(write_map={"MODEL": "PFB_MODEL_DATA"}) the new column cannot be seen by
any other table handle until the DataTree is closed:
For a single-process script this is invisible. For any distributed consumer it is a trap: the
natural shape is driver creates the columns, then fans work out to worker processes, and
those workers open their own handles and will not find the column. We can work around it by
closing the DataTree before dispatch, but that is an undocumented ordering constraint that
every consumer has to rediscover. A flush at the end of sync_msv2 would remove it.

Correct, this is actually an artifact of the CTDS (creating a column in one process isn't automatically observed by other processes). The suggested workaround is correct.

[landmanbester]

You're right, and my finding #1 was wrong. Retracting it properly, since it was the one that
mattered.

I inferred "0.4.0-alpha.8 predates 0.5.2, therefore lacks #211/#212" from the version
ordering alone. That's exactly backwards — the alpha line is cut from later commits than
the 0.5.x line:

tag	date
0.5.2 (coexistence release)	2026-06-12
0.4.0-alpha.6	2026-06-19
0.5.4	2026-07-22
0.4.0-alpha.8	2026-07-23

And gh api repos/ska-sa/arcae/compare/<merge-sha>...0.4.0-alpha.8 reports
status=ahead, behind_by=0 for both #211 and #212, so the fixes are in, just as you said.

I then re-ran everything on the intended stack rather than my patched one:

• python-casacore + arcae 0.4.0-alpha.8 coexist in one process. Table written by
  python-casacore then read by arcae, and vice versa, both fine. This is the property our
  test suite depends on, so the concern that started all this is gone.
• tests/test_write.py: 6 passed on the unmodified branch, including
  test_distributed_write for both processes=True and False. So finding Documentation Updates #5 (the stale
  simulator) is also void — ms_from_descriptor only mismatches on 0.5.x, which is not the
  line this branch targets.
• The full target stack works with no ninstances fiddling at all. xarray-ms
  write-support + arcae 0.4.0-alpha.8 + python-casacore in one process: sync_msv2 plus six
  (time, frequency) region writes over a real MS with the default multi-instance handle,
  verified by reading the column back with python-casacore. safe_multithreaded_writes() is
  True on that line, so the ninstances=1 requirement I hit was purely an artifact of my
  forcing arcae 0.5.4. Both ninstances questions in my previous comment are withdrawn.

What still stands, for the record:

• Correct uvw dimension to uvw_label #2 (sync_msv2 doesn't flush) — thanks for confirming it's a CTDS property and that
  close-before-dispatch is the intended workaround. We'll do that in the driver. If it's a
  CTDS constraint rather than something xarray-ms can paper over, a line in the docstring
  would probably save the next person the experiment.
• Add initial documentation #4 (row_map == -1) — re-verified on 0.4.0-alpha.8: putcol skips negative row
  indices and keeps the payload aligned, so absent (time, baseline) cells are safe.
  Still think it's worth an explicit test, given that partitions with missing baselines are
  the normal case for MeerKAT and the alternative failure would be silent.
• Make partitioning columns configurable #3 (concurrent multi-process region writes) — your test_distributed_write with
  processes=True already covers this better than my ad-hoc version did.

One genuine question left, about the release plan rather than the code. tricolour can sit on
0.4.0 <= xarray-ms < 0.5.0 happily because it only needs the flagging path. pfb-imaging
would be pinning its imager there too — that's our most performance-sensitive read path and
we've tuned it fairly hard against current behaviour. Checking before worrying: write-support
is 74 commits ahead of main and 0 behind, so today the alpha line is main plus writes and we
lose nothing by pinning down into it. Is keeping that invariant the plan (periodic merges from
main), and is there an intended convergence point where write support lands on the
mainline? That's the only thing that affects how we sequence ratt-ru/pfb-imaging#278 — happy
either way, just want to pin with our eyes open.

Thanks for the correction, and for pointing us at ratt-ru/tricolour#106 — the
WorkItem(path, region) + Multiton[DataTree] + load/compute/write split is exactly the
shape we need, and we intend to follow it rather than invent a parallel one.

[sjperkins]

Checking before worrying: write-support is 74 commits ahead of main and 0 behind, so today the alpha line is main plus writes and we lose nothing by pinning down into it. Is keeping that invariant the plan (periodic merges from main), and is there an intended convergence point where write support lands on the mainline?

The plan is that:

1. Write support will be merged into main once this functionality has been sufficiently stress-tested to confirm there are no issues.
2. Any changes in main are merged into write-support. In principle it is kept in a state of 0 commits behind main.