Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,10 @@ class ExpressionCompiler:
sge.LT,
sge.EQ,
sge.NEQ,
sge.Like,
sge.RegexpLike,
sge.In,
sge.Between,
# Logical operations
sge.And,
sge.Or,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ def _(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:

@register_unary_op(ops.ArrayToStringOp, pass_op=True)
def _(expr: TypedExpr, op: ops.ArrayToStringOp) -> sge.Expression:
return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
return sge.ArrayToString(this=expr.expr, expression=sge.convert(op.delimiter))


@register_nary_op(ops.ToArrayOp)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -373,9 +373,11 @@ def _(expr: TypedExpr, op: ops.ToDatetimeOp) -> sge.Expression:
)
return sge.Cast(this=result, to="DATETIME")

if expr.dtype == dtypes.TIMESTAMP_DTYPE:
return sge.func("DATETIME", expr.expr, sge.convert("UTC"))

if expr.dtype in (
dtypes.STRING_DTYPE,
dtypes.TIMESTAMP_DTYPE,
dtypes.DATETIME_DTYPE,
dtypes.DATE_DTYPE,
):
Expand All @@ -387,9 +389,10 @@ def _(expr: TypedExpr, op: ops.ToDatetimeOp) -> sge.Expression:
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",
return sge.func(
"DATETIME",
sge.func("TIMESTAMP_MICROS", sge.Cast(this=value, to="INT64")),
sge.convert("UTC"),
)


Expand All @@ -400,7 +403,7 @@ def _(expr: TypedExpr, op: ops.ToTimestampOp) -> sge.Expression:
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")
"PARSE_TIMESTAMP", sge.convert(op.format), result, sge.convert("UTC")
)

if expr.dtype in (
Expand Down Expand Up @@ -683,3 +686,49 @@ def _integer_label_to_datetime_op_yearly_freq(
this=next_month_date, expression=sge.Interval(this=one, unit="DAY")
)
return sge.Cast(this=x_label, to=sqlglot_types.from_bigframes_dtype(y.dtype))


@register_binary_op(ops.timestamp_add_op)
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
return sge.TimestampAdd(
this=left.expr, expression=right.expr, unit=sge.Var(this="MICROSECOND")
)


@register_binary_op(ops.timestamp_sub_op)
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
return sge.TimestampSub(
this=left.expr, expression=right.expr, unit=sge.Var(this="MICROSECOND")
)


@register_binary_op(ops.timestamp_diff_op)
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
return sge.TimestampDiff(
this=left.expr, expression=right.expr, unit=sge.Var(this="MICROSECOND")
)


@register_binary_op(ops.date_add_op)
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
left_expr = sge.Cast(this=left.expr, to="TIMESTAMP")
return sge.TimestampAdd(
this=left_expr, expression=right.expr, unit=sge.Var(this="MICROSECOND")
)


@register_binary_op(ops.date_sub_op)
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
left_expr = sge.Cast(this=left.expr, to="TIMESTAMP")
return sge.TimestampSub(
this=left_expr, expression=right.expr, unit=sge.Var(this="MICROSECOND")
)


@register_binary_op(ops.date_diff_op)
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
diff = sge.DateDiff(this=left.expr, expression=right.expr, unit=sge.Var(this="DAY"))
return sge.Mul(
this=diff,
expression=sge.convert(int(UNIT_TO_US_CONVERSION_FACTORS["d"])),
)
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import itertools

import pytest

import bigframes.operations as ops
from bigframes.core import array_value
from bigframes.session import polars_executor
from bigframes.testing.engine_utils import assert_equivalence_execution
from bigframes.testing.engine_utils import SPEC, assert_equivalence_execution

pytest.importorskip("polars")

Expand Down Expand Up @@ -68,3 +69,16 @@ def test_engines_project_comparison_op(
# bool col actually doesn't work properly for bq engine
arr = apply_op_pairwise(scalars_array_value, op, excluded_cols=["string_col"])
assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)


@pytest.mark.parametrize("engine", ["bq-sqlglot"], indirect=True)
def test_engines_precedence_like_and_in(
scalars_array_value: array_value.ArrayValue, engine
):
exprs = [
ops.eq_op.as_expr("bool_col", ops.StrContainsOp("a").as_expr("string_col")),
]
arr, _ = scalars_array_value.compute_values(exprs)
res = asyncio.run(engine.execute(arr.node, SPEC))
assert res is not None
assert len(res.batches().to_pandas()) > 0
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,17 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio

import pandas as pd
import pytest

import bigframes.operations as ops
from bigframes import dtypes
from bigframes.core import array_value
from bigframes.core import expression as ex
from bigframes.session import polars_executor
from bigframes.testing.engine_utils import assert_equivalence_execution
from bigframes.testing.engine_utils import SPEC, assert_equivalence_execution

pytest.importorskip("polars")

Expand Down Expand Up @@ -64,3 +69,45 @@ def test_engines_date_accessors(scalars_array_value: array_value.ArrayValue, eng

arr, _ = scalars_array_value.compute_values(exprs)
assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)


@pytest.mark.parametrize("engine", ["bq", "bq-sqlglot"], indirect=True)
def test_engines_temporal_arithmetic(
scalars_array_value: array_value.ArrayValue, engine
):
exprs = [
ops.timestamp_add_op.as_expr(
"timestamp_col", ex.const(pd.Timedelta(seconds=1), dtypes.TIMEDELTA_DTYPE)
),
ops.timestamp_sub_op.as_expr(
"timestamp_col", ex.const(pd.Timedelta(seconds=1), dtypes.TIMEDELTA_DTYPE)
),
ops.date_add_op.as_expr(
"date_col", ex.const(pd.Timedelta(days=1), dtypes.TIMEDELTA_DTYPE)
),
ops.date_sub_op.as_expr(
"date_col", ex.const(pd.Timedelta(days=1), dtypes.TIMEDELTA_DTYPE)
),
ops.timestamp_diff_op.as_expr("timestamp_col", "timestamp_col"),
ops.date_diff_op.as_expr("date_col", "date_col"),
]

arr, _ = scalars_array_value.compute_values(exprs)
res = asyncio.run(engine.execute(arr.node, SPEC))
assert res is not None
assert len(res.batches().to_pandas()) > 0


@pytest.mark.parametrize("engine", ["bq", "bq-sqlglot"], indirect=True)
def test_engines_to_datetime(scalars_array_value: array_value.ArrayValue, engine):
exprs = [
ops.ToDatetimeOp().as_expr("timestamp_col"),
]
arr, _ = scalars_array_value.compute_values(exprs)
res = asyncio.run(engine.execute(arr.node, SPEC))
assert res is not None
df = res.batches().to_pandas()
# The input timestamp was: TIMESTAMP('2021-07-21T17:43:43.945289+00:00')
# The output should be naive DATETIME('2021-07-21T17:43:43.945289')
val = df.iloc[0, -1]
assert pd.Timestamp(val) == pd.Timestamp("2021-07-21T17:43:43.945289")
Loading