From f43f1acfc7e050218f78d1774638a91189cbb47b Mon Sep 17 00:00:00 2001
From: Martin Durant <martin.durant@alumni.utoronto.ca>
Date: Mon, 11 Nov 2024 09:47:46 -0500
Subject: [PATCH] string exception

---
 src/akimbo/apply_tree.py | 5 +++++
 src/akimbo/strings.py    | 5 +----
 src/akimbo/utils.py      | 6 ++++++
 tests/test_pandas.py     | 1 +
 4 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/akimbo/apply_tree.py b/src/akimbo/apply_tree.py
index ab7d7a5..a6531a7 100644
--- a/src/akimbo/apply_tree.py
+++ b/src/akimbo/apply_tree.py
@@ -35,6 +35,8 @@ def run_with_transform(
     **kw,
 ) -> ak.Array:
     def func(layout, **kwargs):
+        from akimbo.utils import match_string
+
         if not isinstance(layout, tuple):
             layout = (layout,)
         if all(match(lay, **(match_kwargs or {})) for lay in layout):
@@ -57,6 +59,9 @@ def func(layout, **kwargs):
                 return out.layout
             else:
                 return out
+        if match_string(*layout):
+            # non-string op may fail to descend into string
+            return layout[0]
 
     return ak.transform(func, arr, *others, allow_records=True)
 
diff --git a/src/akimbo/strings.py b/src/akimbo/strings.py
index 0cf167b..1460154 100644
--- a/src/akimbo/strings.py
+++ b/src/akimbo/strings.py
@@ -8,10 +8,7 @@
 
 from akimbo.apply_tree import dec
 from akimbo.mixin import Accessor
-
-
-def match_string(*layout):
-    return layout[0].is_list and layout[0].parameter("__array__") == "string"
+from akimbo.utils import match_string
 
 
 def _encode(layout):
diff --git a/src/akimbo/utils.py b/src/akimbo/utils.py
index f0037cf..a5f1b89 100644
--- a/src/akimbo/utils.py
+++ b/src/akimbo/utils.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import awkward as ak
 
 
@@ -34,3 +36,7 @@ def to_ak_layout(ar):
         return ar
     else:
         return ak.Array(ak.to_layout(ar))
+
+
+def match_string(*layout):
+    return layout[0].is_list and layout[0].parameter("__array__") == "string"
diff --git a/tests/test_pandas.py b/tests/test_pandas.py
index 3e944bd..2580e59 100644
--- a/tests/test_pandas.py
+++ b/tests/test_pandas.py
@@ -99,6 +99,7 @@ def test_mixed_ufunc():
     df2 = df.ak == df.ak
     expected = [[True, True, True], [True, True], [True]]
     assert df2["b"].tolist() == expected
+    assert df2["a"].tolist() == df["a"].tolist()
 
 
 def test_to_autoarrow():