Add table append method

Author: benjeffery

docs/conf.py

@@ -311,6 +311,7 @@ def handle_item(fieldarg, content):
     # TODO these have been triaged here to make the docs compile, but we should
     # sort them out properly. https://github.com/tskit-dev/tskit/issues/336
     ("py:class", "array_like"),
+    ("py:class", "row-like"),
     ("py:class", "array-like"),
     ("py:class", "dtype=np.uint32"),
     ("py:class", "dtype=np.uint32."),

python/tests/simplify.py

@@ -24,6 +24,7 @@
 Python implementation of the simplify algorithm.
 """
 import sys
+from dataclasses import replace
 
 import numpy as np
 import portion
@@ -158,12 +159,11 @@ def record_node(self, input_id, is_sample=False):
         flags &= ~tskit.NODE_IS_SAMPLE
         if is_sample:
             flags |= tskit.NODE_IS_SAMPLE
-        output_id = self.tables.nodes.add_row(
-            flags=flags,
-            time=node.time,
-            population=node.population,
-            metadata=node.metadata,
-            individual=node.individual,
+        output_id = self.tables.nodes.append(
+            replace(
+                node,
+                flags=flags,
+            )
         )
         self.node_id_map[input_id] = output_id
         return output_id
@@ -186,9 +186,7 @@ def flush_edges(self):
         num_edges = 0
         for child in sorted(self.edge_buffer.keys()):
             for edge in self.edge_buffer[child]:
-                self.tables.edges.add_row(
-                    edge.left, edge.right, edge.parent, edge.child
-                )
+                self.tables.edges.append(edge)
                 num_edges += 1
         self.edge_buffer.clear()
         return num_edges
@@ -413,19 +411,15 @@ def finalise_sites(self):
                         mapped_parent = -1
                         if mut.parent != -1:
                             mapped_parent = mutation_id_map[mut.parent]
-                        self.tables.mutations.add_row(
-                            site=len(self.tables.sites),
-                            node=self.mutation_node_map[mut.id],
-                            time=mut.time,
-                            parent=mapped_parent,
-                            derived_state=mut.derived_state,
-                            metadata=mut.metadata,
+                        self.tables.mutations.append(
+                            replace(
+                                mut,
+                                site=len(self.tables.sites),
+                                node=self.mutation_node_map[mut.id],
+                                parent=mapped_parent,
+                            )
                         )
-                self.tables.sites.add_row(
-                    position=site.position,
-                    ancestral_state=site.ancestral_state,
-                    metadata=site.metadata,
-                )
+                self.tables.sites.append(site)
 
     def finalise_references(self):
         input_populations = self.ts.tables.populations
@@ -455,17 +449,12 @@ def finalise_references(self):
         for input_id, count in enumerate(population_ref_count):
             if count > 0:
                 row = input_populations[input_id]
-                output_id = self.tables.populations.add_row(metadata=row.metadata)
+                output_id = self.tables.populations.append(row)
                 population_id_map[input_id] = output_id
         for input_id, count in enumerate(individual_ref_count):
             if count > 0:
                 row = input_individuals[input_id]
-                output_id = self.tables.individuals.add_row(
-                    flags=row.flags,
-                    location=row.location,
-                    parents=row.parents,
-                    metadata=row.metadata,
-                )
+                output_id = self.tables.individuals.append(row)
                 individual_id_map[input_id] = output_id
 
         # Remap the population ID references for nodes.
@@ -489,11 +478,11 @@ def finalise_references(self):
                     mapped_parents.append(-1)
                 else:
                     mapped_parents.append(individual_id_map[p])
-            self.tables.individuals.add_row(
-                flags=row.flags,
-                location=row.location,
-                parents=mapped_parents,
-                metadata=row.metadata,
+            self.tables.individuals.append(
+                replace(
+                    row,
+                    parents=mapped_parents,
+                )
             )
 
         # We don't support migrations for now. We'll need to remap these as well.
@@ -710,7 +699,7 @@ def flush_edges(self):
         num_edges = 0
         for child in sorted(self.edge_buffer.keys()):
             for edge in self.edge_buffer[child]:
-                self.table.add_row(edge.left, edge.right, edge.parent, edge.child)
+                self.table.append(edge)
                 num_edges += 1
         self.edge_buffer.clear()
         return num_edges

python/tests/test_highlevel.py

@@ -40,6 +40,7 @@
 import unittest
 import uuid as _uuid
 import warnings
+from dataclasses import replace
 
 import kastore
 import msprime
@@ -2854,15 +2855,14 @@ def verify_random_permutation(self, ts):
         inv_node_map = {v: k for k, v in node_map.items()}
         for j in range(ts.num_nodes):
             node = ts.node(inv_node_map[j])
-            other_tables.nodes.add_row(
-                flags=node.flags, time=node.time, population=node.population
-            )
+            other_tables.nodes.append(node)
         for e in ts.edges():
-            other_tables.edges.add_row(
-                left=e.left,
-                right=e.right,
-                parent=node_map[e.parent],
-                child=node_map[e.child],
+            other_tables.edges.append(
+                replace(
+                    e,
+                    parent=node_map[e.parent],
+                    child=node_map[e.child],
+                )
             )
         for _ in range(ts.num_populations):
             other_tables.populations.add_row()

python/tests/test_parsimony.py

@@ -24,6 +24,7 @@
 """
 import io
 import itertools
+from dataclasses import replace
 
 import attr
 import Bio.Phylo.TreeConstruction
@@ -549,18 +550,20 @@ def verify(self, ts):
                 ancestral_state, mutations = self.do_map_mutations(
                     tree, G[site.id], alleles[site.id]
                 )
-                site_id = tables.sites.add_row(site.position, ancestral_state)
+                site_id = tables.sites.append(
+                    replace(site, ancestral_state=ancestral_state)
+                )
                 parent_offset = len(tables.mutations)
                 for mutation in mutations:
                     parent = mutation.parent
                     if parent != tskit.NULL:
                         parent += parent_offset
-                    tables.mutations.add_row(
-                        site_id,
-                        node=mutation.node,
-                        time=mutation.time,
-                        parent=parent,
-                        derived_state=mutation.derived_state,
+                    tables.mutations.append(
+                        replace(
+                            mutation,
+                            site=site_id,
+                            parent=parent,
+                        )
                     )
         other_ts = tables.tree_sequence()
         for h1, h2 in zip(
@@ -715,18 +718,20 @@ def verify(self, ts):
                 ancestral_state, mutations = self.do_map_mutations(
                     tree, G[site.id], alleles[site.id]
                 )
-                site_id = tables.sites.add_row(site.position, ancestral_state)
+                site_id = tables.sites.append(
+                    replace(site, ancestral_state=ancestral_state)
+                )
                 parent_offset = len(tables.mutations)
                 for m in mutations:
                     parent = m.parent
                     if m.parent != tskit.NULL:
                         parent = m.parent + parent_offset
-                    tables.mutations.add_row(
-                        site_id,
-                        node=m.node,
-                        time=m.time,
-                        parent=parent,
-                        derived_state=m.derived_state,
+                    tables.mutations.append(
+                        replace(
+                            m,
+                            site=site_id,
+                            parent=parent,
+                        )
                     )
         other_ts = tables.tree_sequence()
         assert ts.num_samples == other_ts.num_samples
@@ -1206,12 +1211,12 @@ def verify(self, ts, k):
                 parent = mutation.parent
                 if parent != tskit.NULL:
                     parent += parent_offset
-                tables.mutations.add_row(
-                    j,
-                    node=mutation.node,
-                    time=mutation.time,
-                    parent=parent,
-                    derived_state=mutation.derived_state,
+                tables.mutations.append(
+                    replace(
+                        mutation,
+                        site=j,
+                        parent=parent,
+                    )
                 )
 
         ts2 = tables.tree_sequence()

python/tests/test_tables.py

@@ -386,6 +386,20 @@ def test_add_row_round_trip(self):
                 t2.add_row(**dataclasses.asdict(row))
             assert t1 == t2
 
+    def test_append_row(self):
+        for num_rows in [0, 10, 100]:
+            table = self.table_class()
+            for j, row in enumerate(self.make_transposed_input_data(num_rows)):
+                k = table.append(table.row_class(**row))
+                assert k == j
+            for colname, input_array in self.make_input_data(num_rows).items():
+                output_array = getattr(table, colname)
+                assert input_array.shape == output_array.shape
+                assert np.all(input_array == output_array)
+            table.clear()
+            assert table.num_rows == 0
+            assert len(table) == 0
+
     def test_set_columns_data(self):
         for num_rows in [0, 10, 100, 1000]:
             input_data = {col.name: col.get_input(num_rows) for col in self.columns}
@@ -1627,7 +1641,7 @@ def verify_edge_sort_offset(self, ts):
             all_edges = keep + reversed_edges
             tables.edges.clear()
             for e in all_edges:
-                tables.edges.add_row(e.left, e.right, e.parent, e.child)
+                tables.edges.append(e)
             # Verify that import fails for randomised edges
             with pytest.raises(_tskit.LibraryError):
                 tables.tree_sequence()
@@ -1638,7 +1652,7 @@ def verify_edge_sort_offset(self, ts):
             # Sorting from the correct index should give us back the original table.
             tables.edges.clear()
             for e in all_edges:
-                tables.edges.add_row(e.left, e.right, e.parent, e.child)
+                tables.edges.append(e)
             tables.sort(edge_start=start)
             # Verify the new and old edges are equal.
             assert edges == tables.edges

python/tests/test_topology.py

@@ -31,6 +31,7 @@
 import random
 import sys
 import unittest
+from dataclasses import replace
 
 import msprime
 import numpy as np
@@ -66,23 +67,18 @@ def simple_keep_intervals(tables, intervals, simplify=True, record_provenance=Tr
             if not (edge.right <= interval_left or edge.left >= interval_right):
                 left = max(interval_left, edge.left)
                 right = min(interval_right, edge.right)
-                tables.edges.add_row(
-                    left, right, edge.parent, edge.child, edge.metadata
-                )
+                tables.edges.append(replace(edge, left=left, right=right))
     for site in ts.sites():
         for interval_left, interval_right in intervals:
             if interval_left <= site.position < interval_right:
-                site_id = tables.sites.add_row(
-                    site.position, site.ancestral_state, site.metadata
-                )
+                site_id = tables.sites.append(site)
                 for m in site.mutations:
-                    tables.mutations.add_row(
-                        site=site_id,
-                        node=m.node,
-                        derived_state=m.derived_state,
-                        parent=tskit.NULL,
-                        time=m.time,
-                        metadata=m.metadata,
+                    tables.mutations.append(
+                        replace(
+                            m,
+                            site=site_id,
+                            parent=tskit.NULL,
+                        )
                     )
     tables.build_index()
     tables.compute_mutation_parents()
@@ -2392,9 +2388,7 @@ def verify_unary_tree_sequence(self, ts):
         edges.sort(key=lambda e: node_times[e.parent])
         tables.edges.reset()
         for e in edges:
-            tables.edges.add_row(
-                left=e.left, right=e.right, child=e.child, parent=e.parent
-            )
+            tables.edges.append(e)
         ts_new = tables.tree_sequence()
         assert ts_new.num_edges > ts.num_edges
         self.assert_haplotypes_equal(ts, ts_new)
@@ -3471,8 +3465,8 @@ def test_unary_non_sample_external_nodes(self):
         next_node = ts.num_nodes
         tables.edges.reset()
         for e in ts.edges():
-            tables.edges.add_row(e.left, e.right, e.parent, e.child)
-            tables.edges.add_row(e.left, e.right, e.parent, next_node)
+            tables.edges.append(e)
+            tables.edges.append(replace(e, child=next_node))
             tables.nodes.add_row(time=0)
             next_node += 1
         tables.sort()
@@ -6684,7 +6678,7 @@ def do_squash(self, ts, compare_lib=True):
             squashed_list = squash_edges(ts)
             squashed_py = tskit.EdgeTable()
             for e in squashed_list:
-                squashed_py.add_row(e.left, e.right, e.parent, e.child)
+                squashed_py.append(e)
             # Check the Python and C implementations produce the same output.
             assert squashed_py == squashed
         return squashed
@@ -6888,7 +6882,7 @@ def verify_slice_and_squash(self, ts):
         random.shuffle(sliced_edges)
         sliced_table = tskit.EdgeTable()
         for e in sliced_edges:
-            sliced_table.add_row(e.left, e.right, e.parent, e.child)
+            sliced_table.append(e)
 
         # Squash the edges and check against input table.
         sliced_table.squash()
@@ -6949,7 +6943,7 @@ def add_edge(left, right, parent, child):
                 # Squash
                 edge.right = right
             else:
-                tables.edges.add_row(edge.left, edge.right, edge.parent, edge.child)
+                tables.edges.append(edge)
                 edge_map[child] = new_edge
 
     tables.edges.clear()
@@ -6981,7 +6975,7 @@ def add_edge(left, right, parent, child):
         add_edge(left, tables.sequence_length, parent, child)
     # Flush the remaining edges to the table
     for edge in edge_map.values():
-        tables.edges.add_row(edge.left, edge.right, edge.parent, edge.child)
+        tables.edges.append(edge)
     tables.sort()
     ts = tables.tree_sequence()
     # Now simplify to remove redundant nodes.
@@ -8057,7 +8051,7 @@ def ts_missing_middle(self):
                     missing_to = e.left
                 else:
                     continue  # omit this edge => node is isolated
-            tables.edges.add_row(e.left, e.right, e.parent, e.child)
+            tables.edges.append(e)
         # Check we have non-missing to L & R
         assert 0.0 < missing_from < 1.0
         assert 0.0 < missing_to < 1.0