feat: Series.str.__getitem__

bigframes/core/compile/scalar_op_compiler.py

@@ -902,6 +902,24 @@ def array_to_string_op_impl(x: ibis_types.Value, op: ops.ArrayToStringOp):
     return typing.cast(ibis_types.ArrayValue, x).join(op.delimiter)
 
 
+@scalar_op_compiler.register_unary_op(ops.ArrayIndexOp, pass_op=True)
+def array_index_op_impl(x: ibis_types.Value, op: ops.ArrayIndexOp):
+    res = typing.cast(ibis_types.ArrayValue, x)[op.index]
+    if x.type().is_string():
+        return _null_or_value(res, res != ibis.literal(""))
+    else:
+        return res
+
+
+@scalar_op_compiler.register_unary_op(ops.ArraySliceOp, pass_op=True)
+def array_slice_op_impl(x: ibis_types.Value, op: ops.ArraySliceOp):
+    res = typing.cast(ibis_types.ArrayValue, x)[op.start : op.stop : op.step]
+    if x.type().is_string():
+        return _null_or_value(res, res != ibis.literal(""))
+    else:
+        return res
+
+
 # JSON Ops
 @scalar_op_compiler.register_binary_op(ops.JSONSet, pass_op=True)
 def json_set_op_impl(x: ibis_types.Value, y: ibis_types.Value, op: ops.JSONSet):
@@ -984,7 +1002,7 @@ def ne_op(
 
 
 def _null_or_value(value: ibis_types.Value, where_value: ibis_types.BooleanValue):
-    return ibis.where(
+    return ibis.ifelse(
         where_value,
         value,
         ibis.null(),

bigframes/operations/__init__.py

@@ -602,6 +602,40 @@ def output_type(self, *input_types):
         return dtypes.STRING_DTYPE
 
 
+@dataclasses.dataclass(frozen=True)
+class ArrayIndexOp(UnaryOp):
+    name: typing.ClassVar[str] = "array_index"
+    index: int
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if dtypes.is_string_like(input_type):
+            return dtypes.STRING_DTYPE
+        elif dtypes.is_array_like(input_type):
+            return dtypes.arrow_dtype_to_bigframes_dtype(
+                input_type.pyarrow_dtype.value_type
+            )
+        else:
+            raise TypeError("Input type must be an array or string-like type.")
+
+
+@dataclasses.dataclass(frozen=True)
+class ArraySliceOp(UnaryOp):
+    name: typing.ClassVar[str] = "array_slice"
+    start: int
+    stop: typing.Optional[int] = None
+    step: typing.Optional[int] = None
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if dtypes.is_string_like(input_type):
+            return dtypes.STRING_DTYPE
+        elif dtypes.is_array_like(input_type):
+            return input_type
+        else:
+            raise TypeError("Input type must be an array or string-like type.")
+
+
 ## JSON Ops
 @dataclasses.dataclass(frozen=True)
 class JSONExtract(UnaryOp):

bigframes/operations/strings.py

@@ -38,6 +38,33 @@
 class StringMethods(bigframes.operations.base.SeriesMethods, vendorstr.StringMethods):
     __doc__ = vendorstr.StringMethods.__doc__
 
+    def __getitem__(self, key: Union[int, slice]) -> series.Series:
+        if isinstance(key, int):
+            if key < 0:
+                raise NotImplementedError("Negative indexing is not supported.")
+            return self._apply_unary_op(ops.ArrayIndexOp(index=key))
+        elif isinstance(key, slice):
+            if key.step is not None and key.step != 1:
+                raise NotImplementedError(
+                    f"Only a step of 1 is allowed, got {key.step}"
+                )
+            if (key.start is not None and key.start < 0) or (
+                key.stop is not None and key.stop < 0
+            ):
+                raise NotImplementedError(
+                    "Slicing with negative numbers is not allowed."
+                )
+
+            return self._apply_unary_op(
+                ops.ArraySliceOp(
+                    start=key.start if key.start is not None else 0,
+                    stop=key.stop,
+                    step=key.step,
+                )
+            )
+        else:
+            raise ValueError(f"key must be an int or slice, got {type(key).__name__}")
+
     def find(
         self,
         sub: str,