Add initial test, fix ragged lengths

Author: benjeffery

c/tests/test_tables.c

@@ -4049,6 +4049,87 @@ test_individual_table(void)
     free(metadata_offset);
 }
 
+static void
+test_individual_table_takeset(void)
+{
+    int ret = 0;
+    tsk_id_t ret_id;
+    tsk_individual_table_t source_table, table;
+    tsk_size_t num_rows = 100;
+    tsk_id_t j;
+    tsk_size_t k;
+    tsk_flags_t *flags;
+    double *location;
+    tsk_id_t *parents;
+    char *metadata;
+    tsk_size_t *metadata_offset;
+    tsk_size_t *parents_offset;
+    tsk_size_t *location_offset;
+    tsk_size_t spatial_dimension = 2;
+    tsk_size_t num_parents = 3;
+    const char *test_metadata = "test";
+    tsk_size_t test_metadata_length = 4;
+    double test_location[spatial_dimension];
+    tsk_id_t test_parents[num_parents];
+
+    /* Make a table to copy from */
+    ret = tsk_individual_table_init(&source_table, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    for (k = 0; k < spatial_dimension; k++) {
+        test_location[k] = (double) k;
+    }
+    for (k = 0; k < num_parents; k++) {
+        test_parents[k] = (tsk_id_t) k + 42;
+    }
+    for (j = 0; j < (tsk_id_t) num_rows; j++) {
+        ret_id = tsk_individual_table_add_row(&source_table, (tsk_flags_t) j,
+            test_location, spatial_dimension, test_parents, num_parents, test_metadata,
+            test_metadata_length);
+        CU_ASSERT_EQUAL_FATAL(ret_id, j);
+    }
+
+    /* Prepare arrays to be taken */
+    flags = tsk_malloc(num_rows * sizeof(tsk_flags_t));
+    CU_ASSERT_FATAL(flags != NULL);
+    tsk_memcpy(flags, source_table.flags, num_rows * sizeof(tsk_flags_t));
+    location = tsk_malloc(spatial_dimension * num_rows * sizeof(double));
+    CU_ASSERT_FATAL(location != NULL);
+    tsk_memcpy(
+        location, source_table.location, spatial_dimension * num_rows * sizeof(double));
+    location_offset = tsk_malloc((num_rows + 1) * sizeof(tsk_size_t));
+    CU_ASSERT_FATAL(location_offset != NULL);
+    tsk_memcpy(location_offset, source_table.location_offset,
+        (num_rows + 1) * sizeof(tsk_size_t));
+    parents = tsk_malloc(num_parents * num_rows * sizeof(tsk_id_t));
+    CU_ASSERT_FATAL(parents != NULL);
+    tsk_memcpy(parents, source_table.parents, num_parents * num_rows * sizeof(tsk_id_t));
+    parents_offset = tsk_malloc((num_rows + 1) * sizeof(tsk_size_t));
+    CU_ASSERT_FATAL(parents_offset != NULL);
+    tsk_memcpy(parents_offset, source_table.parents_offset,
+        (num_rows + 1) * sizeof(tsk_size_t));
+    metadata = tsk_malloc(num_rows * test_metadata_length * sizeof(char));
+    CU_ASSERT_FATAL(metadata != NULL);
+    tsk_memcpy(
+        metadata, source_table.metadata, num_rows * test_metadata_length * sizeof(char));
+    metadata_offset = tsk_malloc((num_rows + 1) * sizeof(tsk_size_t));
+    CU_ASSERT_FATAL(metadata_offset != NULL);
+    tsk_memcpy(metadata_offset, source_table.metadata_offset,
+        (num_rows + 1) * sizeof(tsk_size_t));
+
+    ret = tsk_individual_table_init(&table, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = tsk_individual_table_takeset_columns(&table, num_rows, flags, location,
+        location_offset, parents, parents_offset, metadata, metadata_offset);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_individual_table_equals(&source_table, &table, 0));
+
+    ret = tsk_individual_table_free(&table);
+    CU_ASSERT_EQUAL(ret, 0);
+    ret = tsk_individual_table_free(&source_table);
+    CU_ASSERT_EQUAL(ret, 0);
+}
+
 static void
 test_individual_table_update_row(void)
 {
@@ -9235,6 +9316,7 @@ main(int argc, char **argv)
         { "test_migration_table", test_migration_table },
         { "test_migration_table_update_row", test_migration_table_update_row },
         { "test_individual_table", test_individual_table },
+        { "test_individual_table_takeset", test_individual_table_takeset },
         { "test_individual_table_update_row", test_individual_table_update_row },
         { "test_population_table", test_population_table },
         { "test_population_table_update_row", test_population_table_update_row },

c/tskit/tables.c

@@ -643,7 +643,7 @@ alloc_empty_ragged_column(tsk_size_t num_rows, void **data_col, tsk_size_t **off
 
 static int
 takeset_ragged_column(tsk_size_t num_rows, void *data, tsk_size_t *offset,
-    void **data_dest, tsk_size_t **offset_dest)
+    void **data_dest, tsk_size_t **offset_dest, tsk_size_t *length_dest)
 {
     int ret = 0;
 
@@ -664,6 +664,7 @@ takeset_ragged_column(tsk_size_t num_rows, void *data, tsk_size_t *offset,
         *data_dest = data;
         *offset_dest = offset;
     }
+    *length_dest = (*offset_dest)[num_rows];
 out:
     return ret;
 }
@@ -1087,17 +1088,17 @@ tsk_individual_table_takeset_columns(tsk_individual_table_t *self, tsk_size_t nu
     }
 
     ret = takeset_ragged_column(num_rows, location, location_offset,
-        (void *) &self->location, &self->location_offset);
+        (void *) &self->location, &self->location_offset, &self->location_length);
     if (ret != 0) {
         goto out;
     }
     ret = takeset_ragged_column(num_rows, parents, parents_offset,
-        (void *) &self->parents, &self->parents_offset);
+        (void *) &self->parents, &self->parents_offset, &self->parents_length);
     if (ret != 0) {
         goto out;
     }
     ret = takeset_ragged_column(num_rows, metadata, metadata_offset,
-        (void *) &self->metadata, &self->metadata_offset);
+        (void *) &self->metadata, &self->metadata_offset, &self->metadata_length);
     if (ret != 0) {
         goto out;
     }
@@ -1837,7 +1838,7 @@ tsk_node_table_takeset_columns(tsk_node_table_t *self, tsk_size_t num_rows,
         goto out;
     }
     ret = takeset_ragged_column(num_rows, metadata, metadata_offset,
-        (void *) &self->metadata, &self->metadata_offset);
+        (void *) &self->metadata, &self->metadata_offset, &self->metadata_length);
     if (ret != 0) {
         goto out;
     }