Bulk parsing json metadata

[jeromekelleher]

For large tables, accessing metadata row-by-row using the tskit API can be prohibitively slow. Sometimes we want columnar access to the data.

One way to do this is to use the PyArrow read_json function, although it requires some jiggery-pokery.

import numpy as np
import numba
import pyarrow.json

@numba.njit
def insert_newlines(metadata, metadata_offset):
    output = np.zeros(metadata.shape[0] + metadata_offset.shape[0] - 1, dtype=np.int8)
    i = 0
    for j in range(metadata_offset.shape[0] - 1):
        for k in range(metadata_offset[j], metadata_offset[j + 1]):
            output[i] = metadata[k]
            i += 1
        output[i] = ord('\n')
        i += 1
    return output

nodes = tables.nodes
buff = insert_newlines(nodes.metadata, nodes.metadata_offset)
# remove the table reference - only doing this as we don't have zero-copy access to the 
# metadata and metadata offset cols
del nodes 

with open("/tmp/tmp.json", "wb+") as f:
    f.write(buff.tobytes())
    f.seek(0)
    table = pyarrow.json.read_json(f)

# Now you can access metadata columns like table["Viridian_Pangolin"].to_numpy()

First we need to insert newlines between all the metadata records from the table, so that PyArrow can read it. Then, we have to write that to a file.

It is pretty fast, at about 5 seconds end-to-end to parse about 1G of quite loosely structured JSON metadata. I think the read_json bit could be a good bit faster if the parser was given an explicit schema, and told which specific fields to pull out.

However, it's still pretty slow as there's quite a lot of copying going on. Unfortunately Arrow only seems to support reading from actual files, so we have to go out to the file system. If there was some way of getting arrow to read directly from the buffer we're making, then maybe it would be worth trying to make this path smoother.

This isn't quite what I need currently, so I'm not going to pursue any further. However, I thought that this might prove useful to someone in the future so I posted it here.

If it is, we should totally add support for inserting newlines between the metadata records in tskit, as this is a very reasonable thing to want.

[benjeffery]

Your use case here is to fetch a single column as a numpy array? Or possibly "fetch a list of keys as a set of numpy arrays"?

If we make it a requirement that there is a schema and only support basic types I think those use cases would be doable by doing the schema checking in Python, then passing down the columns-with-types request to a C json parser, maybe Cjson or YAJL (both MIT-style License). The C would error out if the wrong type was encountered, and pass back the numpy arrays.