feat: Can cast locally in hybrid engine (#1944)

Author: TrevorBergeron

bigframes/core/compile/polars/compiler.py

@@ -33,7 +33,9 @@
 import bigframes.operations.aggregations as agg_ops
 import bigframes.operations.bool_ops as bool_ops
 import bigframes.operations.comparison_ops as comp_ops
+import bigframes.operations.datetime_ops as dt_ops
 import bigframes.operations.generic_ops as gen_ops
+import bigframes.operations.json_ops as json_ops
 import bigframes.operations.numeric_ops as num_ops
 import bigframes.operations.string_ops as string_ops
 
@@ -280,6 +282,30 @@ def _(self, op: ops.ScalarOp, l_input: pl.Expr, r_input: pl.Expr) -> pl.Expr:
  assert isinstance(op, string_ops.StrConcatOp)
  return pl.concat_str(l_input, r_input)
 
+ @compile_op.register(dt_ops.StrftimeOp)
+ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+ assert isinstance(op, dt_ops.StrftimeOp)
+ return input.dt.strftime(op.date_format)
+
+ @compile_op.register(dt_ops.ParseDatetimeOp)
+ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+ assert isinstance(op, dt_ops.ParseDatetimeOp)
+ return input.str.to_datetime(
+ time_unit="us", time_zone=None, ambiguous="earliest"
+ )
+
+ @compile_op.register(dt_ops.ParseTimestampOp)
+ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+ assert isinstance(op, dt_ops.ParseTimestampOp)
+ return input.str.to_datetime(
+ time_unit="us", time_zone="UTC", ambiguous="earliest"
+ )
+
+ @compile_op.register(json_ops.JSONDecode)
+ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+ assert isinstance(op, json_ops.JSONDecode)
+ return input.str.json_decode(_DTYPE_MAPPING[op.to_type])
+
  @dataclasses.dataclass(frozen=True)
  class PolarsAggregateCompiler:
  scalar_compiler = PolarsExpressionCompiler()

bigframes/core/compile/polars/lowering.py

@@ -17,7 +17,7 @@
 from bigframes import dtypes
 from bigframes.core import bigframe_node, expression
 from bigframes.core.rewrite import op_lowering
-from bigframes.operations import comparison_ops, numeric_ops
+from bigframes.operations import comparison_ops, datetime_ops, json_ops, numeric_ops
 import bigframes.operations as ops
 
 # TODO: Would be more precise to actually have separate op set for polars ops (where they diverge from the original ops)
@@ -278,6 +278,16 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
  return wo_bools
 
 
+class LowerAsTypeRule(op_lowering.OpLoweringRule):
+ @property
+ def op(self) -> type[ops.ScalarOp]:
+ return ops.AsTypeOp
+
+ def lower(self, expr: expression.OpExpression) -> expression.Expression:
+ assert isinstance(expr.op, ops.AsTypeOp)
+ return _lower_cast(expr.op, expr.inputs[0])
+
+
 def _coerce_comparables(
  expr1: expression.Expression,
  expr2: expression.Expression,
@@ -299,12 +309,57 @@ def _coerce_comparables(
  return expr1, expr2
 
 
-# TODO: Need to handle bool->string cast to get capitalization correct
 def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
+ if arg.output_type == cast_op.to_type:
+ return arg
+
+ if arg.output_type == dtypes.JSON_DTYPE:
+ return json_ops.JSONDecode(cast_op.to_type).as_expr(arg)
+ if (
+ arg.output_type == dtypes.STRING_DTYPE
+ and cast_op.to_type == dtypes.DATETIME_DTYPE
+ ):
+ return datetime_ops.ParseDatetimeOp().as_expr(arg)
+ if (
+ arg.output_type == dtypes.STRING_DTYPE
+ and cast_op.to_type == dtypes.TIMESTAMP_DTYPE
+ ):
+ return datetime_ops.ParseTimestampOp().as_expr(arg)
+ # date -> string casting
+ if (
+ arg.output_type == dtypes.DATETIME_DTYPE
+ and cast_op.to_type == dtypes.STRING_DTYPE
+ ):
+ return datetime_ops.StrftimeOp("%Y-%m-%d %H:%M:%S").as_expr(arg)
+ if arg.output_type == dtypes.TIME_DTYPE and cast_op.to_type == dtypes.STRING_DTYPE:
+ return datetime_ops.StrftimeOp("%H:%M:%S.%6f").as_expr(arg)
+ if (
+ arg.output_type == dtypes.TIMESTAMP_DTYPE
+ and cast_op.to_type == dtypes.STRING_DTYPE
+ ):
+ return datetime_ops.StrftimeOp("%Y-%m-%d %H:%M:%S%.6f%:::z").as_expr(arg)
+ if arg.output_type == dtypes.BOOL_DTYPE and cast_op.to_type == dtypes.STRING_DTYPE:
+ # bool -> decimal needs two-step cast
+ new_arg = ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(arg)
+ is_true_cond = ops.eq_op.as_expr(arg, expression.const(True))
+ is_false_cond = ops.eq_op.as_expr(arg, expression.const(False))
+ return ops.CaseWhenOp().as_expr(
+ is_true_cond,
+ expression.const("True"),
+ is_false_cond,
+ expression.const("False"),
+ )
  if arg.output_type == dtypes.BOOL_DTYPE and dtypes.is_numeric(cast_op.to_type):
  # bool -> decimal needs two-step cast
  new_arg = ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(arg)
  return cast_op.as_expr(new_arg)
+ if arg.output_type == dtypes.TIME_DTYPE and dtypes.is_numeric(cast_op.to_type):
+ # polars cast gives nanoseconds, so convert to microseconds
+ return numeric_ops.floordiv_op.as_expr(
+ cast_op.as_expr(arg), expression.const(1000)
+ )
+ if dtypes.is_numeric(arg.output_type) and cast_op.to_type == dtypes.TIME_DTYPE:
+ return cast_op.as_expr(ops.mul_op.as_expr(expression.const(1000), arg))
  return cast_op.as_expr(arg)
 
 
@@ -329,6 +384,7 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
  LowerDivRule(),
  LowerFloorDivRule(),
  LowerModRule(),
+ LowerAsTypeRule(),
 )
 
 

bigframes/operations/datetime_ops.py

@@ -39,6 +39,28 @@
 time_op = TimeOp()
 
 
+@dataclasses.dataclass(frozen=True)
+class ParseDatetimeOp(base_ops.UnaryOp):
+ # TODO: Support strict format
+ name: typing.ClassVar[str] = "parse_datetime"
+
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ if input_types[0] != dtypes.STRING_DTYPE:
+ raise TypeError("expected string input")
+ return pd.ArrowDtype(pa.timestamp("us", tz=None))
+
+
+@dataclasses.dataclass(frozen=True)
+class ParseTimestampOp(base_ops.UnaryOp):
+ # TODO: Support strict format
+ name: typing.ClassVar[str] = "parse_timestamp"
+
+ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+ if input_types[0] != dtypes.STRING_DTYPE:
+ raise TypeError("expected string input")
+ return pd.ArrowDtype(pa.timestamp("us", tz="UTC"))
+
+
 @dataclasses.dataclass(frozen=True)
 class ToDatetimeOp(base_ops.UnaryOp):
  name: typing.ClassVar[str] = "to_datetime"