Replacing metadata in rows

[hyanwong]

Following on from #2953, I wanted to be able to change metadata in an individual row (or rows) without having to validate/reencode the existing metadata. I think this does it:

def change_ragged_row(index, replacement_row, ragged_arr, ragged_arr_offset):
    # insert a single new replacement_row into the ragged_arr at index pos
    # returns an updated ragged_col and a new offset array
    if replacement_row.dtype != ragged_arr.dtype:
        raise ValueError(f"`replacement_row` must be an array of dtype {ragged_col.dtype}")
    start = ragged_arr_offset[index]
    end = ragged_arr_offset[index + 1]
    len_diff = (end - start) - len(replacement_row)
    ragged_arr_offset[index + 1:] = ragged_arr_offset[index + 1:] - len_diff
    return (
        np.concatenate((ragged_arr[:start], replacement_row, ragged_arr[end:])),
        ragged_arr_offset
    )

def change_metadata_row(index, md, table):
    if table.metadata_schema.schema is not None:
       md = table.metadata_schema.validate_and_encode_row(md)
    byte_md = np.array(bytearray(md), dtype=table.metadata.dtype)
    d = table.asdict()
    d["metadata"], d["metadata_offset"] = change_ragged_row(
        index, byte_md, table.metadata, table.metadata_offset)
    table.set_columns(**d)

Here's a test.

import tskit
import numpy as np
tables = tskit.TableCollection(10)
s = tables.sites
s.add_row(0, "A")
s.add_row(1, "")
s.add_row(2, "TTT")
s.add_row(3, "G")
print(s)

change_metadata_row(0,  b'{"id": 0}', tables.sites)
print(tables.sites)
change_metadata_row(1,  b'{"id": "one"}', tables.sites)
print(tables.sites)
change_metadata_row(0,  b'{"id": "zero"}', tables.sites)
print(tables.sites)
change_metadata_row(1,  b'', tables.sites)
print(tables.sites)

A useful improvement to this would be to be able to replace multiple rows at once (maybe input could be a dict mapping {index: md})

[hyanwong]

Actually, the most efficient way to change multiple rows is probably to add them onto the end and then run the code in #2953:

import tskit
import numpy as np

def change_ragged_order(ragged_arr, ragged_offset, new_order):
    # returns a tuple of a new ragged_col in a new order and a new offset array
    ranges = np.array([ragged_offset[:-1], ragged_offset[1:]]).T
    new_ranges = ranges[new_order, :]
    idx = [np.arange(l, r, dtype=ragged_offset.dtype) for l, r in new_ranges if l != r]
    select = [] if len(idx) == 0 else np.concatenate(idx)
    return ragged_arr[select], np.insert(np.cumsum(np.diff(new_ranges, axis=1)), 0, 0)

def change_metadata(new_md_dict, table):
    if table.metadata_schema.schema is not None:
        for k, v in new_md_dict.items():
            new_md_dict[k] = table.metadata_schema.validate_and_encode_row(v)
    data = [table.metadata]
    # add a list of new byte arrays, then concat
    data += [np.array(bytearray(v), dtype=table.metadata.dtype) for v in new_md_dict.values()]
    tmp_offset = np.cumsum([len(d) for d in data], dtype=table.metadata_offset.dtype)[1:]
    tmp_offset = np.concatenate((table.metadata_offset, tmp_offset))
    tmp_md = np.concatenate(data)
    new_row_ids = np.arange(len(new_md_dict)) + table.num_rows
    idx = np.arange(table.num_rows)
    idx[list(new_md_dict.keys())] = new_row_ids
    d = table.asdict()
    d["metadata"], d["metadata_offset"] = change_ragged_order(tmp_md, tmp_offset, idx)
    table.set_columns(**d)

Test:

# TEST
tables = tskit.TableCollection(10)
s = tables.sites
s.add_row(0, "A")
s.add_row(1, "")
s.add_row(2, "TTT")
s.add_row(3, "G")

change_metadata_row(0,  b'{"id": 0}', tables.sites)
change_metadata_row(1,  b'{"id": "one"}', tables.sites)
print(tables.sites)
change_metadata({2: b"replacement for row id 2"}, tables.sites)
print(tables.sites)
change_metadata({0: b"replacement for row id 0"}, tables.sites)
print(tables.sites)