refactor: fix ops.ZfillOp

Author: chelsea-lin

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

@@ -288,27 +288,29 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
 @register_unary_op(ops.ZfillOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.ZfillOp) -> sge.Expression:
+    length_expr = sge.Greatest(
+        expressions=[sge.Length(this=expr.expr), sge.convert(op.width)]
+    )
     return sge.Case(
         ifs=[
             sge.If(
-                this=sge.EQ(
-                    this=sge.Substring(
-                        this=expr.expr, start=sge.convert(1), length=sge.convert(1)
-                    ),
-                    expression=sge.convert("-"),
+                this=sge.func(
+                    "STARTS_WITH",
+                    expr.expr,
+                    sge.convert("-"),
                 ),
                 true=sge.Concat(
                     expressions=[
                         sge.convert("-"),
                         sge.func(
                             "LPAD",
-                            sge.Substring(this=expr.expr, start=sge.convert(1)),
-                            sge.convert(op.width - 1),
+                            sge.Substring(this=expr.expr, start=sge.convert(2)),
+                            length_expr - 1,
                             sge.convert("0"),
                         ),
                     ]
                 ),
             )
         ],
-        default=sge.func("LPAD", expr.expr, sge.convert(op.width), sge.convert("0")),
+        default=sge.func("LPAD", expr.expr, length_expr, sge.convert("0")),
     )

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

@@ -6,9 +6,9 @@ WITH `bfcte_0` AS (
   SELECT
     *,
     CASE
-      WHEN SUBSTRING(`string_col`, 1, 1) = '-'
-      THEN CONCAT('-', LPAD(SUBSTRING(`string_col`, 1), 9, '0'))
-      ELSE LPAD(`string_col`, 10, '0')
+      WHEN STARTS_WITH(`string_col`, '-')
+      THEN CONCAT('-', LPAD(SUBSTRING(`string_col`, 2), GREATEST(LENGTH(`string_col`), 10) - 1, '0'))
+      ELSE LPAD(`string_col`, GREATEST(LENGTH(`string_col`), 10), '0')
     END AS `bfcol_1`
   FROM `bfcte_0`
 )