From 2375c070f0d66b2ea2c3726dc600e199566b21ad Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Fri, 22 Aug 2025 19:08:21 +0000 Subject: [PATCH 1/4] GeoStBufferOp compiler --- .../compile/sqlglot/expressions/unary_compiler.py | 11 +++++++++++ .../test_unary_compiler/test_geo_st_buffer/out.sql | 13 +++++++++++++ .../sqlglot/expressions/test_unary_compiler.py | 7 +++++++ 3 files changed, 31 insertions(+) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_buffer/out.sql diff --git a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py index ddaf04ae97..40606257e8 100644 --- a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py +++ b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py @@ -262,6 +262,17 @@ 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_geogfromtext_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.func("SAFE.ST_GEOGFROMTEXT", 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/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index 4a5b586c77..38078db83f 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,13 @@ 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_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") From 3d2a33b6c7661210c15a93a52c3b2d39554cb3ec Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Fri, 22 Aug 2025 20:14:09 +0000 Subject: [PATCH 2/4] geo_st_centroid_op compiler --- .../compile/sqlglot/expressions/unary_compiler.py | 5 +++++ .../test_geo_st_centroid/out.sql | 13 +++++++++++++ .../sqlglot/expressions/test_unary_compiler.py | 7 +++++++ 3 files changed, 25 insertions(+) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_centroid/out.sql diff --git a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py index 40606257e8..a7eee679df 100644 --- a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py +++ b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py @@ -273,6 +273,11 @@ def _(op: ops.GeoStBufferOp, expr: TypedExpr) -> sge.Expression: ) +@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_geogfromtext_op) def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression: return sge.func("SAFE.ST_GEOGFROMTEXT", expr.expr) 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/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index 38078db83f..0181abc7c0 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py @@ -181,6 +181,13 @@ def test_geo_st_buffer(scalar_types_df: bpd.DataFrame, snapshot): 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_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") From 364aaef28e1603cc8dd58a6acb41d44d09cd424c Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Fri, 22 Aug 2025 20:26:51 +0000 Subject: [PATCH 3/4] geo_st_convexhull_op compiler --- .../compile/sqlglot/expressions/unary_compiler.py | 5 +++++ .../test_geo_st_convexhull/out.sql | 13 +++++++++++++ .../sqlglot/expressions/test_unary_compiler.py | 7 +++++++ 3 files changed, 25 insertions(+) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_convexhull/out.sql diff --git a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py index a7eee679df..611cb22672 100644 --- a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py +++ b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py @@ -278,6 +278,11 @@ 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) 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/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index 0181abc7c0..97d974704e 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py @@ -188,6 +188,13 @@ def test_geo_st_centroid(scalar_types_df: bpd.DataFrame, snapshot): 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") From 7d7b5d49af6db5735e87f02635e74a0b2521a6a1 Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Fri, 22 Aug 2025 20:37:10 +0000 Subject: [PATCH 4/4] MapOp compiler --- .../compile/sqlglot/expressions/unary_compiler.py | 11 +++++++++++ .../snapshots/test_unary_compiler/test_map/out.sql | 13 +++++++++++++ .../sqlglot/expressions/test_unary_compiler.py | 9 +++++++++ 3 files changed, 33 insertions(+) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_map/out.sql diff --git a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py index 611cb22672..42edf0158a 100644 --- a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py +++ b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py @@ -455,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_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 97d974704e..dae1599636 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py @@ -391,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")