feat: Series.str.__getitem__ (#897)

- [X] Make sure to open an issue as internal issue: 358459166 - [X] Ensure the tests and linter pass - [X] Code coverage does not decrease (if any source code was changed) - [X] Appropriate docs were updated (if necessary) Fixes internal issue: 358459166 🦕

Author: chelsea-lin

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,