fixup str extract method

Author: TrevorBergeron

bigframes/operations/strings.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import re
-from typing import Generic, Literal, Optional, TypeVar, Union
+from typing import Generic, Hashable, Literal, Optional, TypeVar, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.strings.accessor as vendorstr
@@ -209,23 +209,19 @@ def extract(self, pat: str, flags: int = 0) -> df.DataFrame:
         if compiled.groups == 0:
             raise ValueError("No capture groups in 'pat'")
 
-        results: list[str] = []
-        block = self._data._block
+        results: dict[Hashable, T] = {}
         for i in range(compiled.groups):
             labels = [
                 label
                 for label, groupn in compiled.groupindex.items()
                 if i + 1 == groupn
             ]
-            label = labels[0] if labels else str(i)
-            block, id = block.apply_unary_op(
-                self._data._value_column,
+            label = labels[0] if labels else i
+            result = self._data._apply_unary_op(
                 ops.StrExtractOp(pat=pat, n=i + 1),
-                result_label=label,
             )
-            results.append(id)
-        block = block.select_columns(results)
-        return df.DataFrame(block)
+            results[label] = series.Series(result)
+        return df.DataFrame(results)
 
     def replace(
         self,
@@ -287,10 +283,10 @@ def split(
             )
         return self._data._apply_unary_op(ops.StringSplitOp(pat=pat))
 
-    def zfill(self, width: int) -> series.Series:
+    def zfill(self, width: int) -> T:
         return self._data._apply_unary_op(ops.ZfillOp(width=width))
 
-    def center(self, width: int, fillchar: str = " ") -> series.Series:
+    def center(self, width: int, fillchar: str = " ") -> T:
         return self._data._apply_unary_op(
             ops.StrPadOp(length=width, fillchar=fillchar, side="both")
         )

scripts/publish_api_coverage.py

@@ -30,6 +30,7 @@
 import bigframes.core.groupby
 import bigframes.core.window
 import bigframes.operations.datetimes
+import bigframes.operations.strings
 import bigframes.pandas as bpd
 
 REPO_ROOT = pathlib.Path(__file__).parent.parent

tests/system/small/operations/test_strings.py

@@ -78,8 +78,6 @@ def test_str_extract(scalars_dfs, pat):
     bf_result = bf_series.str.extract(pat).to_pandas()
     pd_result = scalars_pandas_df[col_name].str.extract(pat)
 
-    # Pandas produces int col labels, while bq df only supports str labels at present
-    pd_result = pd_result.set_axis(pd_result.columns.astype(str), axis=1)
     pd.testing.assert_frame_equal(
         pd_result,
         bf_result,

tests/system/small/test_index.py

@@ -707,3 +707,17 @@ def test_index_str_accessor_binary(scalars_df_index, scalars_pandas_df_index):
     pd_result = pd_index.str.cat(pd_index.str[:4])
 
     pd.testing.assert_index_equal(bf_result, pd_result)
+
+
+@pytest.mark.parametrize(
+    ("pat"),
+    [(r"(ell)(lo)"), (r"(?P<somename>h..)"), (r"(?P<somename>e.*o)([g-l]+)")],
+)
+def test_index_str_extract(scalars_df_index, scalars_pandas_df_index, pat):
+    bf_index = scalars_df_index.set_index("string_col").index
+    pd_index = scalars_pandas_df_index.set_index("string_col").index
+
+    bf_result = bf_index.str.extract(pat).to_pandas()
+    pd_result = pd_index.str.extract(pat)
+
+    pd.testing.assert_frame_equal(pd_result, bf_result, check_index_type=False)

third_party/bigframes_vendored/pandas/core/indexes/base.py

@@ -366,6 +366,36 @@ def T(self) -> Index:
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+    @property
+    def str(self):
+        """
+        Vectorized string functions for Series and Index.
+
+        NAs stay NA unless handled otherwise by a particular method. Patterned
+        after Python’s string methods, with some inspiration from R’s stringr package.
+
+        **Examples:**
+
+            >>> import bigframes.pandas as bpd
+            >>> s = bpd.Series(["A_Str_Series"])
+            >>> s
+            0    A_Str_Series
+            dtype: string
+
+            >>> s.str.lower()
+            0    a_str_series
+            dtype: string
+
+            >>> s.str.replace("_", "")
+            0    AStrSeries
+            dtype: string
+
+        Returns:
+            bigframes.operations.strings.StringMethods:
+                An accessor containing string methods.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
     def copy(
         self,
         name=None,