refactor: fix str.extract doctest

Author: chelsea-lin

bigframes/core/compile/sqlglot/expressions/string_ops.py

@@ -48,12 +48,14 @@ def _(expr: TypedExpr, op: ops.StrExtractOp) -> sge.Expression:
     # Cannot use BigQuery's REGEXP_EXTRACT function, which only allows one
     # capturing group.
     pat_expr = sge.convert(op.pat)
-    if op.n != 0:
-        pat_expr = sge.func("CONCAT", sge.convert(".*?"), pat_expr, sge.convert(".*"))
-    else:
+    if op.n == 0:
         pat_expr = sge.func("CONCAT", sge.convert(".*?("), pat_expr, sge.convert(").*"))
+        n = 1
+    else:
+        pat_expr = sge.func("CONCAT", sge.convert(".*?"), pat_expr, sge.convert(".*"))
+        n = op.n
 
-    rex_replace = sge.func("REGEXP_REPLACE", expr.expr, pat_expr, sge.convert(r"\1"))
+    rex_replace = sge.func("REGEXP_REPLACE", expr.expr, pat_expr, sge.convert(f"\\{n}"))
     rex_contains = sge.func("REGEXP_CONTAINS", expr.expr, sge.convert(op.pat))
     return sge.If(this=rex_contains, true=rex_replace, false=sge.null())
 

tests/unit/core/compile/sqlglot/expressions/snapshots/test_string_ops/test_str_extract/out.sql

@@ -5,13 +5,19 @@ WITH `bfcte_0` AS (
 ), `bfcte_1` AS (
   SELECT
     *,
+    IF(
+      REGEXP_CONTAINS(`string_col`, '([a-z]*)'),
+      REGEXP_REPLACE(`string_col`, CONCAT('.*?(', '([a-z]*)', ').*'), '\\1'),
+      NULL
+    ) AS `bfcol_1`,
     IF(
       REGEXP_CONTAINS(`string_col`, '([a-z]*)'),
       REGEXP_REPLACE(`string_col`, CONCAT('.*?', '([a-z]*)', '.*'), '\\1'),
       NULL
-    ) AS `bfcol_1`
+    ) AS `bfcol_2`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `string_col`
+  `bfcol_1` AS `zero`,
+  `bfcol_2` AS `one`
 FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/test_string_ops.py

@@ -260,9 +260,11 @@ def test_str_contains_regex(scalar_types_df: bpd.DataFrame, snapshot):
 def test_str_extract(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "string_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_ops_to_sql(
-        bf_df, [ops.StrExtractOp(r"([a-z]*)", 1).as_expr(col_name)], [col_name]
-    )
+    ops_map = {
+        "zero": ops.StrExtractOp(r"([a-z]*)", 0).as_expr(col_name),
+        "one": ops.StrExtractOp(r"([a-z]*)", 1).as_expr(col_name),
+    }
+    sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
 
     snapshot.assert_match(sql, "out.sql")