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

Author: TrevorBergeron

CHANGELOG.md

@@ -4,6 +4,23 @@
 
 [1]: https://pypi.org/project/bigframes/#history
 
+## [2.23.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v2.22.0...v2.23.0) (2025-09-29)
+
+
+### Features
+
+* Add ai.generate_double to bigframes.bigquery package ([#2111](https://github.com/googleapis/python-bigquery-dataframes/issues/2111)) ([6b8154c](https://github.com/googleapis/python-bigquery-dataframes/commit/6b8154c578bb1a276e9cf8fe494d91f8cd6260f2))
+
+
+### Bug Fixes
+
+* Prevent invalid syntax for no-op .replace ops ([#2112](https://github.com/googleapis/python-bigquery-dataframes/issues/2112)) ([c311876](https://github.com/googleapis/python-bigquery-dataframes/commit/c311876b2adbc0b66ae5e463c6e56466c6a6a495))
+
+
+### Documentation
+
+* Add timedelta notebook sample ([#2124](https://github.com/googleapis/python-bigquery-dataframes/issues/2124)) ([d1a9888](https://github.com/googleapis/python-bigquery-dataframes/commit/d1a9888a2b47de6aca5dddc94d0c8f280344b58a))
+
 ## [2.22.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v2.21.0...v2.22.0) (2025-09-25)
 
 

bigframes/bigquery/_operations/ai.py

@@ -35,6 +35,80 @@
 ]
 
 
+@log_adapter.method_logger(custom_base_name="bigquery_ai")
+def generate(
+    prompt: PROMPT_TYPE,
+    *,
+    connection_id: str | None = None,
+    endpoint: str | None = None,
+    request_type: Literal["dedicated", "shared", "unspecified"] = "unspecified",
+    model_params: Mapping[Any, Any] | None = None,
+    # TODO(b/446974666) Add output_schema parameter
+) -> series.Series:
+    """
+    Returns the AI analysis based on the prompt, which can be any combination of text and unstructured data.
+
+    **Examples:**
+
+        >>> import bigframes.pandas as bpd
+        >>> import bigframes.bigquery as bbq
+        >>> bpd.options.display.progress_bar = None
+        >>> country = bpd.Series(["Japan", "Canada"])
+        >>> bbq.ai.generate(("What's the capital city of ", country, " one word only"))
+        0    {'result': 'Tokyo\\n', 'full_response': '{"cand...
+        1    {'result': 'Ottawa\\n', 'full_response': '{"can...
+        dtype: struct<result: string, full_response: extension<dbjson<JSONArrowType>>, status: string>[pyarrow]
+
+        >>> bbq.ai.generate(("What's the capital city of ", country, " one word only")).struct.field("result")
+        0     Tokyo\\n
+        1    Ottawa\\n
+        Name: result, dtype: string
+
+    Args:
+        prompt (Series | List[str|Series] | Tuple[str|Series, ...]):
+            A mixture of Series and string literals that specifies the prompt to send to the model. The Series can be BigFrames Series
+            or pandas Series.
+        connection_id (str, optional):
+            Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
+            If not provided, the connection from the current session will be used.
+        endpoint (str, optional):
+            Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
+            generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
+            uses the full endpoint of the model. If you don't specify an ENDPOINT value, BigQuery ML selects a recent stable
+            version of Gemini to use.
+        request_type (Literal["dedicated", "shared", "unspecified"]):
+            Specifies the type of inference request to send to the Gemini model. The request type determines what quota the request uses.
+            * "dedicated": function only uses Provisioned Throughput quota. The function returns the error Provisioned throughput is not
+            purchased or is not active if Provisioned Throughput quota isn't available.
+            * "shared": the function only uses dynamic shared quota (DSQ), even if you have purchased Provisioned Throughput quota.
+            * "unspecified": If you haven't purchased Provisioned Throughput quota, the function uses DSQ quota.
+            If you have purchased Provisioned Throughput quota, the function uses the Provisioned Throughput quota first.
+            If requests exceed the Provisioned Throughput quota, the overflow traffic uses DSQ quota.
+        model_params (Mapping[Any, Any]):
+            Provides additional parameters to the model. The MODEL_PARAMS value must conform to the generateContent request body format.
+
+    Returns:
+        bigframes.series.Series: A new struct Series with the result data. The struct contains these fields:
+        * "result": a STRING value containing the model's response to the prompt. The result is None if the request fails or is filtered by responsible AI.
+        * "full_response": a JSON value containing the response from the projects.locations.endpoints.generateContent call to the model.
+        The generated text is in the text element.
+        * "status": a STRING value that contains the API response status for the corresponding row. This value is empty if the operation was successful.
+    """
+
+    prompt_context, series_list = _separate_context_and_series(prompt)
+    assert len(series_list) > 0
+
+    operator = ai_ops.AIGenerate(
+        prompt_context=tuple(prompt_context),
+        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        endpoint=endpoint,
+        request_type=request_type,
+        model_params=json.dumps(model_params) if model_params else None,
+    )
+
+    return series_list[0]._apply_nary_op(operator, series_list[1:])
+
+
 @log_adapter.method_logger(custom_base_name="bigquery_ai")
 def generate_bool(
     prompt: PROMPT_TYPE,

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1173,6 +1173,10 @@ def udf(*inputs):
 
 @scalar_op_compiler.register_unary_op(ops.MapOp, pass_op=True)
 def map_op_impl(x: ibis_types.Value, op: ops.MapOp):
+    # this should probably be handled by a rewriter
+    if len(op.mappings) == 0:
+        return x
+
     case = ibis_api.case()
     for mapping in op.mappings:
         case = case.when(x == mapping[0], mapping[1])
@@ -1970,6 +1974,20 @@ def struct_op_impl(
     return ibis_types.struct(data)
 
 
+@scalar_op_compiler.register_nary_op(ops.AIGenerate, pass_op=True)
+def ai_generate(
+    *values: ibis_types.Value, op: ops.AIGenerate
+) -> ibis_types.StructValue:
+
+    return ai_ops.AIGenerate(
+        _construct_prompt(values, op.prompt_context),  # type: ignore
+        op.connection_id,  # type: ignore
+        op.endpoint,  # type: ignore
+        op.request_type.upper(),  # type: ignore
+        op.model_params,  # type: ignore
+    ).to_expr()
+
+
 @scalar_op_compiler.register_nary_op(ops.AIGenerateBool, pass_op=True)
 def ai_generate_bool(
     *values: ibis_types.Value, op: ops.AIGenerateBool

bigframes/core/compile/sqlglot/__init__.py

@@ -17,6 +17,7 @@
 import bigframes.core.compile.sqlglot.expressions.ai_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.array_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.blob_ops  # noqa: F401
+import bigframes.core.compile.sqlglot.expressions.bool_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.comparison_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.date_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.datetime_ops  # noqa: F401

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -38,6 +38,17 @@ def compile(
     return UNARY_OP_REGISTRATION[op](op, column, window=window)
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.AllOp)
+def _(
+    op: agg_ops.AllOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    # BQ will return null for empty column, result would be false in pandas.
+    result = apply_window_if_present(sge.func("LOGICAL_AND", column.expr), window)
+    return sge.func("IFNULL", result, sge.true())
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.ApproxQuartilesOp)
 def _(
     op: agg_ops.ApproxQuartilesOp,
@@ -69,6 +80,15 @@ def _(
     return sge.func("APPROX_TOP_COUNT", column.expr, sge.convert(op.number))
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.AnyValueOp)
+def _(
+    op: agg_ops.AnyValueOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    return apply_window_if_present(sge.func("ANY_VALUE", column.expr), window)
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.CountOp)
 def _(
     op: agg_ops.CountOp,

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

@@ -26,6 +26,13 @@
 register_nary_op = scalar_compiler.scalar_op_compiler.register_nary_op
 
 
+@register_nary_op(ops.AIGenerate, pass_op=True)
+def _(*exprs: TypedExpr, op: ops.AIGenerate) -> sge.Expression:
+    args = [_construct_prompt(exprs, op.prompt_context)] + _construct_named_args(op)
+
+    return sge.func("AI.GENERATE", *args)
+
+
 @register_nary_op(ops.AIGenerateBool, pass_op=True)
 def _(*exprs: TypedExpr, op: ops.AIGenerateBool) -> sge.Expression:
     args = [_construct_prompt(exprs, op.prompt_context)] + _construct_named_args(op)

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

@@ -0,0 +1,47 @@
+# 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
+
+import sqlglot.expressions as sge
+
+from bigframes import dtypes
+from bigframes import operations as ops
+from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
+import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+
+register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
+
+
+@register_binary_op(ops.and_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE:
+        return sge.And(this=left.expr, expression=right.expr)
+    return sge.BitwiseAnd(this=left.expr, expression=right.expr)
+
+
+@register_binary_op(ops.or_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE:
+        return sge.Or(this=left.expr, expression=right.expr)
+    return sge.BitwiseOr(this=left.expr, expression=right.expr)
+
+
+@register_binary_op(ops.xor_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.dtype == dtypes.BOOL_DTYPE and right.dtype == dtypes.BOOL_DTYPE:
+        left_expr = sge.And(this=left.expr, expression=sge.Not(this=right.expr))
+        right_expr = sge.And(this=sge.Not(this=left.expr), expression=right.expr)
+        return sge.Or(this=left_expr, expression=right_expr)
+    return sge.BitwiseXor(this=left.expr, expression=right.expr)

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

@@ -78,6 +78,8 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_unary_op(ops.MapOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.MapOp) -> sge.Expression:
+    if len(op.mappings) == 0:
+        return expr.expr
     return sge.Case(
         this=expr.expr,
         ifs=[

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

@@ -323,6 +323,49 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return result
 
 
+@register_binary_op(ops.mod_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    # In BigQuery returned value has the same sign as X. In pandas, the sign of y is used, so we need to flip the result if sign(x) != sign(y)
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+
+    # BigQuery MOD function doesn't support float types, so cast to BIGNUMERIC
+    if left.dtype == dtypes.FLOAT_DTYPE or right.dtype == dtypes.FLOAT_DTYPE:
+        left_expr = sge.Cast(this=left_expr, to="BIGNUMERIC")
+        right_expr = sge.Cast(this=right_expr, to="BIGNUMERIC")
+
+    # MOD(N, 0) will error in bigquery, but needs to return null
+    bq_mod = sge.Mod(this=left_expr, expression=right_expr)
+    zero_result = (
+        constants._NAN
+        if (left.dtype == dtypes.FLOAT_DTYPE or right.dtype == dtypes.FLOAT_DTYPE)
+        else constants._ZERO
+    )
+    return sge.Case(
+        ifs=[
+            sge.If(
+                this=sge.EQ(this=right_expr, expression=constants._ZERO),
+                true=zero_result * left_expr,
+            ),
+            sge.If(
+                this=sge.and_(
+                    right_expr < constants._ZERO,
+                    bq_mod > constants._ZERO,
+                ),
+                true=right_expr + bq_mod,
+            ),
+            sge.If(
+                this=sge.and_(
+                    right_expr > constants._ZERO,
+                    bq_mod < constants._ZERO,
+                ),
+                true=right_expr + bq_mod,
+            ),
+        ],
+        default=bq_mod,
+    )
+
+
 @register_binary_op(ops.mul_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     left_expr = _coerce_bool_to_int(left)

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -175,7 +175,7 @@ def from_union(
         ), f"At least two select expressions must be provided, but got {selects}."
 
         existing_ctes: list[sge.CTE] = []
-        union_selects: list[sge.Select] = []
+        union_selects: list[sge.Expression] = []
         for select in selects:
             assert isinstance(
                 select, sge.Select
@@ -204,10 +204,14 @@ def from_union(
                 sge.Select().select(*selections).from_(sge.Table(this=new_cte_name))
             )
 
-        union_expr = sg.union(
-            *union_selects,
-            distinct=False,
-            copy=False,
+        union_expr = typing.cast(
+            sge.Select,
+            functools.reduce(
+                lambda x, y: sge.Union(
+                    this=x, expression=y, distinct=False, copy=False
+                ),
+                union_selects,
+            ),
         )
         final_select_expr = sge.Select().select(sge.Star()).from_(union_expr.subquery())
         final_select_expr.set("with", sge.With(expressions=existing_ctes))