Merge branch 'main' into shuowei-blob-image

Author: shuoweil

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/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

bigframes/core/indexes/multi.py

@@ -19,6 +19,7 @@
 import bigframes_vendored.pandas.core.indexes.multi as vendored_pandas_multindex
 import pandas
 
+from bigframes._tools import docs
 from bigframes.core import blocks
 from bigframes.core import expression as ex
 from bigframes.core.indexes.base import Index
@@ -27,9 +28,8 @@
     import bigframes.session
 
 
-class MultiIndex(Index, vendored_pandas_multindex.MultiIndex):
-    __doc__ = vendored_pandas_multindex.MultiIndex.__doc__
-
+@docs.inherit_docs(vendored_pandas_multindex.MultiIndex)
+class MultiIndex(Index):
     @classmethod
     def from_tuples(
         cls,

bigframes/core/window/rolling.py

@@ -22,6 +22,7 @@
 import pandas
 
 from bigframes import dtypes
+from bigframes._tools import docs
 from bigframes.core import agg_expressions
 from bigframes.core import expression as ex
 from bigframes.core import ordering, utils, window_spec
@@ -36,9 +37,8 @@
 
 
 @log_adapter.class_logger
-class Window(vendored_pandas_rolling.Window):
-    __doc__ = vendored_pandas_rolling.Window.__doc__
-
+@docs.inherit_docs(vendored_pandas_rolling.Window)
+class Window:
     def __init__(
         self,
         block: blocks.Block,