Basic delete_rows operation for node table

Author: jeromekelleher

c/tests/test_tables.c

@@ -1490,6 +1490,64 @@ test_node_table_update_row(void)
     tsk_node_table_free(&table);
 }
 
+static void
+test_node_table_delete_rows(void)
+{
+    int ret;
+    tsk_id_t ret_id;
+    tsk_node_table_t table;
+    tsk_node_table_t copy;
+    tsk_node_t row;
+    bool delete[3] = { 0, 0, 0 };
+    tsk_id_t id_map[3];
+    const char *metadata = "ABC";
+
+    ret = tsk_node_table_init(&table, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret_id = tsk_node_table_add_row(&table, 0, 1.0, 2, 3, metadata, 1);
+    CU_ASSERT_FATAL(ret_id >= 0);
+    ret_id = tsk_node_table_add_row(&table, 1, 2.0, 3, 4, metadata, 2);
+    CU_ASSERT_FATAL(ret_id >= 0);
+    ret_id = tsk_node_table_add_row(&table, 2, 3.0, 4, 5, metadata, 3);
+    CU_ASSERT_FATAL(ret_id >= 0);
+
+    ret = tsk_node_table_copy(&table, &copy, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = tsk_node_table_delete_rows(&copy, delete, 0, id_map);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_node_table_equals(&copy, &table, 0));
+
+    delete[0] = 1;
+    delete[1] = 1;
+    delete[2] = 1;
+    ret = tsk_node_table_delete_rows(&copy, delete, 0, id_map);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_EQUAL_FATAL(copy.num_rows, 0);
+
+    ret = tsk_node_table_copy(&table, &copy, TSK_NO_INIT);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    delete[0] = 1;
+    delete[1] = 0;
+    delete[2] = 1;
+    ret = tsk_node_table_delete_rows(&copy, delete, 0, id_map);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_EQUAL_FATAL(copy.num_rows, 1);
+
+    ret = tsk_node_table_get_row(&copy, 0, &row);
+    CU_ASSERT_EQUAL_FATAL(row.flags, 1);
+    CU_ASSERT_EQUAL_FATAL(row.time, 2.0);
+    CU_ASSERT_EQUAL_FATAL(row.population, 3);
+    CU_ASSERT_EQUAL_FATAL(row.individual, 4);
+    CU_ASSERT_EQUAL_FATAL(row.metadata_length, 2);
+    CU_ASSERT_EQUAL_FATAL(row.metadata[0], 'A');
+    CU_ASSERT_EQUAL_FATAL(row.metadata[1], 'B');
+
+    tsk_node_table_free(&copy);
+    tsk_node_table_free(&table);
+}
+
 static void
 test_edge_table_with_options(tsk_flags_t options)
 {
@@ -10426,6 +10484,7 @@ main(int argc, char **argv)
     CU_TestInfo tests[] = {
         { "test_node_table", test_node_table },
         { "test_node_table_update_row", test_node_table_update_row },
+        { "test_node_table_delete_rows", test_node_table_delete_rows },
         { "test_node_table_takeset", test_node_table_takeset },
         { "test_edge_table", test_edge_table },
         { "test_edge_table_update_row", test_edge_table_update_row },

c/tskit/tables.c

@@ -2271,6 +2271,46 @@ tsk_node_table_get_row(const tsk_node_table_t *self, tsk_id_t index, tsk_node_t
     return ret;
 }
 
+int
+tsk_node_table_delete_rows(tsk_node_table_t *self, bool *delete_rows,
+    tsk_flags_t TSK_UNUSED(options), tsk_id_t *id_map)
+{
+    int ret;
+    tsk_node_table_t copy;
+    tsk_node_t node;
+    tsk_size_t j;
+    tsk_id_t ret_id;
+
+    ret = tsk_node_table_copy(self, &copy, 0);
+    if (ret != 0) {
+        goto out;
+    }
+    ret = tsk_node_table_clear(self);
+    if (ret != 0) {
+        goto out;
+    }
+    for (j = 0; j < copy.num_rows; j++) {
+        if (id_map != NULL) {
+            id_map[j] = TSK_NULL;
+        }
+        if (!delete_rows[j]) {
+            tsk_node_table_get_row_unsafe(&copy, (tsk_id_t) j, &node);
+            ret_id = tsk_node_table_add_row(self, node.flags, node.time, node.population,
+                node.individual, node.metadata, node.metadata_length);
+            if (ret_id < 0) {
+                ret = (int) ret_id;
+                goto out;
+            }
+            if (id_map != NULL) {
+                id_map[j] = ret_id;
+            }
+        }
+    }
+out:
+    tsk_node_table_free(&copy);
+    return ret;
+}
+
 static int
 tsk_node_table_dump(const tsk_node_table_t *self, kastore_t *store, tsk_flags_t options)
 {

c/tskit/tables.h

@@ -1389,6 +1389,35 @@ and is not checked for compatibility with any existing schema on this table.
 int tsk_node_table_extend(tsk_node_table_t *self, const tsk_node_table_t *other,
     tsk_size_t num_rows, const tsk_id_t *row_indexes, tsk_flags_t options);
 
+/**
+@brief Deletes rows specified by a boolean array.
+
+@rst
+Deletes rows from this table and (optionally) return the mapping from IDs
+in the current table to the updated table. The rows to delete are specified
+by a boolean array, such that for each row ``j`` if ``delete_rows[j]`` is
+true (or more generally non-zero), then row ``j`` will be removed.
+
+If the ``id_map`` argument is non-null, this array will be updated
+to represent the mapping between IDs before and after row deletion.
+For row ``j``, ``id_map[j]`` will contain the new ID for row ``j``
+if it is retained, or :c:macro:`TSK_NULL` if the row has been removed.
+Thus, ``id_map`` must be an array of at least ``num_rows`` :c:type:`tsk_id_t`
+values.
+@endrst
+
+@param self A pointer to a tsk_node_table_t object.
+@param delete_rows Array of boolean flags describing whether a particular
+    row should be deleted or not. Must be at least ``num_rows`` long.
+@param options Bitwise option flags. Currently unused; should be
+    set to zero to ensure compatibility with later versions of tskit.
+@param id_map An array in which to store the mapping between new
+    and old IDs. If NULL, this will be ignored.
+@return Return 0 on success or a negative value on failure.
+*/
+int tsk_node_table_delete_rows(
+    tsk_node_table_t *self, bool *delete_rows, tsk_flags_t options, tsk_id_t *id_map);
+
 /**
 @brief Returns true if the data in the specified table is identical to the data
        in this table.