Use takeset to make loading zero-copy

[benjeffery]

Fixes #2131

Here is a comparison of this main and this branch's memory use (main first):

Almost all of the remaining overhead is due to casting offsets up to 64bit, which uses 3x their size, with that off we get:

[petrelharp]

Awesome to see this moving. What's memtestmsp.py doing?

[benjeffery]

Awesome to see this moving. What's memtestmsp.py doing?

import tskit
tables = tskit.TableCollection.load("big.trees")

This skips the extra copy for a tree sequence, which I'll deal with separately.

[petrelharp]

Nice! So we're seeing a factor of 2 reduction in memory except for this offset casting thing?

[codecov]

Codecov Report

Merging #2143 (df4f635) into main (13499ff) will decrease coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #2143      +/-   ##
==========================================
- Coverage   93.30%   93.29%   -0.02%     
==========================================
  Files          27       27              
  Lines       25818    25895      +77     
  Branches     1163     1163              
==========================================
+ Hits        24090    24159      +69     
- Misses       1698     1706       +8     
  Partials       30       30              

Flag	Coverage Δ
c-tests	92.21% <100.00%> (-0.02%)	⬇️
lwt-tests	89.05% <ø> (ø)
python-c-tests	72.09% <ø> (ø)
python-tests	98.89% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
c/tskit/tables.c	90.29% <100.00%> (-0.01%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 13499ff...df4f635. Read the comment docs.

[benjeffery]

All the tests finally pass here. Some feedback on how I've solved the edge cases (e.g. if after out:) would be good.

[jeromekelleher]

Let's get #2032 squared away first, will make it easier to see the new stuff here. I think we're pretty sold on the takeset idea basically working anyway, right?

[benjeffery]

Here's a clean example that we're getting what we expect in a very simple case - this is a tree sequence with only one row in the populations table: tc.populations.add_row(metadata=b"a"*10_000_000_000) and the memory usage when we load it to a TableCollection:

[benjeffery]

Here is the case for an msprime generated file:

Summarising events on the top (main) plot:

• 0 python initialises and imports tskit.
• 0.15 Loading starts, kastore allocs
• ~0.19 kastore completes, calls to set_columns happen along with cast_offsets (0.22s blip is a cast buffer being freed)
• 0.24 kastore closes and frees
  (time.sleep(0.1) happens here, added so we can see the settled value for the table collection)
• 0.36 table collection freed
• 0.4 python frees its stuff.

This branch is similar, but there are no set column allocs or kastore frees so it is shorter. You can still see the cast_offset allocations at the peak.

[jeromekelleher]

Brilliant! So all working as expected now then? (Analysis super helpful btw)

[benjeffery]

So all working as expected now then?

Yes, I believe so! It has been useful to break all this down and check we understand it all. Basically, the 0.05seconds after 0.2 are snipped out and don't happen, so as well as the memory usage falling the load is quicker!

[benjeffery]

@Mergifyio rebase

[mergify]

rebase

✅ Branch has been successfully rebased

[benjeffery]

Now that I've got a clean diff I can see a couple of errors, will fix today so don't review yet!

[benjeffery]

@jeromekelleher Ok! Ready for review.

[benjeffery]

@jeromekelleher I've made the changes we discussed - in doing so it revealed that read_table_ragged_col_t had two references to the offset array, one of which was direct and therefore non-NULLable by the caller of read_tables. I couldn't see why this was needed so have removed the second reference, it was added when we added support for 64bit.

[jeromekelleher]

Looks great, let's merge!