Merge pull request #95 from martindurant/str_sugar

Add mul/add convenience to .str

Author: martindurant

src/akimbo/apply_tree.py

@@ -71,11 +71,12 @@ def dec(
     match: Callable[[ak.contents.Content], bool] = leaf,
     outtype: Callable[[ak.contents.Content], ak.contents.Content] | None = None,
     inmode: Literal["arrow", "numpy", "ak", "other"] = "ak",
+    match_kwargs=None,
 ):
     """Make a nested/ragged version of an operation to apply throughout a tree"""
 
     @functools.wraps(func)
-    def f(arr, *args, where=None, match_kwargs=None, **kwargs):
+    def f(arr, *args, where=None, **kwargs):
         others = []
         if args:
             sig = list(inspect.signature(func).parameters)[1:]

src/akimbo/dask.py

@@ -57,7 +57,7 @@ def to_dask_awkward(self):
         )
 
     def __getattr__(self, item):
-        if self.subaccessor:
+        if self.subaccessor and isinstance(item, str):
             item = getattr(self.subaccessors[self.subaccessor], item)
         elif isinstance(item, str) and item in self.subaccessors:
             return DaskAwkwardAccessor(

src/akimbo/mixin.py

@@ -176,6 +176,10 @@ def unexplode(self, *cols: tuple[str, ...], outname="grouped") -> ak.Array:
     @classmethod
     def _create_op(cls, op):
         def run(self, *args, **kwargs):
+            if self.subaccessor:
+                # defer
+                op2 = op(self.subaccessors[self.subaccessor](), None)
+                return self.__getattr__(op2)(*args, **kwargs)
             args = [
                 to_ak_layout(_) if isinstance(_, (str, int, float, np.number)) else _
                 for _ in args

src/akimbo/polars.py

@@ -50,7 +50,7 @@ def __getattr__(self, item: str, **flags) -> callable:
             )
 
         def select(*inargs, subaccessor=self.subaccessor, where=None, **kwargs):
-            if subaccessor:
+            if subaccessor and isinstance(item, str):
                 func0 = getattr(self.subaccessors[subaccessor](), item)
             elif callable(item):
                 func0 = item
@@ -100,18 +100,6 @@ def f(batch):
                 raise NotImplementedError
             else:
                 obj = self._obj
-            # def map_batches(
-            #     self,
-            #     function: Callable[[DataFrame], DataFrame],
-            #     *,
-            #     predicate_pushdown: bool = True,
-            #     projection_pushdown: bool = True,
-            #     slice_pushdown: bool = True,
-            #     no_optimizations: bool = False,
-            #     schema: None | SchemaDict = None,
-            #     validate_output_schema: bool = True,
-            #     streamable: bool = False,
-            # ) -> LazyFrame:
             arrow_type = polars_to_arrow_schema(obj.collect_schema())
             arr = pa.table([[]] * len(arrow_type), schema=arrow_type)
             out1 = f(pl.from_arrow(arr))

src/akimbo/ray.py

@@ -34,7 +34,7 @@ def __getattr__(self, item: str) -> callable:
             return RayAccessor(self._obj, subaccessor=item, behavior=self._behavior)
 
         def select(*inargs, subaccessor=self.subaccessor, where=None, **kwargs):
-            if subaccessor:
+            if subaccessor and isinstance(item, str):
                 func0 = getattr(self.subaccessors[subaccessor](), item)
             elif callable(item):
                 func0 = item

src/akimbo/spark.py

@@ -38,7 +38,7 @@ def __getattr__(self, item: str) -> sdf:
             return SparkAccessor(self._obj, subaccessor=item, behavior=self._behavior)
 
         def select(*inargs, subaccessor=self.subaccessor, where=None, **kwargs):
-            if subaccessor:
+            if subaccessor and isinstance(item, str):
                 func0 = getattr(self.subaccessors[subaccessor](), item)
             elif callable(item):
                 func0 = item

src/akimbo/strings.py

@@ -3,6 +3,7 @@
 import functools
 
 import awkward as ak
+import pyarrow as pa
 import pyarrow.compute as pc
 
 from akimbo.apply_tree import dec
@@ -59,6 +60,16 @@ def strptime(*args, format="%FT%T", unit="us", error_is_null=True, **kw):
     return out
 
 
+def repeat(arr, count):
+    return pc.binary_repeat(arr, count)
+
+
+def concat(arr, arr2, sep=""):
+    return pc.binary_join_element_wise(
+        arr.cast(pa.string()), arr2.cast(pa.string()), sep
+    )
+
+
 class StringAccessor:
     """String operations on nested/var-length data"""
 
@@ -92,6 +103,14 @@ def __getattr__(self, attr: str) -> callable:
         return getattr(ak.str, attr)
 
     strptime = staticmethod(dec(strptime, match=match_string, inmode="arrow"))
+    repeat = staticmethod(dec(repeat, match=match_string, inmode="arrow"))
+    join_el = staticmethod(dec(concat, match=match_string, inmode="arrow"))
+
+    def __add__(self, *_):
+        return dec(concat, match=match_string, inmode="arrow")
+
+    def __mul__(self, *_):
+        return dec(repeat, match=match_string, inmode="arrow")
 
     def __dir__(self) -> list[str]:
         return sorted(methods + ["strptime", "encode", "decode"])

tests/test_polars.py

@@ -75,6 +75,19 @@ def test_binary():
     assert s3.to_list() == [[0, 0, 0], [], [0, 0]]
 
 
+def test_str_sugar():
+    s = pl.Series([["hay", "hi", "hola"], [], ["bye", "yo"]])
+    out = s.ak.str.repeat(3)
+    out2 = s.ak.str * 3
+    expected = [["hayhayhay", "hihihi", "holaholahola"], [], ["byebyebye", "yoyoyo"]]
+    assert out.ak.tolist() == out2.ak.tolist() == expected
+
+    out = s.ak.str.join_el(s)
+    out2 = s.ak.str + s
+    expected = [["hayhay", "hihi", "holahola"], [], ["byebye", "yoyo"]]
+    assert out.ak.tolist() == out2.ak.tolist() == expected
+
+
 def test_unexplode():
     df = pl.DataFrame(
         {

tests/test_ray.py

@@ -150,3 +150,16 @@ def f2(data):
     out = df.ak.apply(f2)
     result = out.ak.to_output()
     assert result.ak.tolist() == [6, None] * 100
+
+
+def test_str_sugar(df):
+    s = df.ak["y"]
+    out = s.ak.str.repeat(3)
+    out2 = s.ak.str * 3
+    expected = [["heyheyhey", None], ["hihihi", "hohoho"]] * 100
+    assert out.ak.to_output().tolist() == out2.ak.to_output().tolist() == expected
+
+    out = s.ak.str.join_el(s)
+    out2 = s.ak.str + s
+    expected = [["heyhey", None], ["hihi", "hoho"]] * 100
+    assert out.ak.to_output().tolist() == out2.ak.to_output().tolist() == expected

tests/test_str.py

@@ -40,3 +40,16 @@ def test_split():
     s = pd.Series([b"hello world", b"oio", b""])
     s2 = s.ak.str.split_whitespace()
     assert s2.tolist() == [[b"hello", b"world"], [b"oio"], [b""]]
+
+
+def test_str_sugar():
+    s = pd.Series([["hay", "hi", "hola"], [], ["bye", "yo"]])
+    out = s.ak.str.repeat(3)
+    out2 = s.ak.str * 3
+    expected = [["hayhayhay", "hihihi", "holaholahola"], [], ["byebyebye", "yoyoyo"]]
+    assert out.ak.tolist() == out2.ak.tolist() == expected
+
+    out = s.ak.str.join_el(s)
+    out2 = s.ak.str + s
+    expected = [["hayhay", "hihi", "holahola"], [], ["byebye", "yoyo"]]
+    assert out.ak.tolist() == out2.ak.tolist() == expected