fix the sqlgot compiler

Author: tswast

bigframes/bigquery/_operations/geo.py

@@ -685,7 +685,7 @@ def st_regionstats(
     band: Optional[str] = None,
     include: Optional[str] = None,
     options: Optional[Mapping[str, Union[str, int, float]]] = None,
-) -> bigframes.dataframe.DataFrame:
+) -> bigframes.series.Series:
     """Returns statistics summarizing the pixel values of the raster image
     referenced by raster_id that intersect with geography.
 
@@ -726,14 +726,13 @@ def st_regionstats(
             documentation for a list of available options.
 
     Returns:
-        bigframes.dataframe.DataFrame:
-            A dataframe containing the computed statistics.
+        bigframes.pandas.Series:
+            A STRUCT Series containing the computed statistics.
     """
     op = ops.StRegionStatsOp(
         raster_id=raster_id,
         band=band,
         include=include,
         options=json.dumps(options) if options else None,
     )
-    df = geography._apply_unary_op(op)
-    return df[df.columns[0]].struct.explode()
+    return geography._apply_unary_op(op)

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

@@ -89,15 +89,11 @@ def _(
     geography: TypedExpr,
     op: ops.StRegionStatsOp,
 ):
-    args = [geography.expr]  # TODO: get raster, band, include from op.
+    args = [geography.expr, sge.convert(op.raster_id)]
+    if op.band:
+        args.append(sge.Kwarg(this="band", expression=sge.convert(op.band)))
+    if op.include:
+        args.append(sge.Kwarg(this="include", expression=sge.convert(op.include)))
     if op.options:
-        args.append(
-            sge.Anonymous(
-                this="_",
-                expressions=[
-                    sge.Identifier(this="OPTIONS"),
-                    sge.Anonymous(this="JSON", expressions=[sge.convert(op.options)]),
-                ],
-            )
-        )
+        args.append(sge.Kwarg(this="options", expression=sge.convert(op.options)))
     return sge.func("ST_REGIONSTATS", *args)

tests/unit/core/compile/sqlglot/snapshots/test_geo_compiler/test_st_regionstats/out.sql

@@ -1,44 +1,36 @@
-WITH `bfcte_1` AS (
+WITH `bfcte_0` AS (
   SELECT
     *
-  FROM UNNEST(ARRAY<STRUCT<`bfcol_0` INT64, `bfcol_1` STRING>>[STRUCT(0, 'POINT(1 1)')])
-), `bfcte_0` AS (
-  SELECT
-    *
-  FROM UNNEST(ARRAY<STRUCT<`bfcol_2` INT64, `bfcol_3` STRING>>[STRUCT(0, 'raster_uri')])
-), `bfcte_2` AS (
-  SELECT
-    `bfcol_2` AS `bfcol_4`,
-    `bfcol_3` AS `bfcol_5`
-  FROM `bfcte_0`
-), `bfcte_3` AS (
-  SELECT
-    *
-  FROM `bfcte_1`
-  LEFT JOIN `bfcte_2`
-    ON COALESCE(`bfcol_0`, 0) = COALESCE(`bfcol_4`, 0)
-    AND COALESCE(`bfcol_0`, 1) = COALESCE(`bfcol_4`, 1)
-), `bfcte_4` AS (
+  FROM UNNEST(ARRAY<STRUCT<`bfcol_0` STRING, `bfcol_1` INT64>>[STRUCT('POINT(1 1)', 0)])
+), `bfcte_1` AS (
   SELECT
     *,
-    ST_REGIONSTATS(`bfcol_1`, `bfcol_5`, 'band1', _(OPTIONS, JSON('{"scale": 100}'))) AS `bfcol_8`
-  FROM `bfcte_3`
-), `bfcte_5` AS (
+    ST_REGIONSTATS(
+      `bfcol_0`,
+      'ee://some/raster/uri',
+      band => 'band1',
+      include => 'some equation',
+      options => '{"scale": 100}'
+    ) AS `bfcol_2`
+  FROM `bfcte_0`
+), `bfcte_2` AS (
   SELECT
     *,
-    `bfcol_8`.`min` AS `bfcol_10`,
-    `bfcol_8`.`max` AS `bfcol_11`,
-    `bfcol_8`.`sum` AS `bfcol_12`,
-    `bfcol_8`.`count` AS `bfcol_13`,
-    `bfcol_8`.`mean` AS `bfcol_14`,
-    `bfcol_8`.`area` AS `bfcol_15`
-  FROM `bfcte_4`
+    `bfcol_2`.`min` AS `bfcol_5`,
+    `bfcol_2`.`max` AS `bfcol_6`,
+    `bfcol_2`.`sum` AS `bfcol_7`,
+    `bfcol_2`.`count` AS `bfcol_8`,
+    `bfcol_2`.`mean` AS `bfcol_9`,
+    `bfcol_2`.`area` AS `bfcol_10`
+  FROM `bfcte_1`
 )
 SELECT
-  `bfcol_10` AS `min`,
-  `bfcol_11` AS `max`,
-  `bfcol_12` AS `sum`,
-  `bfcol_13` AS `count`,
-  `bfcol_14` AS `mean`,
-  `bfcol_15` AS `area`
-FROM `bfcte_5`
+  `bfcol_5` AS `min`,
+  `bfcol_6` AS `max`,
+  `bfcol_7` AS `sum`,
+  `bfcol_8` AS `count`,
+  `bfcol_9` AS `mean`,
+  `bfcol_10` AS `area`
+FROM `bfcte_2`
+ORDER BY
+  `bfcol_1` ASC NULLS LAST

tests/unit/core/compile/sqlglot/test_geo_compiler.py

@@ -16,14 +16,18 @@
 
 import bigframes.bigquery as bbq
 import bigframes.geopandas as gpd
-import bigframes.pandas as bpd
 
 pytest.importorskip("pytest_snapshot")
 
 
 def test_st_regionstats(compiler_session, snapshot):
     geos = gpd.GeoSeries(["POINT(1 1)"], session=compiler_session)
-    rasters = bpd.Series(["raster_uri"], dtype="string", session=compiler_session)
-    df = bbq.st_regionstats(geos, rasters, "band1", {"scale": 100})
-    assert "area" in df.columns
-    snapshot.assert_match(df.sql, "out.sql")
+    result = bbq.st_regionstats(
+        geos,
+        "ee://some/raster/uri",
+        band="band1",
+        include="some equation",
+        options={"scale": 100},
+    )
+    assert "area" in result.struct.dtypes.index
+    snapshot.assert_match(result.struct.explode().sql, "out.sql")