attempt at ibis compiler

Author: tswast

bigframes/core/compile/ibis_compiler/__init__.py

@@ -21,4 +21,5 @@
 from __future__ import annotations
 
 import bigframes.core.compile.ibis_compiler.operations.generic_ops  # noqa: F401
+import bigframes.core.compile.ibis_compiler.operations.geo_ops  # noqa: F401
 import bigframes.core.compile.ibis_compiler.scalar_op_registry  # noqa: F401

bigframes/core/compile/ibis_compiler/operations/geo_ops.py

@@ -0,0 +1,55 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from bigframes_vendored import ibis
+from bigframes_vendored.ibis.expr import types as ibis_types
+import bigframes_vendored.ibis.expr.datatypes as ibis_dtypes
+import bigframes_vendored.ibis.expr.operations.geospatial as ibis_geo
+
+from bigframes.core.compile.ibis_compiler import scalar_op_compiler
+from bigframes.operations import geo_ops
+
+register_unary_op = scalar_op_compiler.scalar_op_compiler.register_unary_op
+
+
+@register_unary_op(geo_ops.StRegionStatsOp, pass_op=True)
+def st_regionstats(
+    geography: ibis_types.Value,
+    op: geo_ops.StRegionStatsOp,
+):
+
+    if op.band:
+        band = ibis.literal(op.band, type=ibis_dtypes.string())
+    else:
+        band = None
+
+    if op.include:
+        include = ibis.literal(op.include, type=ibis_dtypes.string())
+    else:
+        include = None
+
+    if op.options:
+        options = ibis.literal(op.options, type=ibis_dtypes.string())
+    else:
+        options = None
+
+    return ibis_geo.GeoRegionStats(
+        arg=geography,
+        raster_id=ibis.literal(op.raster_id, type=ibis_dtypes.string()),
+        band=band,
+        include=include,
+        options=options,
+    )

bigframes/core/compile/ibis_compiler/scalar_op_compiler.py

@@ -26,7 +26,7 @@
 from bigframes.core import agg_expressions, ordering
 import bigframes.core.compile.ibis_types
 import bigframes.core.expression as ex
-from bigframes.operations import geo_ops, numeric_ops
+from bigframes.operations import numeric_ops
 
 if TYPE_CHECKING:
     import bigframes.operations as ops
@@ -286,19 +286,3 @@ def isnanornull(arg):
 @scalar_op_compiler.register_unary_op(numeric_ops.isfinite_op)
 def isfinite(arg):
     return arg.isinf().negate() & arg.isnan().negate()
-
-
-@scalar_op_compiler.register_unary_op(geo_ops.StRegionStatsOp, pass_op=True)
-def st_regionstats(
-    geography: ibis_types.Value,
-    op: geo_ops.StRegionStatsOp,
-):
-    args = [geography]  # TODO: get band, include, and other properies from op.
-    if op.options:
-        args.append(bigframes_vendored.ibis.literal(op.options, type="json"))
-    # TODO: We may need a custom ibis op so that we can pass arguments by name instead of position.
-    return bigframes_vendored.ibis.remote_function(
-        "st_regionstats",
-        args,
-        output_type="struct<min: float, max: float, sum: float, count: int, mean: float, area: float>",  # type: ignore
-    )

tests/system/small/geopandas/test_geoseries.py

@@ -30,6 +30,7 @@
     Polygon,
 )
 
+import bigframes.bigquery
 import bigframes.geopandas
 import bigframes.pandas
 import bigframes.series
@@ -557,9 +558,8 @@ def test_st_regionstats(session: bigframes.session.Session):
         ]
     )
     geos = bigframes.geopandas.GeoSeries([polygon], session=session)
-    rasters = bigframes.pandas.Series([raster], dtype="string", session=session)
     result = bigframes.bigquery.st_regionstats(
-        geos, rasters, "loss", options={"scale": 1000}
+        geos, raster, "loss", options={"scale": 1000}
     ).to_pandas()
     assert result is not None
     assert "mean" in result.columns

third_party/bigframes_vendored/ibis/backends/sql/compilers/bigquery/__init__.py

@@ -261,6 +261,16 @@ def visit_BoundingBox(self, op, *, arg):
 
     visit_GeoXMax = visit_GeoXMin = visit_GeoYMax = visit_GeoYMin = visit_BoundingBox
 
+    def visit_GeoRegionStats(self, op, *, arg, raster_id, band, include, options):
+        args = [arg, raster_id]
+        if op.band:
+            args.append(sge.Kwarg(this="band", expression=band))
+        if op.include:
+            args.append(sge.Kwarg(this="include", expression=include))
+        if op.options:
+            args.append(sge.Kwarg(this="options", expression=options))
+        return sge.func("ST_REGIONSTATS", *args)
+
     def visit_GeoSimplify(self, op, *, arg, tolerance, preserve_collapsed):
         if (
             not isinstance(op.preserve_collapsed, ops.Literal)

third_party/bigframes_vendored/ibis/expr/operations/geospatial.py

@@ -343,6 +343,28 @@ class GeoNRings(GeoSpatialUnOp):
     dtype = dt.int64
 
 
+@public
+class GeoRegionStats(GeoSpatialUnOp):
+    """Returns results of ST_REGIONSTATS."""
+
+    raster_id: Value[dt.String]
+    band: Value[dt.String]
+    include: Value[dt.String]
+    options: Value[dt.String]
+
+    dtype = dt.Struct(
+        fields={
+            "count": dt.int64,
+            "min": dt.float64,
+            "max": dt.float64,
+            "stdDev": dt.float64,
+            "sum": dt.float64,
+            "mean": dt.float64,
+            "area": dt.float64,
+        }
+    )
+
+
 @public
 class GeoSRID(GeoSpatialUnOp):
     """Returns the spatial reference identifier for the ST_Geometry."""