Merge branch 'main' into tswast-doctest-boilerplate

Author: tswast

CHANGELOG.md

@@ -4,6 +4,29 @@
 
 [1]: https://pypi.org/project/bigframes/#history
 
+## [2.24.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v2.23.0...v2.24.0) (2025-10-07)
+
+
+### Features
+
+* Add ai.classify() to bigframes.bigquery package ([#2137](https://github.com/googleapis/python-bigquery-dataframes/issues/2137)) ([56e5033](https://github.com/googleapis/python-bigquery-dataframes/commit/56e50331d198b7f517f85695c208f893ab9389d2))
+* Add ai.generate() to bigframes.bigquery module ([#2128](https://github.com/googleapis/python-bigquery-dataframes/issues/2128)) ([3810452](https://github.com/googleapis/python-bigquery-dataframes/commit/3810452f16d8d6c9d3eb9075f1537177d98b4725))
+* Add ai.if_() and ai.score() to bigframes.bigquery package ([#2132](https://github.com/googleapis/python-bigquery-dataframes/issues/2132)) ([32502f4](https://github.com/googleapis/python-bigquery-dataframes/commit/32502f4195306d262788f39d1ab4206fc84ae50e))
+
+
+### Bug Fixes
+
+* Fix internal type errors with temporal accessors ([#2125](https://github.com/googleapis/python-bigquery-dataframes/issues/2125)) ([c390da1](https://github.com/googleapis/python-bigquery-dataframes/commit/c390da11b7c2aa710bc2fbc692efb9f06059e4c4))
+* Fix row count local execution bug ([#2133](https://github.com/googleapis/python-bigquery-dataframes/issues/2133)) ([ece0762](https://github.com/googleapis/python-bigquery-dataframes/commit/ece07623e354a1dde2bd37020349e13f682e863f))
+* Join on, how args are now positional ([#2140](https://github.com/googleapis/python-bigquery-dataframes/issues/2140)) ([b711815](https://github.com/googleapis/python-bigquery-dataframes/commit/b7118152bfecc6ecf67aa4df23ec3f0a2b08aa30))
+* Only show JSON dtype warning when accessing dtypes directly ([#2136](https://github.com/googleapis/python-bigquery-dataframes/issues/2136)) ([eca22ee](https://github.com/googleapis/python-bigquery-dataframes/commit/eca22ee3104104cea96189391e527cad09bd7509))
+* Remove noisy AmbiguousWindowWarning from partial ordering mode ([#2129](https://github.com/googleapis/python-bigquery-dataframes/issues/2129)) ([4607f86](https://github.com/googleapis/python-bigquery-dataframes/commit/4607f86ebd77b916aafc37f69725b676e203b332))
+
+
+### Performance Improvements
+
+* Scale read stream workers to cpu count ([#2135](https://github.com/googleapis/python-bigquery-dataframes/issues/2135)) ([67e46cd](https://github.com/googleapis/python-bigquery-dataframes/commit/67e46cd47933b84b55808003ed344b559e47c498))
+
 ## [2.23.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v2.22.0...v2.23.0) (2025-09-29)
 
 

bigframes/bigquery/_operations/ai.py

@@ -25,7 +25,7 @@
 
 from bigframes import clients, dtypes, series, session
 from bigframes.core import convert, log_adapter
-from bigframes.operations import ai_ops
+from bigframes.operations import ai_ops, output_schemas
 
 PROMPT_TYPE = Union[
     series.Series,
@@ -43,7 +43,7 @@ def generate(
     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
+    output_schema: Mapping[str, str] | None = None,
 ) -> series.Series:
     """
     Returns the AI analysis based on the prompt, which can be any combination of text and unstructured data.
@@ -63,6 +63,14 @@ def generate(
         1    Ottawa\\n
         Name: result, dtype: string
 
+        You get structured output when the `output_schema` parameter is set:
+
+        >>> animals = bpd.Series(["Rabbit", "Spider"])
+        >>> bbq.ai.generate(animals, output_schema={"number_of_legs": "INT64", "is_herbivore": "BOOL"})
+        0    {'is_herbivore': True, 'number_of_legs': 4, 'f...
+        1    {'is_herbivore': False, 'number_of_legs': 8, '...
+        dtype: struct<is_herbivore: bool, number_of_legs: int64, full_response: extension<dbjson<JSONArrowType>>, status: string>[pyarrow]
+
     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
@@ -85,10 +93,14 @@ def generate(
             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.
+        output_schema (Mapping[str, str]):
+            A mapping value that specifies the schema of the output, in the form {field_name: data_type}. Supported data types include
+            `STRING`, `INT64`, `FLOAT64`, `BOOL`, `ARRAY`, and `STRUCT`.
 
     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.
+        If you specify an output schema then result is replaced by your custom schema.
         * "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.
@@ -97,12 +109,22 @@ def generate(
     prompt_context, series_list = _separate_context_and_series(prompt)
     assert len(series_list) > 0
 
+    if output_schema is None:
+        output_schema_str = None
+    else:
+        output_schema_str = ", ".join(
+            [f"{name} {sql_type}" for name, sql_type in output_schema.items()]
+        )
+        # Validate user input
+        output_schemas.parse_sql_fields(output_schema_str)
+
     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,
+        output_schema=output_schema_str,
     )
 
     return series_list[0]._apply_nary_op(operator, series_list[1:])

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1985,6 +1985,7 @@ def ai_generate(
         op.endpoint,  # type: ignore
         op.request_type.upper(),  # type: ignore
         op.model_params,  # type: ignore
+        op.output_schema,  # type: ignore
     ).to_expr()
 
 

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

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 from dataclasses import asdict
-import typing
 
 import sqlglot.expressions as sge
 
@@ -105,24 +104,24 @@ def _construct_named_args(op: ops.NaryOp) -> list[sge.Kwarg]:
 
     op_args = asdict(op)
 
-    connection_id = typing.cast(str, op_args["connection_id"])
+    connection_id = op_args["connection_id"]
     args.append(
         sge.Kwarg(this="connection_id", expression=sge.Literal.string(connection_id))
     )
 
-    endpoit = typing.cast(str, op_args.get("endpoint", None))
+    endpoit = op_args.get("endpoint", None)
     if endpoit is not None:
         args.append(sge.Kwarg(this="endpoint", expression=sge.Literal.string(endpoit)))
 
-    request_type = typing.cast(str, op_args.get("request_type", None))
+    request_type = op_args.get("request_type", None)
     if request_type is not None:
         args.append(
             sge.Kwarg(
                 this="request_type", expression=sge.Literal.string(request_type.upper())
             )
         )
 
-    model_params = typing.cast(str, op_args.get("model_params", None))
+    model_params = op_args.get("model_params", None)
     if model_params is not None:
         args.append(
             sge.Kwarg(
@@ -133,4 +132,13 @@ def _construct_named_args(op: ops.NaryOp) -> list[sge.Kwarg]:
             )
         )
 
+    output_schema = op_args.get("output_schema", None)
+    if output_schema is not None:
+        args.append(
+            sge.Kwarg(
+                this="output_schema",
+                expression=sge.Literal.string(output_schema),
+            )
+        )
+
     return args

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

@@ -18,9 +18,9 @@
 
 from bigframes import dtypes
 from bigframes import operations as ops
+from bigframes.core.compile.sqlglot import sqlglot_types
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
-from bigframes.core.compile.sqlglot.sqlglot_types import SQLGlotType
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
 
@@ -29,7 +29,7 @@
 def _(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
     from_type = expr.dtype
     to_type = op.to_type
-    sg_to_type = SQLGlotType.from_bigframes_dtype(to_type)
+    sg_to_type = sqlglot_types.from_bigframes_dtype(to_type)
     sg_expr = expr.expr
 
     if to_type == dtypes.JSON_DTYPE:

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -79,7 +79,7 @@ def from_pyarrow(
             expressions=[
                 sge.ColumnDef(
                     this=sge.to_identifier(field.column, quoted=True),
-                    kind=sgt.SQLGlotType.from_bigframes_dtype(field.dtype),
+                    kind=sgt.from_bigframes_dtype(field.dtype),
                 )
                 for field in schema.items
             ],
@@ -620,7 +620,7 @@ def _select_to_cte(expr: sge.Select, cte_name: sge.Identifier) -> sge.Select:
 
 
 def _literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
-    sqlglot_type = sgt.SQLGlotType.from_bigframes_dtype(dtype)
+    sqlglot_type = sgt.from_bigframes_dtype(dtype)
     if value is None:
         return _cast(sge.Null(), sqlglot_type)
     elif dtype == dtypes.BYTES_DTYPE:

bigframes/core/compile/sqlglot/sqlglot_types.py

@@ -25,62 +25,57 @@
 import bigframes.dtypes
 
 
-class SQLGlotType:
-    @classmethod
-    def from_bigframes_dtype(
-        cls,
-        bigframes_dtype: typing.Union[
-            bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype, np.dtype[typing.Any]
-        ],
-    ) -> str:
-        if bigframes_dtype == bigframes.dtypes.INT_DTYPE:
-            return "INT64"
-        elif bigframes_dtype == bigframes.dtypes.FLOAT_DTYPE:
-            return "FLOAT64"
-        elif bigframes_dtype == bigframes.dtypes.STRING_DTYPE:
-            return "STRING"
-        elif bigframes_dtype == bigframes.dtypes.BOOL_DTYPE:
-            return "BOOLEAN"
-        elif bigframes_dtype == bigframes.dtypes.DATE_DTYPE:
-            return "DATE"
-        elif bigframes_dtype == bigframes.dtypes.TIME_DTYPE:
-            return "TIME"
-        elif bigframes_dtype == bigframes.dtypes.DATETIME_DTYPE:
-            return "DATETIME"
-        elif bigframes_dtype == bigframes.dtypes.TIMESTAMP_DTYPE:
-            return "TIMESTAMP"
-        elif bigframes_dtype == bigframes.dtypes.BYTES_DTYPE:
-            return "BYTES"
-        elif bigframes_dtype == bigframes.dtypes.NUMERIC_DTYPE:
-            return "NUMERIC"
-        elif bigframes_dtype == bigframes.dtypes.BIGNUMERIC_DTYPE:
-            return "BIGNUMERIC"
-        elif bigframes_dtype == bigframes.dtypes.JSON_DTYPE:
-            return "JSON"
-        elif bigframes_dtype == bigframes.dtypes.GEO_DTYPE:
-            return "GEOGRAPHY"
-        elif bigframes_dtype == bigframes.dtypes.TIMEDELTA_DTYPE:
-            return "INT64"
-        elif isinstance(bigframes_dtype, pd.ArrowDtype):
-            if pa.types.is_list(bigframes_dtype.pyarrow_dtype):
-                inner_bigframes_dtype = bigframes.dtypes.arrow_dtype_to_bigframes_dtype(
-                    bigframes_dtype.pyarrow_dtype.value_type
+def from_bigframes_dtype(
+    bigframes_dtype: typing.Union[
+        bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype, np.dtype[typing.Any]
+    ],
+) -> str:
+    if bigframes_dtype == bigframes.dtypes.INT_DTYPE:
+        return "INT64"
+    elif bigframes_dtype == bigframes.dtypes.FLOAT_DTYPE:
+        return "FLOAT64"
+    elif bigframes_dtype == bigframes.dtypes.STRING_DTYPE:
+        return "STRING"
+    elif bigframes_dtype == bigframes.dtypes.BOOL_DTYPE:
+        return "BOOLEAN"
+    elif bigframes_dtype == bigframes.dtypes.DATE_DTYPE:
+        return "DATE"
+    elif bigframes_dtype == bigframes.dtypes.TIME_DTYPE:
+        return "TIME"
+    elif bigframes_dtype == bigframes.dtypes.DATETIME_DTYPE:
+        return "DATETIME"
+    elif bigframes_dtype == bigframes.dtypes.TIMESTAMP_DTYPE:
+        return "TIMESTAMP"
+    elif bigframes_dtype == bigframes.dtypes.BYTES_DTYPE:
+        return "BYTES"
+    elif bigframes_dtype == bigframes.dtypes.NUMERIC_DTYPE:
+        return "NUMERIC"
+    elif bigframes_dtype == bigframes.dtypes.BIGNUMERIC_DTYPE:
+        return "BIGNUMERIC"
+    elif bigframes_dtype == bigframes.dtypes.JSON_DTYPE:
+        return "JSON"
+    elif bigframes_dtype == bigframes.dtypes.GEO_DTYPE:
+        return "GEOGRAPHY"
+    elif bigframes_dtype == bigframes.dtypes.TIMEDELTA_DTYPE:
+        return "INT64"
+    elif isinstance(bigframes_dtype, pd.ArrowDtype):
+        if pa.types.is_list(bigframes_dtype.pyarrow_dtype):
+            inner_bigframes_dtype = bigframes.dtypes.arrow_dtype_to_bigframes_dtype(
+                bigframes_dtype.pyarrow_dtype.value_type
+            )
+            return f"ARRAY<{from_bigframes_dtype(inner_bigframes_dtype)}>"
+        elif pa.types.is_struct(bigframes_dtype.pyarrow_dtype):
+            struct_type = typing.cast(pa.StructType, bigframes_dtype.pyarrow_dtype)
+            inner_fields: list[str] = []
+            for i in range(struct_type.num_fields):
+                field = struct_type.field(i)
+                key = sg.to_identifier(field.name).sql("bigquery")
+                dtype = from_bigframes_dtype(
+                    bigframes.dtypes.arrow_dtype_to_bigframes_dtype(field.type)
                 )
-                return (
-                    f"ARRAY<{SQLGlotType.from_bigframes_dtype(inner_bigframes_dtype)}>"
-                )
-            elif pa.types.is_struct(bigframes_dtype.pyarrow_dtype):
-                struct_type = typing.cast(pa.StructType, bigframes_dtype.pyarrow_dtype)
-                inner_fields: list[str] = []
-                for i in range(struct_type.num_fields):
-                    field = struct_type.field(i)
-                    key = sg.to_identifier(field.name).sql("bigquery")
-                    dtype = SQLGlotType.from_bigframes_dtype(
-                        bigframes.dtypes.arrow_dtype_to_bigframes_dtype(field.type)
-                    )
-                    inner_fields.append(f"{key} {dtype}")
-                return "STRUCT<{}>".format(", ".join(inner_fields))
+                inner_fields.append(f"{key} {dtype}")
+            return "STRUCT<{}>".format(", ".join(inner_fields))
 
-        raise ValueError(
-            f"Unsupported type for {bigframes_dtype}. {constants.FEEDBACK_LINK}"
-        )
+    raise ValueError(
+        f"Unsupported type for {bigframes_dtype}. {constants.FEEDBACK_LINK}"
+    )

bigframes/core/indexes/base.py

@@ -761,6 +761,58 @@ def item(self):
         # Docstring is in third_party/bigframes_vendored/pandas/core/indexes/base.py
         return self.to_series().peek(2).item()
 
+    def __eq__(self, other) -> Index:  # type: ignore
+        return self._apply_binop(other, ops.eq_op)
+
+    def _apply_binop(self, other, op: ops.BinaryOp) -> Index:
+        # TODO: Handle local objects, or objects not implicitly alignable? Gets ambiguous with partial ordering though
+        if isinstance(other, (bigframes.series.Series, Index)):
+            other = Index(other)
+            if other.nlevels != self.nlevels:
+                raise ValueError("Dimensions do not match")
+
+            lexpr = self._block.expr
+            rexpr = other._block.expr
+            join_result = lexpr.try_row_join(rexpr)
+            if join_result is None:
+                raise ValueError("Cannot align objects")
+
+            expr, (lmap, rmap) = join_result
+
+            expr, res_ids = expr.compute_values(
+                [
+                    op.as_expr(lmap[lid], rmap[rid])
+                    for lid, rid in zip(lexpr.column_ids, rexpr.column_ids)
+                ]
+            )
+            return Index(
+                blocks.Block(
+                    expr.select_columns(res_ids),
+                    index_columns=res_ids,
+                    column_labels=[],
+                    index_labels=[None] * len(res_ids),
+                )
+            )
+        elif (
+            isinstance(other, bigframes.dtypes.LOCAL_SCALAR_TYPES) and self.nlevels == 1
+        ):
+            block, id = self._block.project_expr(
+                op.as_expr(self._block.index_columns[0], ex.const(other))
+            )
+            return Index(block.select_column(id))
+        elif isinstance(other, tuple) and len(other) == self.nlevels:
+            block = self._block.project_exprs(
+                [
+                    op.as_expr(self._block.index_columns[i], ex.const(other[i]))
+                    for i in range(self.nlevels)
+                ],
+                labels=[None] * self.nlevels,
+                drop=True,
+            )
+            return Index(block.set_index(block.value_columns))
+        else:
+            return NotImplemented
+
 
 def _should_create_datetime_index(block: blocks.Block) -> bool:
     if len(block.index.dtypes) != 1: