Add backend support for str dtypes

Author: michaelosthege

mcbackend/backends/clickhouse.py

@@ -30,6 +30,7 @@
     "float32": "Float32",
     "float64": "Float64",
     "bool": "UInt8",
+    "str": "String",
 }
 
 
@@ -155,6 +156,7 @@ def __init__(
         self._client = client
         # The following attributes belong to the batched insert mechanism.
         # Inserting in batches is much faster than inserting single rows.
+        self._str_cols = set()
         self._insert_query: str = ""
         self._insert_queue: List[Dict[str, Any]] = []
         self._last_insert = time.time()
@@ -166,12 +168,22 @@ def append(
         self, draw: Dict[str, numpy.ndarray], stats: Optional[Dict[str, numpy.ndarray]] = None
     ):
         stat = {f"__stat_{sname}": svals for sname, svals in (stats or {}).items()}
-        params: Dict[str, Any] = {"_draw_idx": self._draw_idx, **draw, **stat}
-        self._draw_idx += 1
+        params: Dict[str, numpy.ndarray] = {**draw, **stat}
+
+        # On first append create a query to be used for the batched insert
         if not self._insert_query:
             names = "`,`".join(params.keys())
-            self._insert_query = f"INSERT INTO {self.cid} (`{names}`) VALUES"
+            self._insert_query = f"INSERT INTO {self.cid} (`_draw_idx`,`{names}`) VALUES"
+            self._str_cols = {k for k, v in params.items() if "str" in numpy.asarray(v).dtype.name}
+
+        # Convert str ndarrays to lists
+        for col in self._str_cols:
+            params[col] = params[col].tolist()
+
+        # Queue up for insertion
+        params["_draw_idx"] = self._draw_idx
         self._insert_queue.append(params)
+        self._draw_idx += 1
 
         if (
             len(self._insert_queue) >= self._insert_every
@@ -235,7 +247,9 @@ def _get_rows(
             return numpy.array([], dtype=object)
 
         # The unpacking must also account for non-rigid shapes
-        if is_rigid(nshape):
+        # and str-dtyped empty arrays default to fixed length 1 strings.
+        # The [None] list is slower, but more flexible in this regard.
+        if is_rigid(nshape) and dtype != "str":
             assert nshape is not None
             buffer = numpy.empty((draws, *nshape), dtype)
         else:

mcbackend/backends/numpy.py

@@ -65,7 +65,7 @@ def __init__(self, cmeta: ChainMeta, rmeta: RunMeta, *, preallocate: int) -> Non
             (self._stats, self._stat_is_rigid, rmeta.sample_stats),
         ]:
             for var in variables:
-                rigid = is_rigid(var.shape) and not var.undefined_ndim
+                rigid = is_rigid(var.shape) and not var.undefined_ndim and var.dtype != "str"
                 rigid_dict[var.name] = rigid
                 if preallocate > 0 and rigid:
                     reserve = (preallocate, *var.shape)
@@ -88,13 +88,19 @@ def __len__(self) -> int:
         return self._draw_idx
 
     def get_draws(self, var_name: str, slc: slice = slice(None)) -> numpy.ndarray:
-        return self._samples[var_name][: self._draw_idx][slc]
+        data = self._samples[var_name][: self._draw_idx][slc]
+        if self.variables[var_name].dtype == "str":
+            return data.astype(str)
+        return data
 
     def get_draws_at(self, idx: int, var_names: Sequence[str]) -> Dict[str, numpy.ndarray]:
         return {vn: numpy.asarray(self._samples[vn][idx]) for vn in var_names}
 
     def get_stats(self, stat_name: str, slc: slice = slice(None)) -> numpy.ndarray:
-        return self._stats[stat_name][: self._draw_idx][slc]
+        data = self._stats[stat_name][: self._draw_idx][slc]
+        if self.sample_stats[stat_name].dtype == "str":
+            return data.astype(str)
+        return data
 
     def get_stats_at(self, idx: int, stat_names: Sequence[str]) -> Dict[str, numpy.ndarray]:
         return {sn: numpy.asarray(self._stats[sn][idx]) for sn in stat_names}

mcbackend/test_utils.py

@@ -29,6 +29,8 @@ def make_runmeta(*, flexibility: bool = False, **kwargs) -> RunMeta:
             Variable("accepted", "bool", list((3,)), dims=["sampler"]),
             # But some stats may refer to the iteration.
             Variable("logp", "float64", []),
+            # String dtypes may be used for more complex information
+            Variable("message", "str"),
         ],
         data=[
             DataVariable(
@@ -60,8 +62,16 @@ def make_draw(variables: Sequence[Variable]):
         )
         if "float" in var.dtype:
             draw[var.name] = numpy.random.normal(size=dshape).astype(var.dtype)
+        elif var.dtype == "str":
+            alphabet = tuple("abcdef#+*/'")
+            words = [
+                "".join(numpy.random.choice(alphabet, size=numpy.random.randint(3, 10)))
+                for _ in range(int(numpy.prod(dshape)))
+            ]
+            draw[var.name] = numpy.array(words, dtype=var.dtype).reshape(dshape)
         else:
             draw[var.name] = numpy.random.randint(low=0, high=100, size=dshape).astype(var.dtype)
+        assert draw[var.name].shape == dshape
     return draw
 
 
@@ -149,7 +159,7 @@ def test__append_get_with_changelings(self, with_stats):
             expected = [draw[var.name] for draw in draws]
             actual = chain.get_draws(var.name)
             assert isinstance(actual, numpy.ndarray)
-            if var.name == "changeling":
+            if not is_rigid(var.shape) or var.dtype == "str":
                 # Non-ridid variables are returned as object-arrays.
                 assert actual.shape == (len(expected),)
                 assert actual.dtype == object
@@ -166,9 +176,13 @@ def test__append_get_with_changelings(self, with_stats):
                 expected = [stat[var.name] for stat in stats]
                 actual = chain.get_stats(var.name)
                 assert isinstance(actual, numpy.ndarray)
-                if is_rigid(var.shape):
+                if var.dtype == "str":
                     assert tuple(actual.shape) == tuple(numpy.shape(expected))
-                    assert actual.dtype == var.dtype
+                    # String dtypes have strange names
+                    assert "str" in actual.dtype.name
+                elif is_rigid(var.shape):
+                    assert tuple(actual.shape) == tuple(numpy.shape(expected))
+                    assert actual.dtype.name == var.dtype
                     numpy.testing.assert_array_equal(actual, expected)
                 else:
                     # Non-ridid variables are returned as object-arrays.
@@ -200,7 +214,7 @@ def test__get_slicing(self, slc: slice):
         # "A" are just numbers to make diagnosis easier.
         # "B" are dynamically shaped to cover the edge cases.
         rmeta = RunMeta(
-            variables=[Variable("A", "uint8")],
+            variables=[Variable("A", "uint8"), Variable("M", "str", [2, 3])],
             sample_stats=[Variable("B", "uint8", [2, 0])],
             data=[],
         )
@@ -209,7 +223,7 @@ def test__get_slicing(self, slc: slice):
 
         # Generate draws and add them to the chain
         N = 20
-        draws = [dict(A=n) for n in range(N)]
+        draws = [make_draw(rmeta.variables) for n in range(N)]
         stats = [make_draw(rmeta.sample_stats) for n in range(N)]
         for d, s in zip(draws, stats):
             chain.append(d, s)
@@ -218,12 +232,25 @@ def test__get_slicing(self, slc: slice):
         # slc=None in this test means "don't pass it".
         # The implementations should default to slc=slice(None, None, None).
         kwargs = dict(slc=slc) if slc is not None else {}
-        act_draws = chain.get_draws("A", **kwargs)
+        act_draws_A = chain.get_draws("A", **kwargs)
+        act_draws_M = chain.get_draws("M", **kwargs)
         act_stats = chain.get_stats("B", **kwargs)
-        expected_draws = [d["A"] for d in draws][slc or slice(None, None, None)]
+        expected_draws_A = [d["A"] for d in draws][slc or slice(None, None, None)]
+        expected_draws_M = [d["M"] for d in draws][slc or slice(None, None, None)]
         expected_stats = [s["B"] for s in stats][slc or slice(None, None, None)]
+
         # Variable "A" has a rigid shape
-        numpy.testing.assert_array_equal(act_draws, expected_draws)
+        if expected_draws_A:
+            numpy.testing.assert_array_equal(act_draws_A, expected_draws_A)
+        else:
+            assert len(act_draws_A) == 0
+
+        # Variable "M" is a string matrix
+        if expected_draws_M:
+            numpy.testing.assert_array_equal(act_draws_M, expected_draws_M)
+        else:
+            assert len(act_draws_M) == 0
+
         # Stat "B" is dynamically shaped, which means we're dealing with
         # dtype=object arrays. These must be checked elementwise.
         assert len(act_stats) == len(expected_stats)
@@ -256,6 +283,7 @@ def test__to_inferencedata(self):
             sample_stats=[
                 Variable("tune", "bool"),
                 Variable("sampler_0__logp", "float32"),
+                Variable("warning", "str"),
             ],
         )
         run = self.backend.init_run(rmeta)