Merge branch 'main' into shuowei-anywidget-html-repr

Author: shuoweil

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -245,27 +245,6 @@ def _cut_ops_w_intervals(
     return case_expr
 
 
-@UNARY_OP_REGISTRATION.register(agg_ops.DateSeriesDiffOp)
-def _(
-    op: agg_ops.DateSeriesDiffOp,
-    column: typed_expr.TypedExpr,
-    window: typing.Optional[window_spec.WindowSpec] = None,
-) -> sge.Expression:
-    if column.dtype != dtypes.DATE_DTYPE:
-        raise TypeError(f"Cannot perform date series diff on type {column.dtype}")
-    shift_op_impl = UNARY_OP_REGISTRATION[agg_ops.ShiftOp(0)]
-    shifted = shift_op_impl(agg_ops.ShiftOp(op.periods), column, window)
-    # Conversion factor from days to microseconds
-    conversion_factor = 24 * 60 * 60 * 1_000_000
-    return sge.Cast(
-        this=sge.DateDiff(
-            this=column.expr, expression=shifted, unit=sge.Identifier(this="DAY")
-        )
-        * sge.convert(conversion_factor),
-        to="INT64",
-    )
-
-
 @UNARY_OP_REGISTRATION.register(agg_ops.DenseRankOp)
 def _(
     op: agg_ops.DenseRankOp,
@@ -327,13 +306,27 @@ def _(
 ) -> sge.Expression:
     shift_op_impl = UNARY_OP_REGISTRATION[agg_ops.ShiftOp(0)]
     shifted = shift_op_impl(agg_ops.ShiftOp(op.periods), column, window)
-    if column.dtype in (dtypes.BOOL_DTYPE, dtypes.INT_DTYPE, dtypes.FLOAT_DTYPE):
-        if column.dtype == dtypes.BOOL_DTYPE:
-            return sge.NEQ(this=column.expr, expression=shifted)
-        else:
-            return sge.Sub(this=column.expr, expression=shifted)
-    else:
-        raise TypeError(f"Cannot perform diff on type {column.dtype}")
+    if column.dtype == dtypes.BOOL_DTYPE:
+        return sge.NEQ(this=column.expr, expression=shifted)
+
+    if column.dtype in (dtypes.INT_DTYPE, dtypes.FLOAT_DTYPE):
+        return sge.Sub(this=column.expr, expression=shifted)
+
+    if column.dtype == dtypes.TIMESTAMP_DTYPE:
+        return sge.TimestampDiff(
+            this=column.expr,
+            expression=shifted,
+            unit=sge.Identifier(this="MICROSECOND"),
+        )
+
+    if column.dtype == dtypes.DATETIME_DTYPE:
+        return sge.DatetimeDiff(
+            this=column.expr,
+            expression=shifted,
+            unit=sge.Identifier(this="MICROSECOND"),
+        )
+
+    raise TypeError(f"Cannot perform diff on type {column.dtype}")
 
 
 @UNARY_OP_REGISTRATION.register(agg_ops.MaxOp)
@@ -593,23 +586,6 @@ def _(
     return sge.func("IFNULL", expr, ir._literal(zero, column.dtype))
 
 
-@UNARY_OP_REGISTRATION.register(agg_ops.TimeSeriesDiffOp)
-def _(
-    op: agg_ops.TimeSeriesDiffOp,
-    column: typed_expr.TypedExpr,
-    window: typing.Optional[window_spec.WindowSpec] = None,
-) -> sge.Expression:
-    if column.dtype != dtypes.TIMESTAMP_DTYPE:
-        raise TypeError(f"Cannot perform time series diff on type {column.dtype}")
-    shift_op_impl = UNARY_OP_REGISTRATION[agg_ops.ShiftOp(0)]
-    shifted = shift_op_impl(agg_ops.ShiftOp(op.periods), column, window)
-    return sge.TimestampDiff(
-        this=column.expr,
-        expression=shifted,
-        unit=sge.Identifier(this="MICROSECOND"),
-    )
-
-
 @UNARY_OP_REGISTRATION.register(agg_ops.VarOp)
 def _(
     op: agg_ops.VarOp,

bigframes/core/compile/sqlglot/aggregations/windows.py

@@ -62,10 +62,10 @@ def apply_window_if_present(
 
     # This is the key change. Don't create a spec for the default window frame
     # if there's no ordering. This avoids generating an `ORDER BY NULL` clause.
-    if not window.bounds and not order:
+    if window.is_unbounded and not order:
         return sge.Window(this=value, partition_by=group_by)
 
-    if not window.bounds and not include_framing_clauses:
+    if window.is_unbounded and not include_framing_clauses:
         return sge.Window(this=value, partition_by=group_by, order=order)
 
     kind = (

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

@@ -16,7 +16,9 @@
 
 import sqlglot.expressions as sge
 
+from bigframes import dtypes
 from bigframes import operations as ops
+from bigframes.core.compile.constants import UNIT_TO_US_CONVERSION_FACTORS
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
@@ -81,22 +83,73 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_unary_op(ops.StrftimeOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.StrftimeOp) -> sge.Expression:
-    return sge.func("FORMAT_TIMESTAMP", sge.convert(op.date_format), expr.expr)
+    func_name = ""
+    if expr.dtype == dtypes.DATE_DTYPE:
+        func_name = "FORMAT_DATE"
+    elif expr.dtype == dtypes.DATETIME_DTYPE:
+        func_name = "FORMAT_DATETIME"
+    elif expr.dtype == dtypes.TIME_DTYPE:
+        func_name = "FORMAT_TIME"
+    elif expr.dtype == dtypes.TIMESTAMP_DTYPE:
+        func_name = "FORMAT_TIMESTAMP"
+
+    return sge.func(func_name, sge.convert(op.date_format), expr.expr)
 
 
 @register_unary_op(ops.time_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("TIME", expr.expr)
 
 
-@register_unary_op(ops.ToDatetimeOp)
-def _(expr: TypedExpr) -> sge.Expression:
-    return sge.Cast(this=sge.func("TIMESTAMP_SECONDS", expr.expr), to="DATETIME")
-
+@register_unary_op(ops.ToDatetimeOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.ToDatetimeOp) -> sge.Expression:
+    if op.format:
+        result = expr.expr
+        if expr.dtype != dtypes.STRING_DTYPE:
+            result = sge.Cast(this=result, to="STRING")
+        result = sge.func(
+            "PARSE_TIMESTAMP", sge.convert(op.format), result, sge.convert("UTC")
+        )
+        return sge.Cast(this=result, to="DATETIME")
+
+    if expr.dtype == dtypes.STRING_DTYPE:
+        return sge.TryCast(this=expr.expr, to="DATETIME")
+
+    value = expr.expr
+    unit = op.unit or "ns"
+    factor = UNIT_TO_US_CONVERSION_FACTORS[unit]
+    if factor != 1:
+        value = sge.Mul(this=value, expression=sge.convert(factor))
+    value = sge.func("TRUNC", value)
+    return sge.Cast(
+        this=sge.func("TIMESTAMP_MICROS", sge.Cast(this=value, to="INT64")),
+        to="DATETIME",
+    )
+
+
+@register_unary_op(ops.ToTimestampOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.ToTimestampOp) -> sge.Expression:
+    if op.format:
+        result = expr.expr
+        if expr.dtype != dtypes.STRING_DTYPE:
+            result = sge.Cast(this=result, to="STRING")
+        return sge.func(
+            "PARSE_TIMESTAMP", sge.convert(op.format), expr.expr, sge.convert("UTC")
+        )
 
-@register_unary_op(ops.ToTimestampOp)
-def _(expr: TypedExpr) -> sge.Expression:
-    return sge.func("TIMESTAMP_SECONDS", expr.expr)
+    if expr.dtype == dtypes.STRING_DTYPE:
+        return sge.func("TIMESTAMP", expr.expr)
+
+    value = expr.expr
+    unit = op.unit or "ns"
+    factor = UNIT_TO_US_CONVERSION_FACTORS[unit]
+    if factor != 1:
+        value = sge.Mul(this=value, expression=sge.convert(factor))
+    value = sge.func("TRUNC", value)
+    return sge.Cast(
+        this=sge.func("TIMESTAMP_MICROS", sge.Cast(this=value, to="INT64")),
+        to="TIMESTAMP",
+    )
 
 
 @register_unary_op(ops.UnixMicros)

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

@@ -16,6 +16,7 @@
 
 import sqlglot.expressions as sge
 
+from bigframes import dtypes
 from bigframes import operations as ops
 from bigframes.core.compile.constants import UNIT_TO_US_CONVERSION_FACTORS
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
@@ -32,7 +33,12 @@ def _(expr: TypedExpr) -> sge.Expression:
 @register_unary_op(ops.ToTimedeltaOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.ToTimedeltaOp) -> sge.Expression:
     value = expr.expr
+    if expr.dtype == dtypes.TIMEDELTA_DTYPE:
+        return value
+
     factor = UNIT_TO_US_CONVERSION_FACTORS[op.unit]
     if factor != 1:
         value = sge.Mul(this=value, expression=sge.convert(factor))
+    if expr.dtype == dtypes.FLOAT_DTYPE:
+        value = sge.Cast(this=sge.Floor(this=value), to=sge.DataType(this="INT64"))
     return value

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -648,6 +648,8 @@ def _literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
     elif dtype == dtypes.BYTES_DTYPE:
         return _cast(str(value), sqlglot_type)
     elif dtypes.is_time_like(dtype):
+        if isinstance(value, str):
+            return _cast(sge.convert(value), sqlglot_type)
         if isinstance(value, np.generic):
             value = value.item()
         return _cast(sge.convert(value.isoformat()), sqlglot_type)

bigframes/core/local_data.py

@@ -124,12 +124,13 @@ def to_arrow(
         geo_format: Literal["wkb", "wkt"] = "wkt",
         duration_type: Literal["int", "duration"] = "duration",
         json_type: Literal["string"] = "string",
+        max_chunksize: Optional[int] = None,
     ) -> tuple[pa.Schema, Iterable[pa.RecordBatch]]:
         if geo_format != "wkt":
             raise NotImplementedError(f"geo format {geo_format} not yet implemented")
         assert json_type == "string"
 
-        batches = self.data.to_batches()
+        batches = self.data.to_batches(max_chunksize=max_chunksize)
         schema = self.data.schema
         if duration_type == "int":
             schema = _schema_durations_to_ints(schema)