address comments

Author: chelsea-lin

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

@@ -33,46 +33,15 @@ def _(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
     sg_expr = expr.expr
 
     if to_type == dtypes.JSON_DTYPE:
-        if from_type == dtypes.STRING_DTYPE:
-            func_name = "PARSE_JSON_IN_SAFE" if op.safe else "PARSE_JSON"
-            return sge.func(func_name, sg_expr)
-        if from_type in (dtypes.INT_DTYPE, dtypes.BOOL_DTYPE, dtypes.FLOAT_DTYPE):
-            sg_expr = sge.Cast(this=sg_expr, to="STRING")
-            return sge.func("PARSE_JSON", sg_expr)
-        raise TypeError(f"Cannot cast from {from_type} to {to_type}")
+        return _cast_to_json(expr, op)
 
     if from_type == dtypes.JSON_DTYPE:
-        func_name = ""
-        if to_type == dtypes.INT_DTYPE:
-            func_name = "INT64"
-        elif to_type == dtypes.FLOAT_DTYPE:
-            func_name = "FLOAT64"
-        elif to_type == dtypes.BOOL_DTYPE:
-            func_name = "BOOL"
-        elif to_type == dtypes.STRING_DTYPE:
-            func_name = "STRING"
-        if func_name:
-            func_name = "SAFE." + func_name if op.safe else func_name
-            return sge.func(func_name, sg_expr)
-        raise TypeError(f"Cannot cast from {from_type} to {to_type}")
+        return _cast_from_json(expr, op)
 
     if to_type == dtypes.INT_DTYPE:
-        # Cannot cast DATETIME to INT directly so need to convert to TIMESTAMP first.
-        if from_type == dtypes.DATETIME_DTYPE:
-            sg_expr = _cast(sg_expr, "TIMESTAMP", op.safe)
-            return sge.func("UNIX_MICROS", sg_expr)
-        if from_type == dtypes.TIMESTAMP_DTYPE:
-            return sge.func("UNIX_MICROS", sg_expr)
-        if from_type == dtypes.TIME_DTYPE:
-            return sge.func(
-                "TIME_DIFF",
-                _cast(sg_expr, "TIME", op.safe),
-                sge.convert("00:00:00"),
-                "MICROSECOND",
-            )
-        if from_type == dtypes.NUMERIC_DTYPE or from_type == dtypes.FLOAT_DTYPE:
-            sg_expr = sge.func("TRUNC", sg_expr)
-            return _cast(sg_expr, sg_to_type, op.safe)
+        result = _cast_to_int(expr, op)
+        if result is not None:
+            return result
 
     if to_type == dtypes.FLOAT_DTYPE and from_type == dtypes.BOOL_DTYPE:
         sg_expr = _cast(sg_expr, "INT64", op.safe)
@@ -124,6 +93,59 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 
 # Helper functions
+def _cast_to_json(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
+    from_type = expr.dtype
+    sg_expr = expr.expr
+
+    if from_type == dtypes.STRING_DTYPE:
+        func_name = "PARSE_JSON_IN_SAFE" if op.safe else "PARSE_JSON"
+        return sge.func(func_name, sg_expr)
+    if from_type in (dtypes.INT_DTYPE, dtypes.BOOL_DTYPE, dtypes.FLOAT_DTYPE):
+        sg_expr = sge.Cast(this=sg_expr, to="STRING")
+        return sge.func("PARSE_JSON", sg_expr)
+    raise TypeError(f"Cannot cast from {from_type} to {dtypes.JSON_DTYPE}")
+
+
+def _cast_from_json(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
+    to_type = op.to_type
+    sg_expr = expr.expr
+    func_name = ""
+    if to_type == dtypes.INT_DTYPE:
+        func_name = "INT64"
+    elif to_type == dtypes.FLOAT_DTYPE:
+        func_name = "FLOAT64"
+    elif to_type == dtypes.BOOL_DTYPE:
+        func_name = "BOOL"
+    elif to_type == dtypes.STRING_DTYPE:
+        func_name = "STRING"
+    if func_name:
+        func_name = "SAFE." + func_name if op.safe else func_name
+        return sge.func(func_name, sg_expr)
+    raise TypeError(f"Cannot cast from {dtypes.JSON_DTYPE} to {to_type}")
+
+
+def _cast_to_int(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression | None:
+    from_type = expr.dtype
+    sg_expr = expr.expr
+    # Cannot cast DATETIME to INT directly so need to convert to TIMESTAMP first.
+    if from_type == dtypes.DATETIME_DTYPE:
+        sg_expr = _cast(sg_expr, "TIMESTAMP", op.safe)
+        return sge.func("UNIX_MICROS", sg_expr)
+    if from_type == dtypes.TIMESTAMP_DTYPE:
+        return sge.func("UNIX_MICROS", sg_expr)
+    if from_type == dtypes.TIME_DTYPE:
+        return sge.func(
+            "TIME_DIFF",
+            _cast(sg_expr, "TIME", op.safe),
+            sge.convert("00:00:00"),
+            "MICROSECOND",
+        )
+    if from_type == dtypes.NUMERIC_DTYPE or from_type == dtypes.FLOAT_DTYPE:
+        sg_expr = sge.func("TRUNC", sg_expr)
+        return _cast(sg_expr, "INT64", op.safe)
+    return None
+
+
 def _cast(expr: sge.Expression, to: str, safe: bool):
     if safe:
         return sge.TryCast(this=expr, to=to)