more tests

Author: petrelharp

python/tests/test_tables.py

@@ -1112,15 +1112,34 @@ 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):
+        md_vec = table.metadata_vector(key, default_value=default_value, 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:
+                    md = md[k]
+                else:
+                    md = default_value
+                    break
+            assert np.all(np.cast[dtype](md) == x)
+
     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}},
+            {"a": 4, "u": [1, 2]},
+            {"a": 10, "u": [3, 4]},
+            {"a": -3, "b": {"c": 1}, "u": [5, 6]},
+            {"b": {"c": 3.2}, "u": [7, 8]},
+            {"b": {"x": 8.2}},
             {},
         ]
         for md in metadata_list:
@@ -1130,12 +1149,18 @@ def test_metadata_vector(self):
                     "metadata": md,
                 }
             )
-        default_value = -1
-        for key in ["a", ["b", "c"]]:
-            assert np.equal(
-                [md.get(key, default_value) for md in metadata_list],
-                table.metadata_vector(key, default_value=default_value),
+        for dtype in [None, "int", "float", "object"]:
+            self.verify_metadata_vector(table, key="a", dtype=dtype, default_value=-1)
+            self.verify_metadata_vector(table, key=["a"], 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])
+        md_vec = table.metadata_vector("u", default_value=[0, 0], dtype="int")
+        # and finally we should get rectangular arrays when it makes sense
+        assert md_vec.shape == (table.num_rows, 2)
 
 
 class AssertEqualsMixin:

python/tskit/tables.py

@@ -27,6 +27,7 @@
 import collections.abc
 import dataclasses
 import datetime
+import functools
 import itertools
 import json
 import numbers
@@ -599,10 +600,27 @@ def metadata_vector(self, key, *, dtype=None, default_value=None):
             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_value) for row in self],
-            dtype=dtype,
-        )
+
+        if isinstance(key, list):
+            out = np.array(
+                [
+                    functools.reduce(
+                        lambda d, k: d.get(k, default_value)
+                        if isinstance(d, collections.abc.Mapping)
+                        else default_value,
+                        key,
+                        row.metadata,
+                    )
+                    for row in self
+                ],
+                dtype=dtype,
+            )
+        else:
+            out = np.array(
+                [row.metadata.get(key, default_value) for row in self],
+                dtype=dtype,
+            )
+        return out
 
 
 class IndividualTable(BaseTable, MetadataMixin):