add vectorised metadata, closes #1676

Tweaks for docs and an extra test

Author: petrelharp

python/tests/test_tables.py

@@ -1112,6 +1112,150 @@ def test_set_columns_metadata_schema(self):
         table.set_columns(**table2.asdict())
         assert table.metadata_schema == ms
 
+    def verify_metadata_vector(self, table, key, dtype, default_value=9999):
+        # this is just a hack for testing; the actual method
+        # does this more elegantly
+        has_default = default_value != 9999
+        if has_default:
+            md_vec = table.metadata_vector(
+                key, default_value=default_value, dtype=dtype
+            )
+        else:
+            md_vec = table.metadata_vector(key, dtype=dtype)
+        assert isinstance(md_vec, np.ndarray)
+        if dtype is not None:
+            assert md_vec.dtype == np.dtype(dtype)
+        assert len(md_vec) == table.num_rows
+        if not isinstance(key, list):
+            key = [key]
+        for x, row in zip(md_vec, table):
+            md = row.metadata
+            for k in key:
+                if k in md or not has_default:
+                    md = md[k]
+                else:
+                    md = default_value
+                    break
+            assert np.all(np.cast[dtype](md) == x)
+
+    def test_metadata_vector_errors(self):
+        table = self.table_class()
+        ms = tskit.MetadataSchema({"codec": "json"})
+        table.metadata_schema = ms
+        table.add_row(
+            **{
+                **self.input_data_for_add_row(),
+                "metadata": None,
+            }
+        )
+        with pytest.raises(KeyError):
+            _ = table.metadata_vector("x")
+        metadata_list = [
+            {"a": 4, "u": [1, 2]},
+            {},
+        ]
+        for md in metadata_list:
+            table.add_row(
+                **{
+                    **self.input_data_for_add_row(),
+                    "metadata": md,
+                }
+            )
+        with pytest.raises(KeyError):
+            _ = table.metadata_vector("x")
+
+        table.clear()
+        metadata_list = [
+            {"a": {"c": 5}, "u": [1, 2]},
+            {"a": {"b": 6}},
+        ]
+        for md in metadata_list:
+            table.add_row(
+                **{
+                    **self.input_data_for_add_row(),
+                    "metadata": md,
+                }
+            )
+        with pytest.raises(KeyError):
+            _ = table.metadata_vector(["a", "x"])
+
+    def test_metadata_vector_nodefault(self):
+        table = self.table_class()
+        ms = tskit.MetadataSchema({"codec": "json"})
+        table.metadata_schema = ms
+        metadata_list = [
+            {"abc": 4, "u": [1, 2]},
+            {"abc": 10, "u": [3, 4]},
+            {"abc": -3, "b": {"c": 1}, "u": [5, 6]},
+            {"abc": 1},
+        ]
+        for md in metadata_list:
+            table.add_row(
+                **{
+                    **self.input_data_for_add_row(),
+                    "metadata": md,
+                }
+            )
+        # first the totally obvious test
+        md_vec = table.metadata_vector("abc")
+        assert np.all(np.equal(md_vec, [d["abc"] for d in metadata_list]))
+        # now automated ones
+        for dtype in [None, "int", "float", "object"]:
+            self.verify_metadata_vector(
+                table, key="abc", dtype=dtype, default_value=9999
+            )
+            self.verify_metadata_vector(
+                table, key=["abc"], dtype=dtype, default_value=9999
+            )
+
+    def test_metadata_vector(self):
+        table = self.table_class()
+        ms = tskit.MetadataSchema({"codec": "json"})
+        table.metadata_schema = ms
+        metadata_list = [
+            {"abc": 4, "u": [1, 2]},
+            {"abc": 10, "u": [3, 4]},
+            {"abc": -3, "b": {"c": 1}, "u": [5, 6]},
+            {"b": {"c": 3.2}, "u": [7, 8]},
+            {"b": {"x": 8.2}},
+            {},
+            None,
+        ]
+        for md in metadata_list:
+            table.add_row(
+                **{
+                    **self.input_data_for_add_row(),
+                    "metadata": md,
+                }
+            )
+        # first the totally obvious test
+        md_vec = table.metadata_vector("abc", default_value=0)
+        assert np.all(
+            np.equal(
+                md_vec,
+                [
+                    d["abc"] if (d is not None and "abc" in d) else 0
+                    for d in metadata_list
+                ],
+            )
+        )
+
+        # now some automated ones
+        for dtype in [None, "int", "float", "object"]:
+            self.verify_metadata_vector(table, key="abc", dtype=dtype, default_value=-1)
+            self.verify_metadata_vector(
+                table, key=["abc"], dtype=dtype, default_value=-1
+            )
+            self.verify_metadata_vector(table, key=["x"], dtype=dtype, default_value=-1)
+            self.verify_metadata_vector(
+                table, key=["b", "c"], dtype=dtype, default_value=-1
+            )
+        self.verify_metadata_vector(table, key=["b"], dtype="object", default_value=-1)
+        self.verify_metadata_vector(table, key=["u"], dtype="int", default_value=[0, 0])
+        # and finally we should get rectangular arrays when it makes sense
+        md_vec = table.metadata_vector("u", default_value=[0, 0])
+        assert md_vec.shape == (table.num_rows, 2)
+
 
 class AssertEqualsMixin:
     def test_equal(self, table_5row, test_rows):

python/tskit/tables.py

@@ -31,7 +31,9 @@
 import json
 import numbers
 import warnings
+from collections.abc import Mapping
 from dataclasses import dataclass
+from functools import reduce
 from typing import Any
 from typing import Optional
 from typing import Union
@@ -48,6 +50,17 @@
 dataclass_options = {"frozen": True}
 
 
+# Needed for cases where `None` can be an appropriate kwarg value,
+# we override the meta so that it looks good in the docs.
+class NotSetMeta(type):
+    def __repr__(cls):
+        return "Not set"
+
+
+class NOTSET(metaclass=NotSetMeta):
+    pass
+
+
 @metadata.lazy_decode
 @dataclass(**dataclass_options)
 class IndividualTableRow(util.Dataclass):
@@ -585,6 +598,53 @@ def _update_metadata_schema_cache_from_ll(self) -> None:
             self.ll_table.metadata_schema
         )
 
+    def metadata_vector(self, key, *, dtype=None, default_value=NOTSET):
+        """
+        Returns a numpy array of metadata values obtained by extracting ``key``
+        from each metadata entry, and using ``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 str 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. The default behaviour is to raise
+            ``KeyError`` on missing entries.
+        """
+
+        if default_value == NOTSET:
+
+            def getter(d, k):
+                return d[k]
+
+        else:
+
+            def getter(d, k):
+                return (
+                    d.get(k, default_value) if isinstance(d, Mapping) else default_value
+                )
+
+        if isinstance(key, list):
+            out = np.array(
+                [
+                    reduce(
+                        getter,
+                        key,
+                        row.metadata,
+                    )
+                    for row in self
+                ],
+                dtype=dtype,
+            )
+        else:
+            out = np.array(
+                [getter(row.metadata, key) for row in self],
+                dtype=dtype,
+            )
+        return out
+
 
 class IndividualTable(BaseTable, MetadataMixin):
     """