diff --git a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py index ddaf04ae97..42edf0158a 100644 --- a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py +++ b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py @@ -262,6 +262,27 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.func("ST_BOUNDARY", expr.expr) +@UNARY_OP_REGISTRATION.register(ops.GeoStBufferOp) +def _(op: ops.GeoStBufferOp, expr: TypedExpr) -> sge.Expression: + return sge.func( + "ST_BUFFER", + expr.expr, + sge.convert(op.buffer_radius), + sge.convert(op.num_seg_quarter_circle), + sge.convert(op.use_spheroid), + ) + + +@UNARY_OP_REGISTRATION.register(ops.geo_st_centroid_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.func("ST_CENTROID", expr.expr) + + +@UNARY_OP_REGISTRATION.register(ops.geo_st_convexhull_op) +def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: + return sge.func("ST_CONVEXHULL", expr.expr) + + @UNARY_OP_REGISTRATION.register(ops.geo_st_geogfromtext_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.func("SAFE.ST_GEOGFROMTEXT", expr.expr) @@ -434,6 +455,17 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.Lower(this=expr.expr) +@UNARY_OP_REGISTRATION.register(ops.MapOp) +def _(op: ops.MapOp, expr: TypedExpr) -> sge.Expression: + return sge.Case( + this=expr.expr, + ifs=[ + sge.If(this=sge.convert(key), true=sge.convert(value)) + for key, value in op.mappings + ], + ) + + @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) diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_buffer/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_buffer/out.sql new file mode 100644 index 0000000000..9669c39a9f --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_buffer/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_BUFFER(`bfcol_0`, 1.0, 8.0, FALSE) 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_centroid/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_centroid/out.sql new file mode 100644 index 0000000000..97867318ad --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_centroid/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_CENTROID(`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_convexhull/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_convexhull/out.sql new file mode 100644 index 0000000000..8bb5801173 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_convexhull/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_CONVEXHULL(`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_map/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_map/out.sql new file mode 100644 index 0000000000..a17d6584ce --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_map/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 + *, + CASE `bfcol_0` WHEN 'value1' THEN 'mapped1' END 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/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index 4a5b586c77..dae1599636 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py @@ -174,6 +174,27 @@ def test_geo_st_boundary(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_geo_st_buffer(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["geography_col"]] + sql = _apply_unary_op(bf_df, ops.GeoStBufferOp(1.0, 8.0, False), "geography_col") + + snapshot.assert_match(sql, "out.sql") + + +def test_geo_st_centroid(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["geography_col"]] + sql = _apply_unary_op(bf_df, ops.geo_st_centroid_op, "geography_col") + + snapshot.assert_match(sql, "out.sql") + + +def test_geo_st_convexhull(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["geography_col"]] + sql = _apply_unary_op(bf_df, ops.geo_st_convexhull_op, "geography_col") + + snapshot.assert_match(sql, "out.sql") + + def test_geo_st_geogfromtext(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["string_col"]] sql = _apply_unary_op(bf_df, ops.geo_st_geogfromtext_op, "string_col") @@ -370,6 +391,15 @@ def test_lower(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(sql, "out.sql") +def test_map(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["string_col"]] + sql = _apply_unary_op( + bf_df, ops.MapOp(mappings=(("value1", "mapped1"),)), "string_col" + ) + + snapshot.assert_match(sql, "out.sql") + + def test_lstrip(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["string_col"]] sql = _apply_unary_op(bf_df, ops.StrLstripOp(" "), "string_col")