Support cudf-polars str.head and str.tail (rapidsai#19115)

Closes rapidsai#19031
Closes rapidsai#18995

Authors:
  - https://github.com/brandon-b-miller

Approvers:
  - Matthew Roeschke (https://github.com/mroeschke)

URL: rapidsai#19115

Author: brandon-b-miller

python/cudf_polars/cudf_polars/dsl/expressions/string.py

@@ -113,6 +113,7 @@ def _validate_input(self) -> None:
             StringFunction.Name.ConcatVertical,
             StringFunction.Name.Contains,
             StringFunction.Name.EndsWith,
+            StringFunction.Name.Head,
             StringFunction.Name.Lowercase,
             StringFunction.Name.Replace,
             StringFunction.Name.ReplaceMany,
@@ -123,6 +124,7 @@ def _validate_input(self) -> None:
             StringFunction.Name.StripCharsStart,
             StringFunction.Name.StripCharsEnd,
             StringFunction.Name.Uppercase,
+            StringFunction.Name.Tail,
         ):
             raise NotImplementedError(f"String function {self.name!r}")
         if self.name is StringFunction.Name.Contains:
@@ -283,6 +285,60 @@ def do_evaluate(
                 side = plc.strings.SideType.BOTH
             return Column(plc.strings.strip.strip(column.obj, side, chars.obj_scalar))
 
+        elif self.name is StringFunction.Name.Tail:
+            column = self.children[0].evaluate(df, context=context)
+
+            assert isinstance(self.children[1], Literal)
+            if self.children[1].value is None:
+                return Column(
+                    plc.Column.from_scalar(
+                        plc.Scalar.from_py(None, plc.DataType(plc.TypeId.STRING)),
+                        column.size,
+                    )
+                )
+            elif self.children[1].value == 0:
+                result = plc.Column.from_scalar(
+                    plc.Scalar.from_py("", plc.DataType(plc.TypeId.STRING)),
+                    column.size,
+                )
+                if column.obj.null_mask():
+                    result = result.with_mask(
+                        column.obj.null_mask(), column.obj.null_count()
+                    )
+                return Column(result)
+
+            else:
+                start = -(self.children[1].value)
+                end = 2**31 - 1
+                return Column(
+                    plc.strings.slice.slice_strings(
+                        column.obj,
+                        plc.Scalar.from_py(start, plc.DataType(plc.TypeId.INT32)),
+                        plc.Scalar.from_py(end, plc.DataType(plc.TypeId.INT32)),
+                        None,
+                    )
+                )
+        elif self.name is StringFunction.Name.Head:
+            column = self.children[0].evaluate(df, context=context)
+
+            assert isinstance(self.children[1], Literal)
+
+            end = self.children[1].value
+            if end is None:
+                return Column(
+                    plc.Column.from_scalar(
+                        plc.Scalar.from_py(None, plc.DataType(plc.TypeId.STRING)),
+                        column.size,
+                    )
+                )
+            return Column(
+                plc.strings.slice.slice_strings(
+                    column.obj,
+                    plc.Scalar.from_py(0, plc.DataType(plc.TypeId.INT32)),
+                    plc.Scalar.from_py(end, plc.DataType(plc.TypeId.INT32)),
+                )
+            )
+
         columns = [child.evaluate(df, context=context) for child in self.children]
         if self.name is StringFunction.Name.Lowercase:
             (column,) = columns
@@ -352,6 +408,7 @@ def do_evaluate(
             return Column(
                 plc.strings.replace.replace_multiple(column.obj, target.obj, repl.obj)
             )
+
         raise NotImplementedError(
             f"StringFunction {self.name}"
         )  # pragma: no cover; handled by init raising

python/cudf_polars/tests/expressions/test_stringfunction.py

@@ -470,3 +470,15 @@ def test_string_to_numeric_invalid(numeric_type):
 def test_string_join(ldf, ignore_nulls, delimiter):
     q = ldf.select(pl.col("a").str.join(delimiter, ignore_nulls=ignore_nulls))
     assert_gpu_result_equal(q)
+
+
+@pytest.mark.parametrize("tail", [1, 2, 999, -1, 0, None])
+def test_string_tail(ldf, tail):
+    q = ldf.select(pl.col("a").str.tail(tail))
+    assert_gpu_result_equal(q)
+
+
+@pytest.mark.parametrize("head", [1, 2, 999, -1, 0, None])
+def test_string_head(ldf, head):
+    q = ldf.select(pl.col("a").str.head(head))
+    assert_gpu_result_equal(q)