Merge branch 'main' into shuowei-anywidget-cell-visual

Author: shuoweil

README.rst

@@ -1,8 +1,7 @@
-:orphan:
-
 BigQuery DataFrames (BigFrames)
 ===============================
 
+
 |GA| |pypi| |versions|
 
 BigQuery DataFrames (also known as BigFrames) provides a Pythonic DataFrame

bigframes/_tools/docs.py

@@ -0,0 +1,39 @@
+# Copyright 2026 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.
+
+
+def inherit_docs(source_class):
+    """
+    A class decorator that copies docstrings from source_class to the
+    decorated class for any methods or attributes that match names.
+    """
+
+    def decorator(target_class):
+        if not target_class.__doc__ and source_class.__doc__:
+            target_class.__doc__ = source_class.__doc__
+
+        for name, source_item in vars(source_class).items():
+            if name in vars(target_class):
+                target_item = getattr(target_class, name)
+
+                if hasattr(target_item, "__doc__") and not target_item.__doc__:
+                    if hasattr(source_item, "__doc__") and source_item.__doc__:
+                        try:
+                            target_item.__doc__ = source_item.__doc__
+                        except AttributeError:
+                            pass
+
+        return target_class
+
+    return decorator

bigframes/bigquery/_operations/ai.py

@@ -606,7 +606,7 @@ def generate_table(
     model: Union[bigframes.ml.base.BaseEstimator, str, pd.Series],
     data: Union[dataframe.DataFrame, series.Series, pd.DataFrame, pd.Series],
     *,
-    output_schema: str,
+    output_schema: Union[str, Mapping[str, str]],
     temperature: Optional[float] = None,
     top_p: Optional[float] = None,
     max_output_tokens: Optional[int] = None,
@@ -642,8 +642,10 @@ def generate_table(
             treated as the 'prompt' column.  If a DataFrame is provided, it
             must contain a 'prompt' column, or you must rename the column you
             wish to generate table to 'prompt'.
-        output_schema (str):
-            A string defining the output schema (e.g., "col1 STRING, col2 INT64").
+        output_schema (str | Mapping[str, str]):
+            A string defining the output schema (e.g., "col1 STRING, col2 INT64"),
+            or 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`.
         temperature (float, optional):
             A FLOAT64 value that is used for sampling promiscuity. The value
             must be in the range ``[0.0, 1.0]``.
@@ -666,8 +668,17 @@ def generate_table(
     model_name, session = bq_utils.get_model_name_and_session(model, data)
     table_sql = bq_utils.to_sql(data)
 
+    if isinstance(output_schema, Mapping):
+        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)
+    else:
+        output_schema_str = output_schema
+
     struct_fields_bq: Dict[str, bigframes.core.sql.literals.STRUCT_VALUES] = {
-        "output_schema": output_schema
+        "output_schema": output_schema_str
     }
     if temperature is not None:
         struct_fields_bq["temperature"] = temperature

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -962,7 +962,7 @@ def isin_op_impl(x: ibis_types.Value, op: ops.IsInOp):
                 # to actually cast it, as that could be lossy (eg float -> int)
                 item_inferred_type = ibis_types.literal(item).type()
                 if (
-                    x.type() == item_inferred_type
+                    x.type().name == item_inferred_type.name
                     or x.type().is_numeric()
                     and item_inferred_type.is_numeric()
                 ):
@@ -978,7 +978,7 @@ def isin_op_impl(x: ibis_types.Value, op: ops.IsInOp):
 
 @scalar_op_compiler.register_unary_op(ops.ToDatetimeOp, pass_op=True)
 def to_datetime_op_impl(x: ibis_types.Value, op: ops.ToDatetimeOp):
-    if x.type() == ibis_dtypes.str:
+    if x.type() in (ibis_dtypes.str, ibis_dtypes.Timestamp("UTC")):  # type: ignore
         return x.try_cast(ibis_dtypes.Timestamp(None))  # type: ignore
     else:
         # Numerical inputs.
@@ -1001,6 +1001,9 @@ def to_timestamp_op_impl(x: ibis_types.Value, op: ops.ToTimestampOp):
             if op.format
             else timestamp(x)
         )
+    elif x.type() == ibis_dtypes.Timestamp(None):  # type: ignore
+
+        return timestamp(x)
     else:
         # Numerical inputs.
         if op.format:
@@ -2016,8 +2019,8 @@ def _ibis_num(number: float):
 
 
 @ibis_udf.scalar.builtin
-def timestamp(a: str) -> ibis_dtypes.timestamp:  # type: ignore
-    """Convert string to timestamp."""
+def timestamp(a) -> ibis_dtypes.timestamp:  # type: ignore
+    """Convert string or a datetime to timestamp."""
 
 
 @ibis_udf.scalar.builtin

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

@@ -33,16 +33,11 @@
 @register_unary_op(ops.IsInOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.IsInOp) -> sge.Expression:
     values = []
-    is_numeric_expr = dtypes.is_numeric(expr.dtype, include_bool=False)
     for value in op.values:
         if _is_null(value):
             continue
         dtype = dtypes.bigframes_type(type(value))
-        if (
-            expr.dtype == dtype
-            or is_numeric_expr
-            and dtypes.is_numeric(dtype, include_bool=False)
-        ):
+        if dtypes.can_compare(expr.dtype, dtype):
             values.append(sge.convert(value))
 
     if op.match_nulls:

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

@@ -371,7 +371,7 @@ def _(expr: TypedExpr, op: ops.ToDatetimeOp) -> sge.Expression:
         )
         return sge.Cast(this=result, to="DATETIME")
 
-    if expr.dtype == dtypes.STRING_DTYPE:
+    if expr.dtype in (dtypes.STRING_DTYPE, dtypes.TIMESTAMP_DTYPE):
         return sge.TryCast(this=expr.expr, to="DATETIME")
 
     value = expr.expr
@@ -396,7 +396,7 @@ def _(expr: TypedExpr, op: ops.ToTimestampOp) -> sge.Expression:
             "PARSE_TIMESTAMP", sge.convert(op.format), expr.expr, sge.convert("UTC")
         )
 
-    if expr.dtype == dtypes.STRING_DTYPE:
+    if expr.dtype in (dtypes.STRING_DTYPE, dtypes.DATETIME_DTYPE):
         return sge.func("TIMESTAMP", expr.expr)
 
     value = expr.expr

bigframes/core/groupby/dataframe_group_by.py

@@ -24,6 +24,7 @@
 import pandas as pd
 
 from bigframes import session
+from bigframes._tools import docs
 from bigframes.core import agg_expressions
 from bigframes.core import expression as ex
 import bigframes.core.block_transforms as block_ops
@@ -44,9 +45,8 @@
 
 
 @log_adapter.class_logger
-class DataFrameGroupBy(vendored_pandas_groupby.DataFrameGroupBy):
-    __doc__ = vendored_pandas_groupby.GroupBy.__doc__
-
+@docs.inherit_docs(vendored_pandas_groupby.DataFrameGroupBy)
+class DataFrameGroupBy:
     def __init__(
         self,
         block: blocks.Block,

bigframes/core/groupby/series_group_by.py

@@ -24,6 +24,7 @@
 import pandas
 
 from bigframes import session
+from bigframes._tools import docs
 from bigframes.core import expression as ex
 import bigframes.core.block_transforms as block_ops
 import bigframes.core.blocks as blocks
@@ -43,9 +44,8 @@
 
 
 @log_adapter.class_logger
+@docs.inherit_docs(vendored_pandas_groupby.SeriesGroupBy)
 class SeriesGroupBy(vendored_pandas_groupby.SeriesGroupBy):
-    __doc__ = vendored_pandas_groupby.GroupBy.__doc__
-
     def __init__(
         self,
         block: blocks.Block,

bigframes/core/indexes/base.py

@@ -27,6 +27,7 @@
 import pandas
 
 from bigframes import dtypes
+from bigframes._tools import docs
 import bigframes.core.agg_expressions as ex_types
 import bigframes.core.block_transforms as block_ops
 import bigframes.core.blocks as blocks
@@ -47,8 +48,8 @@
     import bigframes.series
 
 
-class Index(vendored_pandas_index.Index):
-    __doc__ = vendored_pandas_index.Index.__doc__
+@docs.inherit_docs(vendored_pandas_index.Index)
+class Index:
     _query_job = None
     _block: blocks.Block
     _linked_frame: Union[
@@ -777,6 +778,11 @@ def to_list(self, *, allow_large_results: Optional[bool] = None) -> list:
     def __len__(self):
         return self.shape[0]
 
+    def __bool__(self):
+        raise ValueError(
+            "Cannot convert Index into bool. Consider using .empty(), .item(), .any(), or .all() methods."
+        )
+
     def item(self):
         # Docstring is in third_party/bigframes_vendored/pandas/core/indexes/base.py
         return self.to_series().peek(2).item()

bigframes/core/indexes/datetimes.py

@@ -20,13 +20,14 @@
     datetimes as vendored_pandas_datetime_index,
 )
 
+from bigframes._tools import docs
 from bigframes.core import expression as ex
 from bigframes.core.indexes.base import Index
 from bigframes.operations import date_ops
 
 
-class DatetimeIndex(Index, vendored_pandas_datetime_index.DatetimeIndex):
-    __doc__ = vendored_pandas_datetime_index.DatetimeIndex.__doc__
+@docs.inherit_docs(vendored_pandas_datetime_index.DatetimeIndex)
+class DatetimeIndex(Index):
 
     # Must be above 5000 for pandas to delegate to bigframes for binops
     __pandas_priority__ = 12000