Merge remote-tracking branch 'github/main' into local_series_index

Author: TrevorBergeron

bigframes/_config/experiment_options.py

@@ -26,7 +26,6 @@ def __init__(self):
         self._semantic_operators: bool = False
         self._ai_operators: bool = False
         self._blob: bool = False
-        self._udf: bool = False
 
     @property
     def semantic_operators(self) -> bool:
@@ -68,17 +67,3 @@ def blob(self, value: bool):
             )
             warnings.warn(msg, category=bfe.PreviewWarning)
         self._blob = value
-
-    @property
-    def udf(self) -> bool:
-        return self._udf
-
-    @udf.setter
-    def udf(self, value: bool):
-        if value is True:
-            msg = bfe.format_message(
-                "BigFrames managed function (udf) is still under experiments. "
-                "It may not work and subject to change in the future."
-            )
-            warnings.warn(msg, category=bfe.PreviewWarning)
-        self._udf = value

bigframes/bigquery/__init__.py

@@ -27,7 +27,7 @@
     unix_millis,
     unix_seconds,
 )
-from bigframes.bigquery._operations.geo import st_area, st_difference
+from bigframes.bigquery._operations.geo import st_area, st_difference, st_intersection
 from bigframes.bigquery._operations.json import (
     json_extract,
     json_extract_array,
@@ -49,6 +49,7 @@
     # geo ops
     "st_area",
     "st_difference",
+    "st_intersection",
     # json ops
     "json_set",
     "json_extract",

bigframes/bigquery/_operations/geo.py

@@ -211,3 +211,95 @@ def st_difference(
             in other.
     """
     return series._apply_binary_op(other, ops.geo_st_difference_op)
+
+
+def st_intersection(
+    series: bigframes.series.Series, other: bigframes.series.Series
+) -> bigframes.series.Series:
+    """
+    Returns a `GEOGRAPHY` that represents the point set intersection of the two
+    input `GEOGRAPHYs`. Thus, every point in the intersection appears in both
+    `geography_1` and `geography_2`.
+
+    .. note::
+        BigQuery's Geography functions, like `st_intersection`, interpret the geometry
+        data type as a point set on the Earth's surface. A point set is a set
+        of points, lines, and polygons on the WGS84 reference spheroid, with
+        geodesic edges. See: https://cloud.google.com/bigquery/docs/geospatial-data
+
+    **Examples:**
+
+        >>> import bigframes as bpd
+        >>> import bigframes.bigquery as bbq
+        >>> import bigframes.geopandas
+        >>> from shapely.geometry import Polygon, LineString, Point
+        >>> bpd.options.display.progress_bar = None
+
+    We can check two GeoSeries against each other, row by row.
+
+        >>> s1 = bigframes.geopandas.GeoSeries(
+        ...    [
+        ...        Polygon([(0, 0), (2, 2), (0, 2)]),
+        ...        Polygon([(0, 0), (2, 2), (0, 2)]),
+        ...        LineString([(0, 0), (2, 2)]),
+        ...        LineString([(2, 0), (0, 2)]),
+        ...        Point(0, 1),
+        ...    ],
+        ... )
+        >>> s2 = bigframes.geopandas.GeoSeries(
+        ...    [
+        ...        Polygon([(0, 0), (1, 1), (0, 1)]),
+        ...        LineString([(1, 0), (1, 3)]),
+        ...        LineString([(2, 0), (0, 2)]),
+        ...        Point(1, 1),
+        ...        Point(0, 1),
+        ...    ],
+        ...    index=range(1, 6),
+        ... )
+
+        >>> s1
+        0    POLYGON ((0 0, 2 2, 0 2, 0 0))
+        1    POLYGON ((0 0, 2 2, 0 2, 0 0))
+        2             LINESTRING (0 0, 2 2)
+        3             LINESTRING (2 0, 0 2)
+        4                       POINT (0 1)
+        dtype: geometry
+
+        >>> s2
+        1    POLYGON ((0 0, 1 1, 0 1, 0 0))
+        2             LINESTRING (1 0, 1 3)
+        3             LINESTRING (2 0, 0 2)
+        4                       POINT (1 1)
+        5                       POINT (0 1)
+        dtype: geometry
+
+        >>> bbq.st_intersection(s1, s2)
+        0                                    None
+        1    POLYGON ((0 0, 0.99954 1, 0 1, 0 0))
+        2                       POINT (1 1.00046)
+        3                   LINESTRING (2 0, 0 2)
+        4                GEOMETRYCOLLECTION EMPTY
+        5                                    None
+        dtype: geometry
+
+    We can also do intersection of each geometry and a single shapely geometry:
+
+        >>> bbq.st_intersection(s1, bigframes.geopandas.GeoSeries([Polygon([(0, 0), (1, 1), (0, 1)])]))
+        0    POLYGON ((0 0, 0.99954 1, 0 1, 0 0))
+        1                                    None
+        2                                    None
+        3                                    None
+        4                                    None
+        dtype: geometry
+
+    Args:
+        other (GeoSeries or geometric object):
+            The Geoseries (elementwise) or geometric object to find the
+            intersection with.
+
+    Returns:
+        bigframes.geopandas.GeoSeries:
+            The Geoseries (elementwise) of the intersection of points in
+            each aligned geometry with other.
+    """
+    return series._apply_binary_op(other, ops.geo_st_intersection_op)

bigframes/core/compile/compiled.py

@@ -665,9 +665,11 @@ def _join_condition(
 
 
 def _as_groupable(value: ibis_types.Value):
-    # Some types need to be converted to string to enable groupby
-    if value.type().is_float64() or value.type().is_geospatial():
+    # Some types need to be converted to another type to enable groupby
+    if value.type().is_float64():
         return value.cast(ibis_dtypes.str)
+    elif value.type().is_geospatial():
+        return typing.cast(ibis_types.GeoSpatialColumn, value).as_binary()
     elif value.type().is_json():
         return scalar_op_compiler.to_json_string(value)
     else:

bigframes/core/compile/scalar_op_compiler.py

@@ -1037,6 +1037,13 @@ def geo_st_geogpoint_op_impl(x: ibis_types.Value, y: ibis_types.Value):
     )
 
 
+@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_x_op)
 def geo_x_op_impl(x: ibis_types.Value):
     return typing.cast(ibis_types.GeoSpatialValue, x).x()

bigframes/dataframe.py

@@ -4170,8 +4170,10 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
         # to the applied function should be a Series, not a scalar.
 
         if utils.get_axis_number(axis) == 1:
-            msg = bfe.format_message("axis=1 scenario is in preview.")
-            warnings.warn(msg, category=bfe.PreviewWarning)
+            msg = bfe.format_message(
+                "DataFrame.apply with parameter axis=1 scenario is in preview."
+            )
+            warnings.warn(msg, category=bfe.FunctionAxisOnePreviewWarning)
 
             if not hasattr(func, "bigframes_bigquery_function"):
                 raise ValueError(

bigframes/exceptions.py

@@ -95,6 +95,10 @@ class ObsoleteVersionWarning(Warning):
     """The BigFrames version is too old."""
 
 
+class FunctionAxisOnePreviewWarning(PreviewWarning):
+    """Remote Function and Managed UDF with axis=1 preview."""
+
+
 def format_message(message: str, fill: bool = True):
     """Formats a warning message with ANSI color codes for the warning color.
 

bigframes/functions/_function_client.py

@@ -53,6 +53,9 @@
     }
 )
 
+# BQ managed functions (@udf) currently only support Python 3.11.
+_MANAGED_FUNC_PYTHON_VERSION = "python-3.11"
+
 
 class FunctionClient:
     # Wait time (in seconds) for an IAM binding to take effect after creation.
@@ -193,11 +196,22 @@ def provision_bq_managed_function(
         name,
         packages,
         is_row_processor,
+        *,
+        capture_references=False,
     ):
         """Create a BigQuery managed function."""
-        import cloudpickle
 
-        pickled = cloudpickle.dumps(func)
+        # TODO(b/406283812): Expose the capability to pass down
+        # capture_references=True in the public udf API.
+        if (
+            capture_references
+            and (python_version := _utils.get_python_version())
+            != _MANAGED_FUNC_PYTHON_VERSION
+        ):
+            raise bf_formatting.create_exception_with_feedback_link(
+                NotImplementedError,
+                f"Capturing references for udf is currently supported only in Python version {_MANAGED_FUNC_PYTHON_VERSION}, you are running {python_version}.",
+            )
 
         # Create BQ managed function.
         bq_function_args = []
@@ -209,13 +223,15 @@ def provision_bq_managed_function(
             bq_function_args.append(f"{name_} {type_}")
 
         managed_function_options = {
-            "runtime_version": _utils.get_python_version(),
+            "runtime_version": _MANAGED_FUNC_PYTHON_VERSION,
             "entry_point": "bigframes_handler",
         }
 
         # Augment user package requirements with any internal package
         # requirements.
-        packages = _utils._get_updated_package_requirements(packages, is_row_processor)
+        packages = _utils._get_updated_package_requirements(
+            packages, is_row_processor, capture_references
+        )
         if packages:
             managed_function_options["packages"] = packages
         managed_function_options_str = self._format_function_options(
@@ -235,20 +251,45 @@ def provision_bq_managed_function(
         persistent_func_id = (
             f"`{self._gcp_project_id}.{self._bq_dataset}`.{bq_function_name}"
         )
-        create_function_ddl = textwrap.dedent(
-            f"""
-            CREATE OR REPLACE FUNCTION {persistent_func_id}({','.join(bq_function_args)})
-            RETURNS {bq_function_return_type}
-            LANGUAGE python
-            OPTIONS ({managed_function_options_str})
-            AS r'''
+
+        udf_name = func.__name__
+        if capture_references:
+            # This code path ensures that if the udf body contains any
+            # references to variables and/or imports outside the body, they are
+            # captured as well.
             import cloudpickle
-            udf = cloudpickle.loads({pickled})
-            def bigframes_handler(*args):
-                return udf(*args)
-            '''
-        """
-        ).strip()
+
+            pickled = cloudpickle.dumps(func)
+            udf_code = textwrap.dedent(
+                f"""
+                import cloudpickle
+                {udf_name} = cloudpickle.loads({pickled})
+            """
+            )
+        else:
+            # This code path ensures that if the udf body is self contained,
+            # i.e. there are no references to variables or imports outside the
+            # body.
+            udf_code = textwrap.dedent(inspect.getsource(func))
+            udf_code = udf_code[udf_code.index("def") :]
+
+        create_function_ddl = (
+            textwrap.dedent(
+                f"""
+                CREATE OR REPLACE FUNCTION {persistent_func_id}({','.join(bq_function_args)})
+                RETURNS {bq_function_return_type}
+                LANGUAGE python
+                OPTIONS ({managed_function_options_str})
+                AS r'''
+                __UDF_PLACE_HOLDER__
+                def bigframes_handler(*args):
+                    return {udf_name}(*args)
+                '''
+            """
+            )
+            .strip()
+            .replace("__UDF_PLACE_HOLDER__", udf_code)
+        )
 
         self._ensure_dataset_exists()
         self._create_bq_function(create_function_ddl)

bigframes/functions/_function_session.py

@@ -58,9 +58,6 @@
 
 from . import _function_client, _utils
 
-# BQ managed functions (@udf) currently only support Python 3.11.
-_MANAGED_FUNC_PYTHON_VERSIONS = ("python-3.11",)
-
 
 class FunctionSession:
     """Session to manage bigframes functions."""
@@ -758,7 +755,13 @@ def udf(
         name: Optional[str] = None,
         packages: Optional[Sequence[str]] = None,
     ):
-        """Decorator to turn a Python udf into a BigQuery managed function.
+        """Decorator to turn a Python user defined function (udf) into a
+        BigQuery managed function.
+
+        .. note::
+            The udf must be self-contained, i.e. it must not contain any
+            references to an import or variable defined outside the function
+            body.
 
         .. note::
             Please have following IAM roles enabled for you:
@@ -809,17 +812,8 @@ def udf(
                 of the form supported in
                 https://pip.pypa.io/en/stable/reference/requirements-file-format/.
         """
-        if not bigframes.options.experiments.udf:
-            raise bf_formatting.create_exception_with_feedback_link(NotImplementedError)
 
-        # Check the Python version.
-        python_version = _utils.get_python_version()
-        if python_version not in _MANAGED_FUNC_PYTHON_VERSIONS:
-            raise bf_formatting.create_exception_with_feedback_link(
-                RuntimeError,
-                f"Python version {python_version} is not supported yet for "
-                "BigFrames managed function.",
-            )
+        warnings.warn("udf is in preview.", category=bfe.PreviewWarning)
 
         # Some defaults may be used from the session if not provided otherwise.
         session = self._resolve_session(session)
@@ -862,7 +856,7 @@ def wrapper(func):
                             ValueError,
                             "'input_types' was not set and parameter "
                             f"'{parameter.name}' is missing a type annotation. "
-                            "Types are required to use managed function.",
+                            "Types are required to use udf.",
                         )
                     input_types.append(param_type)
             elif not isinstance(input_types, collections.abc.Sequence):
@@ -875,8 +869,7 @@ def wrapper(func):
                     raise bf_formatting.create_exception_with_feedback_link(
                         ValueError,
                         "'output_type' was not set and function is missing a "
-                        "return type annotation. Types are required to use "
-                        "managed function.",
+                        "return type annotation. Types are required to use udf",
                     )
 
             # The function will actually be receiving a pandas Series, but allow

bigframes/functions/_utils.py

@@ -64,9 +64,12 @@ def get_remote_function_locations(bq_location):
 
 
 def _get_updated_package_requirements(
-    package_requirements=None, is_row_processor=False
+    package_requirements=None, is_row_processor=False, capture_references=True
 ):
-    requirements = [f"cloudpickle=={cloudpickle.__version__}"]
+    requirements = []
+    if capture_references:
+        requirements.append(f"cloudpickle=={cloudpickle.__version__}")
+
     if is_row_processor:
         # bigframes function will send an entire row of data as json, which
         # would be converted to a pandas series and processed Ensure numpy

bigframes/geopandas/geoseries.py

@@ -96,3 +96,6 @@ def to_wkt(self: GeoSeries) -> bigframes.series.Series:
 
     def difference(self: GeoSeries, other: GeoSeries) -> bigframes.series.Series:  # type: ignore
         return self._apply_binary_op(other, ops.geo_st_difference_op)
+
+    def intersection(self: GeoSeries, other: GeoSeries) -> bigframes.series.Series:  # type: ignore
+        return self._apply_binary_op(other, ops.geo_st_intersection_op)