How to reorder a ragged column, e.g. metadata?

[hyanwong]

I'm wanting to reorder the node table, duplicating some nodes. This is easy for non-ragged columns:

tables.nodes.set_columns(
    flags = tables.nodes.flags[map],
    time=tables.nodes.time[map], 
    ...

But it's less obvious how to reorder the metadata, especially if I don't want to go through the hassle of decoding it etc. Has anyone got a routine to cut up and reorder the bytes of the metadata array given a reordering of the metadata_offset values? The discussion in https://tskit.dev/tutorials/tables_and_editing.html#minor-edits parses it and repacks with packset_metadata, but there should be a more efficient bytewise way, right?

[jeromekelleher]

I don't think there's an easy answer. Basically we want a per table C level function like keep_rows that takes an array of row ids (of any length, allowing duplicates etc), and does the low level stuff to update the table in place.

[hyanwong]

Thanks Jerome. On a similar note, I'm not sure there's a way to get the raw bytes out of the metadata field for each row, so that we can bung them back in using packset_metadata without having to go through the validation etc step (assuming we aren;'t actually altering the metadata, so it doesn't need revalidation / encoding.

That might be a useful function to have. It would help if (as is common) we want to change just one row of the metadata: we can pass the raw metadata back in to the unaltered rows, and just decode/validate/re-encode the one which we are altering.

[hyanwong]

I think this might work for simple reordering (untested). I'm not sure how efficient it is to call np.arange repeatedly inside a list comprehension, however:

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]
    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)

# e.g.
tables = ts.dump_tables()
reorder = np.array([2, 1, 0, 3]) # for example
tables.nodes.set_columns(
    tables.nodes.flags[reorder],
    tables.nodes.time[reorder],
    tables.nodes.population[reorder],
    tables.nodes.individual[reorder],
    *change_ragged_order(
      tables.nodes.metadata, tables.nodes.metadata_offset, reorder
    ),
)

# Also have to change any refs to the new nodes below
...

reordered_ts = table.tree_sequence()

Since many rows might have empty metadata, I wonder if it's worth omitting empty rows inside the list comprehension:

    idx = [np.arange(l, r, dtype=ragged_offset.dtype) for l, r in new_ranges if l != r]