fix: fill flattened pass-through strings (#629)

Close #609.

Signed-off-by: Henry Schreiner <[email protected]>

Author: henryiii

pyproject.toml

@@ -92,6 +92,7 @@ dask = [
 test = [
     { include-group = "test-core" },
     { include-group = "fit" },
+    "awkward>=2.0.7",
 ]
 test-core = [
     "pytest >=8",

src/hist/interop.py

@@ -67,15 +67,22 @@ def destructure(obj: Any) -> dict[str, Any] | None:
     raise TypeError(f"No histogram module found for {obj!r}")
 
 
-def broadcast_and_flatten(args: Sequence[Any]) -> tuple[np.typing.NDArray[Any], ...]:
+def broadcast_and_flatten(
+    args: Sequence[Any],
+) -> tuple[str | np.typing.NDArray[Any], ...]:
     """
     Convert the given histogram-module arrays into a set of consistent 1D NumPy arrays
     for histogram filling. For NumPy this entails broadcasting and flattening.
+
+    This skips passing strings to the backend, they are left inplace.
     """
-    for module in find_histogram_modules(*args):
-        result = module.broadcast_and_flatten(args)
+
+    non_strings = [x for x in args if not isinstance(x, str)]
+    for module in find_histogram_modules(*non_strings):
+        result = module.broadcast_and_flatten(non_strings)
         if result is not NotImplemented:
-            return result
+            it = iter(result)
+            return tuple(next(it) if not isinstance(x, str) else x for x in args)
 
     raise TypeError(f"No histogram module found for {args!r}")
 

tests/test_interop.py

@@ -383,3 +383,17 @@ def test_named_keyword(unnamed_hist):
         axis.Regular(2, 0, 1, name="z"),
     ).fill_flattened(x=x, y=y, z=z)
     assert np.allclose(h.values(), expected)
+
+
+def test_string_fill_flattened():
+    ark = np.array([1, 2, 3, 4, 5])
+    h = hist.new.Reg(10, 0, 10, name="x").StrCat([], growth=True, name="cat").Weight()
+    h.fill_flattened(x=ark, cat="A")
+
+
+def test_ak_fill_flattened():
+    ak = pytest.importorskip("awkward")
+
+    ark = ak.Array([[1, 2, 3], [4, 5, 6], [9]])
+    h = hist.new.Reg(10, 0, 10, name="x").StrCat([], growth=True, name="cat").Weight()
+    h.fill_flattened(x=ark, cat="A")