add vectorised metadata, closes #1676

Author: petrelharp

python/tests/test_tables.py

@@ -1112,6 +1112,31 @@ def test_set_columns_metadata_schema(self):
         table.set_columns(**table2.asdict())
         assert table.metadata_schema == ms
 
+    def test_metadata_vector(self):
+        table = self.table_class()
+        ms = tskit.MetadataSchema({"codec": "json"})
+        table.metadata_schema = ms
+        metadata_list = [
+            {"a": 4},
+            {"a": 10},
+            {"a": -3, "b": {"c": 1}},
+            {"b": {"c": 3.2}},
+            {},
+        ]
+        for md in metadata_list:
+            table.add_row(
+                **{
+                    **self.input_data_for_add_row(),
+                    "metadata": md,
+                }
+            )
+        default_value = -1
+        for key in ["a", ["b", "c"]]:
+            assert np.equal(
+                [md.get(key, default=default_value) for md in metadata_list],
+                table.metadata_vector(key, default_value=default_value),
+            )
+
 
 class AssertEqualsMixin:
     def test_equal(self, table_5row, test_rows):

python/tskit/tables.py

@@ -585,6 +585,25 @@ def _update_metadata_schema_cache_from_ll(self) -> None:
             self.ll_table.metadata_schema
         )
 
+    def metadata_vector(self, key, *, dtype=None, default_value=None):
+        """
+        Returns a numpy array of metadata values obtained by extracting ``key``
+        from each metadata entry, and inserting ``default_value`` if the key is
+        not present. ``key`` may be a list, in which case nested values are returned.
+        For instance, ``key = ["a", "x"]`` will return an array of
+        ``row.metadata["a"]["x"]`` values, iterated over rows in this table.
+
+        :param str key: The name, or a list of names, of metadata entries.
+        :param dtype dtype: The dtype of the result (can usually be omitted).
+        :param object default_value: The value to be inserted if the metadata key
+            is not present. Note that for numeric columns, a default value of None
+            will result in a non-numeric array.
+        """
+        return np.fromiter(
+            [row.metadata.get(key, default=default_value) for row in self],
+            dtype=dtype,
+        )
+
 
 class IndividualTable(BaseTable, MetadataMixin):
     """