Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions bigframes/core/compile/ibis_compiler/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
159 changes: 159 additions & 0 deletions bigframes/core/compile/ibis_compiler/operations/geo_ops.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,159 @@
# 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 typing import cast

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.udf as ibis_udf

from bigframes.core.compile.ibis_compiler import scalar_op_compiler
from bigframes.operations import geo_ops as ops

register_unary_op = scalar_op_compiler.scalar_op_compiler.register_unary_op
register_binary_op = scalar_op_compiler.scalar_op_compiler.register_binary_op


# Geo Ops
@register_unary_op(ops.geo_area_op)
def geo_area_op_impl(x: ibis_types.Value):
return cast(ibis_types.GeoSpatialValue, x).area()


@register_unary_op(ops.geo_st_astext_op)
def geo_st_astext_op_impl(x: ibis_types.Value):
return cast(ibis_types.GeoSpatialValue, x).as_text()


@register_unary_op(ops.geo_st_boundary_op, pass_op=False)
def geo_st_boundary_op_impl(x: ibis_types.Value):
return st_boundary(x)


@register_unary_op(ops.GeoStBufferOp, pass_op=True)
def geo_st_buffer_op_impl(x: ibis_types.Value, op: ops.GeoStBufferOp):
return st_buffer(
x,
op.buffer_radius,
op.num_seg_quarter_circle,
op.use_spheroid,
)


@register_unary_op(ops.geo_st_centroid_op, pass_op=False)
def geo_st_centroid_op_impl(x: ibis_types.Value):
return cast(ibis_types.GeoSpatialValue, x).centroid()


@register_unary_op(ops.geo_st_convexhull_op, pass_op=False)
def geo_st_convexhull_op_impl(x: ibis_types.Value):
return st_convexhull(x)


@register_binary_op(ops.geo_st_difference_op, pass_op=False)
def geo_st_difference_op_impl(x: ibis_types.Value, y: ibis_types.Value):
return cast(ibis_types.GeoSpatialValue, x).difference(
cast(ibis_types.GeoSpatialValue, y)
)


@register_binary_op(ops.GeoStDistanceOp, pass_op=True)
def geo_st_distance_op_impl(
x: ibis_types.Value, y: ibis_types.Value, op: ops.GeoStDistanceOp
):
return st_distance(x, y, op.use_spheroid)


@register_unary_op(ops.geo_st_geogfromtext_op)
def geo_st_geogfromtext_op_impl(x: ibis_types.Value):
# Ibis doesn't seem to provide a dedicated method to cast from string to geography,
# so we use a BigQuery scalar function, st_geogfromtext(), directly.
return st_geogfromtext(x)


@register_binary_op(ops.geo_st_geogpoint_op, pass_op=False)
def geo_st_geogpoint_op_impl(x: ibis_types.Value, y: ibis_types.Value):
return cast(ibis_types.NumericValue, x).point(cast(ibis_types.NumericValue, y))


@register_binary_op(ops.geo_st_intersection_op, pass_op=False)
def geo_st_intersection_op_impl(x: ibis_types.Value, y: ibis_types.Value):
return cast(ibis_types.GeoSpatialValue, x).intersection(
cast(ibis_types.GeoSpatialValue, y)
)


@register_unary_op(ops.geo_st_isclosed_op, pass_op=False)
def geo_st_isclosed_op_impl(x: ibis_types.Value):
return st_isclosed(x)


@register_unary_op(ops.geo_x_op)
def geo_x_op_impl(x: ibis_types.Value):
return cast(ibis_types.GeoSpatialValue, x).x()


@register_unary_op(ops.GeoStLengthOp, pass_op=True)
def geo_length_op_impl(x: ibis_types.Value, op: ops.GeoStLengthOp):
# Call the st_length UDF defined in this file (or imported)
return st_length(x, op.use_spheroid)


@register_unary_op(ops.geo_y_op)
def geo_y_op_impl(x: ibis_types.Value):
return cast(ibis_types.GeoSpatialValue, x).y()


@ibis_udf.scalar.builtin
def st_convexhull(x: ibis_dtypes.geography) -> ibis_dtypes.geography: # type: ignore
"""ST_CONVEXHULL"""
...


@ibis_udf.scalar.builtin
def st_geogfromtext(a: str) -> ibis_dtypes.geography: # type: ignore
"""Convert string to geography."""


@ibis_udf.scalar.builtin
def st_boundary(a: ibis_dtypes.geography) -> ibis_dtypes.geography: # type: ignore
"""Find the boundary of a geography."""


@ibis_udf.scalar.builtin
def st_buffer(
geography: ibis_dtypes.geography, # type: ignore
buffer_radius: ibis_dtypes.Float64,
num_seg_quarter_circle: ibis_dtypes.Float64,
use_spheroid: ibis_dtypes.Boolean,
) -> ibis_dtypes.geography: # type: ignore
...


@ibis_udf.scalar.builtin
def st_distance(a: ibis_dtypes.geography, b: ibis_dtypes.geography, use_spheroid: bool) -> ibis_dtypes.float: # type: ignore
"""Convert string to geography."""


@ibis_udf.scalar.builtin
def st_length(geog: ibis_dtypes.geography, use_spheroid: bool) -> ibis_dtypes.float: # type: ignore
"""ST_LENGTH BQ builtin. This body is never executed."""
pass


@ibis_udf.scalar.builtin
def st_isclosed(a: ibis_dtypes.geography) -> ibis_dtypes.boolean: # type: ignore
"""Checks if a geography is closed."""
134 changes: 0 additions & 134 deletions bigframes/core/compile/ibis_compiler/scalar_op_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -837,98 +837,6 @@ def normalize_op_impl(x: ibis_types.Value):
return result.cast(result_type)


# Geo Ops
@scalar_op_compiler.register_unary_op(ops.geo_area_op)
def geo_area_op_impl(x: ibis_types.Value):
return typing.cast(ibis_types.GeoSpatialValue, x).area()


@scalar_op_compiler.register_unary_op(ops.geo_st_astext_op)
def geo_st_astext_op_impl(x: ibis_types.Value):
return typing.cast(ibis_types.GeoSpatialValue, x).as_text()


@scalar_op_compiler.register_unary_op(ops.geo_st_boundary_op, pass_op=False)
def geo_st_boundary_op_impl(x: ibis_types.Value):
return st_boundary(x)


@scalar_op_compiler.register_unary_op(ops.GeoStBufferOp, pass_op=True)
def geo_st_buffer_op_impl(x: ibis_types.Value, op: ops.GeoStBufferOp):
return st_buffer(
x,
op.buffer_radius,
op.num_seg_quarter_circle,
op.use_spheroid,
)


@scalar_op_compiler.register_unary_op(ops.geo_st_centroid_op, pass_op=False)
def geo_st_centroid_op_impl(x: ibis_types.Value):
return typing.cast(ibis_types.GeoSpatialValue, x).centroid()


@scalar_op_compiler.register_unary_op(ops.geo_st_convexhull_op, pass_op=False)
def geo_st_convexhull_op_impl(x: ibis_types.Value):
return st_convexhull(x)


@scalar_op_compiler.register_binary_op(ops.geo_st_difference_op, pass_op=False)
def geo_st_difference_op_impl(x: ibis_types.Value, y: ibis_types.Value):
return typing.cast(ibis_types.GeoSpatialValue, x).difference(
typing.cast(ibis_types.GeoSpatialValue, y)
)


@scalar_op_compiler.register_binary_op(ops.GeoStDistanceOp, pass_op=True)
def geo_st_distance_op_impl(
x: ibis_types.Value, y: ibis_types.Value, op: ops.GeoStDistanceOp
):
return st_distance(x, y, op.use_spheroid)


@scalar_op_compiler.register_unary_op(ops.geo_st_geogfromtext_op)
def geo_st_geogfromtext_op_impl(x: ibis_types.Value):
# Ibis doesn't seem to provide a dedicated method to cast from string to geography,
# so we use a BigQuery scalar function, st_geogfromtext(), directly.
return st_geogfromtext(x)


@scalar_op_compiler.register_binary_op(ops.geo_st_geogpoint_op, pass_op=False)
def geo_st_geogpoint_op_impl(x: ibis_types.Value, y: ibis_types.Value):
return typing.cast(ibis_types.NumericValue, x).point(
typing.cast(ibis_types.NumericValue, y)
)


@scalar_op_compiler.register_binary_op(ops.geo_st_intersection_op, pass_op=False)
def geo_st_intersection_op_impl(x: ibis_types.Value, y: ibis_types.Value):
return typing.cast(ibis_types.GeoSpatialValue, x).intersection(
typing.cast(ibis_types.GeoSpatialValue, y)
)


@scalar_op_compiler.register_unary_op(ops.geo_st_isclosed_op, pass_op=False)
def geo_st_isclosed_op_impl(x: ibis_types.Value):
return st_isclosed(x)


@scalar_op_compiler.register_unary_op(ops.geo_x_op)
def geo_x_op_impl(x: ibis_types.Value):
return typing.cast(ibis_types.GeoSpatialValue, x).x()


@scalar_op_compiler.register_unary_op(ops.GeoStLengthOp, pass_op=True)
def geo_length_op_impl(x: ibis_types.Value, op: ops.GeoStLengthOp):
# Call the st_length UDF defined in this file (or imported)
return st_length(x, op.use_spheroid)


@scalar_op_compiler.register_unary_op(ops.geo_y_op)
def geo_y_op_impl(x: ibis_types.Value):
return typing.cast(ibis_types.GeoSpatialValue, x).y()


# Parameterized ops
@scalar_op_compiler.register_unary_op(ops.StructFieldOp, pass_op=True)
def struct_field_op_impl(x: ibis_types.Value, op: ops.StructFieldOp):
Expand Down Expand Up @@ -2092,17 +2000,6 @@ def _ibis_num(number: float):
return typing.cast(ibis_types.NumericValue, ibis_types.literal(number))


@ibis_udf.scalar.builtin
def st_convexhull(x: ibis_dtypes.geography) -> ibis_dtypes.geography: # type: ignore
"""ST_CONVEXHULL"""
...


@ibis_udf.scalar.builtin
def st_geogfromtext(a: str) -> ibis_dtypes.geography: # type: ignore
"""Convert string to geography."""


@ibis_udf.scalar.builtin
def timestamp(a: str) -> ibis_dtypes.timestamp: # type: ignore
"""Convert string to timestamp."""
Expand All @@ -2113,32 +2010,6 @@ def unix_millis(a: ibis_dtypes.timestamp) -> int: # type: ignore
"""Convert a timestamp to milliseconds"""


@ibis_udf.scalar.builtin
def st_boundary(a: ibis_dtypes.geography) -> ibis_dtypes.geography: # type: ignore
"""Find the boundary of a geography."""


@ibis_udf.scalar.builtin
def st_buffer(
geography: ibis_dtypes.geography, # type: ignore
buffer_radius: ibis_dtypes.Float64,
num_seg_quarter_circle: ibis_dtypes.Float64,
use_spheroid: ibis_dtypes.Boolean,
) -> ibis_dtypes.geography: # type: ignore
...


@ibis_udf.scalar.builtin
def st_distance(a: ibis_dtypes.geography, b: ibis_dtypes.geography, use_spheroid: bool) -> ibis_dtypes.float: # type: ignore
"""Convert string to geography."""


@ibis_udf.scalar.builtin
def st_length(geog: ibis_dtypes.geography, use_spheroid: bool) -> ibis_dtypes.float: # type: ignore
"""ST_LENGTH BQ builtin. This body is never executed."""
pass


@ibis_udf.scalar.builtin
def unix_micros(a: ibis_dtypes.timestamp) -> int: # type: ignore
"""Convert a timestamp to microseconds"""
Expand Down Expand Up @@ -2272,11 +2143,6 @@ def str_lstrip_op( # type: ignore[empty-body]
"""Remove leading and trailing characters."""


@ibis_udf.scalar.builtin
def st_isclosed(a: ibis_dtypes.geography) -> ibis_dtypes.boolean: # type: ignore
"""Checks if a geography is closed."""


@ibis_udf.scalar.builtin(name="rtrim")
def str_rstrip_op( # type: ignore[empty-body]
x: ibis_dtypes.String, to_strip: ibis_dtypes.String
Expand Down
13 changes: 8 additions & 5 deletions specs/2025-08-04-geoseries-scalars.md
Original file line number Diff line number Diff line change
Expand Up @@ -267,11 +267,14 @@ Raster functions: Functions for analyzing geospatial rasters using geographies.
- [ ] **Export the new operation:**
- [ ] In `bigframes/operations/__init__.py`, import your new operation dataclass and add it to the `__all__` list.
- [ ] **Implement the compilation logic:**
- [ ] In `bigframes/core/compile/scalar_op_compiler.py`:
- [ ] If the BigQuery function has a direct equivalent in Ibis, you can often reuse an existing Ibis method.
- [ ] If not, define a new Ibis UDF using `@ibis_udf.scalar.builtin` to map to the specific BigQuery function signature.
- [ ] Create a new compiler implementation function (e.g., `geo_length_op_impl`).
- [ ] Register this function to your operation dataclass using `@scalar_op_compiler.register_unary_op` or `@scalar_op_compiler.register_binary_op`.
- [ ] In `bigframes/core/compile/ibis_compiler/operations/geo_ops.py`:
- [ ] If the BigQuery function has a direct equivalent in Ibis, you can often reuse an existing Ibis method.
- [ ] If not, define a new Ibis UDF using `@ibis_udf.scalar.builtin` to map to the specific BigQuery function signature.
- [ ] Create a new compiler implementation function (e.g., `geo_length_op_impl`).
- [ ] Register this function to your operation dataclass using `@register_unary_op` or `@register_binary_op`.
- [ ] In `bigframes/core/compile/sqlglot/expressions/geo_ops.py`:
- [ ] Create a new compiler implementation function that generates the appropriate `sqlglot.exp` expression.
- [ ] Register this function to your operation dataclass using `@register_unary_op` or `@register_binary_op`.
- [ ] **Implement the user-facing function or property:**
- [ ] For a `bigframes.bigquery` function:
- [ ] In `bigframes/bigquery/_operations/geo.py`, create the user-facing function (e.g., `st_length`).
Expand Down