diff --git a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py index 011413b67f..609ac374b6 100644 --- a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py +++ b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py @@ -239,6 +239,26 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.Floor(this=expr.expr) +@UNARY_OP_REGISTRATION.register(ops.geo_area_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.func("ST_AREA", expr.expr) + + +@UNARY_OP_REGISTRATION.register(ops.geo_st_astext_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.func("ST_ASTEXT", expr.expr) + + +@UNARY_OP_REGISTRATION.register(ops.geo_x_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.func("SAFE.ST_X", expr.expr) + + +@UNARY_OP_REGISTRATION.register(ops.geo_y_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.func("SAFE.ST_Y", expr.expr) + + @UNARY_OP_REGISTRATION.register(ops.hash_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.func("FARM_FINGERPRINT", expr.expr) @@ -361,6 +381,16 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.Lower(this=expr.expr) +@UNARY_OP_REGISTRATION.register(ops.minute_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.Extract(this=sge.Identifier(this="MINUTE"), expression=expr.expr) + + +@UNARY_OP_REGISTRATION.register(ops.month_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.Extract(this=sge.Identifier(this="MONTH"), expression=expr.expr) + + @UNARY_OP_REGISTRATION.register(ops.StrLstripOp) def _(op: ops.StrLstripOp, expr: TypedExpr) -> sge.Expression: return sge.Trim(this=expr.expr, expression=sge.convert(op.to_strip), side="LEFT") @@ -371,16 +401,31 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.Neg(this=expr.expr) +@UNARY_OP_REGISTRATION.register(ops.normalize_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.TimestampTrunc(this=expr.expr, unit=sge.Identifier(this="DAY")) + + @UNARY_OP_REGISTRATION.register(ops.pos_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return expr.expr +@UNARY_OP_REGISTRATION.register(ops.quarter_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.Extract(this=sge.Identifier(this="QUARTER"), expression=expr.expr) + + @UNARY_OP_REGISTRATION.register(ops.reverse_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.func("REVERSE", expr.expr) +@UNARY_OP_REGISTRATION.register(ops.second_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.Extract(this=sge.Identifier(this="SECOND"), expression=expr.expr) + + @UNARY_OP_REGISTRATION.register(ops.StrRstripOp) def _(op: ops.StrRstripOp, expr: TypedExpr) -> sge.Expression: return sge.Trim(this=expr.expr, expression=sge.convert(op.to_strip), side="RIGHT") @@ -414,6 +459,11 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.Extract(this=sge.Identifier(this="ISOWEEK"), expression=expr.expr) +@UNARY_OP_REGISTRATION.register(ops.iso_year_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.Extract(this=sge.Identifier(this="ISOYEAR"), expression=expr.expr) + + @UNARY_OP_REGISTRATION.register(ops.isnull_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.Is(this=expr.expr, expression=sge.Null()) @@ -442,6 +492,31 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: ) +@UNARY_OP_REGISTRATION.register(ops.StrGetOp) +def _(op: ops.StrGetOp, expr: TypedExpr) -> sge.Expression: + return sge.Substring( + this=expr.expr, + start=sge.convert(op.i + 1), + length=sge.convert(1), + ) + + +@UNARY_OP_REGISTRATION.register(ops.StrSliceOp) +def _(op: ops.StrSliceOp, expr: TypedExpr) -> sge.Expression: + start = op.start + 1 if op.start is not None else None + if op.end is None: + length = None + elif op.start is None: + length = op.end + else: + length = op.end - op.start + return sge.Substring( + this=expr.expr, + start=sge.convert(start) if start is not None else None, + length=sge.convert(length) if length is not None else None, + ) + + @UNARY_OP_REGISTRATION.register(ops.tan_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.func("TAN", expr.expr) @@ -452,6 +527,16 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.func("TANH", expr.expr) +@UNARY_OP_REGISTRATION.register(ops.time_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.func("TIME", expr.expr) + + +@UNARY_OP_REGISTRATION.register(ops.timedelta_floor_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.Floor(this=expr.expr) + + # JSON Ops @UNARY_OP_REGISTRATION.register(ops.JSONExtract) def _(op: ops.JSONExtract, expr: TypedExpr) -> sge.Expression: @@ -501,3 +586,8 @@ def _(op: ops.ToJSONString, expr: TypedExpr) -> sge.Expression: @UNARY_OP_REGISTRATION.register(ops.upper_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.Upper(this=expr.expr) + + +@UNARY_OP_REGISTRATION.register(ops.year_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.Extract(this=sge.Identifier(this="YEAR"), expression=expr.expr) diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_area/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_area/out.sql new file mode 100644 index 0000000000..9b4b6894e0 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_area/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `geography_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + ST_AREA(`bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `geography_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_astext/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_astext/out.sql new file mode 100644 index 0000000000..9557e2f1d6 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_astext/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `geography_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + ST_ASTEXT(`bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `geography_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_x/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_x/out.sql new file mode 100644 index 0000000000..d4c0370ca8 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_x/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `geography_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + SAFE.ST_X(`bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `geography_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_y/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_y/out.sql new file mode 100644 index 0000000000..196c2fcad6 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_y/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `geography_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + SAFE.ST_Y(`bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `geography_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_iso_year/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_iso_year/out.sql new file mode 100644 index 0000000000..13b56f709c --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_iso_year/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + EXTRACT(ISOYEAR FROM `bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_minute/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_minute/out.sql new file mode 100644 index 0000000000..4ef9b8142f --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_minute/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + EXTRACT(MINUTE FROM `bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_month/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_month/out.sql new file mode 100644 index 0000000000..4912622898 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_month/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + EXTRACT(MONTH FROM `bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_normalize/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_normalize/out.sql new file mode 100644 index 0000000000..3c7efd3098 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_normalize/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + TIMESTAMP_TRUNC(`bfcol_0`, DAY) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_quarter/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_quarter/out.sql new file mode 100644 index 0000000000..2be2866661 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_quarter/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + EXTRACT(QUARTER FROM `bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_second/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_second/out.sql new file mode 100644 index 0000000000..144b704788 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_second/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + EXTRACT(SECOND FROM `bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_str_get/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_str_get/out.sql new file mode 100644 index 0000000000..1278c3435d --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_str_get/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `string_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + SUBSTRING(`bfcol_0`, 2, 1) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `string_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_str_slice/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_str_slice/out.sql new file mode 100644 index 0000000000..4f97ab3ac6 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_str_slice/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `string_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + SUBSTRING(`bfcol_0`, 2, 2) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `string_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_time/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_time/out.sql new file mode 100644 index 0000000000..6b74efafd5 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_time/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + TIME(`bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_timedelta_floor/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_timedelta_floor/out.sql new file mode 100644 index 0000000000..1a8b9f4e39 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_timedelta_floor/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `int64_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + FLOOR(`bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `int64_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_year/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_year/out.sql new file mode 100644 index 0000000000..8352a65e9e --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_year/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `timestamp_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + EXTRACT(YEAR FROM `bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `timestamp_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index 5f29e98e81..236f94045f 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py @@ -146,6 +146,34 @@ def test_floor(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_geo_area(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["geography_col"]] + sql = _apply_unary_op(bf_df, ops.geo_area_op, "geography_col") + + snapshot.assert_match(sql, "out.sql") + + +def test_geo_st_astext(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["geography_col"]] + sql = _apply_unary_op(bf_df, ops.geo_st_astext_op, "geography_col") + + snapshot.assert_match(sql, "out.sql") + + +def test_geo_x(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["geography_col"]] + sql = _apply_unary_op(bf_df, ops.geo_x_op, "geography_col") + + snapshot.assert_match(sql, "out.sql") + + +def test_geo_y(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["geography_col"]] + sql = _apply_unary_op(bf_df, ops.geo_y_op, "geography_col") + + snapshot.assert_match(sql, "out.sql") + + def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot): bf_df = repeated_types_df[["string_list_col"]] sql = _apply_unary_op(bf_df, ops.ArrayToStringOp(delimiter="."), "string_list_col") @@ -307,6 +335,20 @@ def test_lstrip(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_minute(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + sql = _apply_unary_op(bf_df, ops.minute_op, "timestamp_col") + + snapshot.assert_match(sql, "out.sql") + + +def test_month(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + sql = _apply_unary_op(bf_df, ops.month_op, "timestamp_col") + + snapshot.assert_match(sql, "out.sql") + + def test_neg(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["float64_col"]] sql = _apply_unary_op(bf_df, ops.neg_op, "float64_col") @@ -314,6 +356,13 @@ def test_neg(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_normalize(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + sql = _apply_unary_op(bf_df, ops.normalize_op, "timestamp_col") + + snapshot.assert_match(sql, "out.sql") + + def test_pos(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["float64_col"]] sql = _apply_unary_op(bf_df, ops.pos_op, "float64_col") @@ -321,6 +370,13 @@ def test_pos(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_quarter(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + sql = _apply_unary_op(bf_df, ops.quarter_op, "timestamp_col") + + snapshot.assert_match(sql, "out.sql") + + def test_reverse(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["string_col"]] sql = _apply_unary_op(bf_df, ops.reverse_op, "string_col") @@ -328,6 +384,13 @@ def test_reverse(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_second(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + sql = _apply_unary_op(bf_df, ops.second_op, "timestamp_col") + + snapshot.assert_match(sql, "out.sql") + + def test_rstrip(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["string_col"]] sql = _apply_unary_op(bf_df, ops.StrRstripOp(" "), "string_col") @@ -342,6 +405,20 @@ def test_sqrt(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_str_get(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["string_col"]] + sql = _apply_unary_op(bf_df, ops.StrGetOp(1), "string_col") + + snapshot.assert_match(sql, "out.sql") + + +def test_str_slice(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["string_col"]] + sql = _apply_unary_op(bf_df, ops.StrSliceOp(1, 3), "string_col") + + snapshot.assert_match(sql, "out.sql") + + def test_str_contains(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["string_col"]] sql = _apply_unary_op(bf_df, ops.StrContainsOp("e"), "string_col") @@ -377,6 +454,13 @@ def test_iso_week(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_iso_year(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + sql = _apply_unary_op(bf_df, ops.iso_year_op, "timestamp_col") + + snapshot.assert_match(sql, "out.sql") + + def test_isnull(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["float64_col"]] sql = _apply_unary_op(bf_df, ops.isnull_op, "float64_col") @@ -419,6 +503,20 @@ def test_tanh(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_time(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + sql = _apply_unary_op(bf_df, ops.time_op, "timestamp_col") + + snapshot.assert_match(sql, "out.sql") + + +def test_timedelta_floor(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["int64_col"]] + sql = _apply_unary_op(bf_df, ops.timedelta_floor_op, "int64_col") + + snapshot.assert_match(sql, "out.sql") + + def test_json_extract(json_types_df: bpd.DataFrame, snapshot): bf_df = json_types_df[["json_col"]] sql = _apply_unary_op(bf_df, ops.JSONExtract(json_path="$"), "json_col") @@ -480,3 +578,10 @@ def test_upper(scalar_types_df: bpd.DataFrame, snapshot): sql = _apply_unary_op(bf_df, ops.upper_op, "string_col") snapshot.assert_match(sql, "out.sql") + + +def test_year(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col"]] + sql = _apply_unary_op(bf_df, ops.year_op, "timestamp_col") + + snapshot.assert_match(sql, "out.sql")