Merge branch 'main' into sycai_ai_doc_fix

Author: sycai

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

@@ -50,4 +50,4 @@ def _(
     if window is None:
         # ROW_NUMBER always needs an OVER clause.
         return sge.Window(this=result)
-    return apply_window_if_present(result, window)
+    return apply_window_if_present(result, window, include_framing_clauses=False)

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

@@ -104,7 +104,51 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    return apply_window_if_present(sge.func("DENSE_RANK"), window)
+    return apply_window_if_present(
+        sge.func("DENSE_RANK"), window, include_framing_clauses=False
+    )
+
+
+@UNARY_OP_REGISTRATION.register(agg_ops.FirstOp)
+def _(
+    op: agg_ops.FirstOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    # FIRST_VALUE in BQ respects nulls by default.
+    return apply_window_if_present(sge.FirstValue(this=column.expr), window)
+
+
+@UNARY_OP_REGISTRATION.register(agg_ops.FirstNonNullOp)
+def _(
+    op: agg_ops.FirstNonNullOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    return apply_window_if_present(
+        sge.IgnoreNulls(this=sge.FirstValue(this=column.expr)), window
+    )
+
+
+@UNARY_OP_REGISTRATION.register(agg_ops.LastOp)
+def _(
+    op: agg_ops.LastOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    # LAST_VALUE in BQ respects nulls by default.
+    return apply_window_if_present(sge.LastValue(this=column.expr), window)
+
+
+@UNARY_OP_REGISTRATION.register(agg_ops.LastNonNullOp)
+def _(
+    op: agg_ops.LastNonNullOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    return apply_window_if_present(
+        sge.IgnoreNulls(this=sge.LastValue(this=column.expr)), window
+    )
 
 
 @UNARY_OP_REGISTRATION.register(agg_ops.MaxOp)
@@ -182,7 +226,9 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    return apply_window_if_present(sge.func("RANK"), window)
+    return apply_window_if_present(
+        sge.func("RANK"), window, include_framing_clauses=False
+    )
 
 
 @UNARY_OP_REGISTRATION.register(agg_ops.SizeUnaryOp)

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

@@ -25,6 +25,7 @@
 def apply_window_if_present(
     value: sge.Expression,
     window: typing.Optional[window_spec.WindowSpec] = None,
+    include_framing_clauses: bool = True,
 ) -> sge.Expression:
     if window is None:
         return value
@@ -64,11 +65,11 @@ def apply_window_if_present(
     if not window.bounds and not order:
         return sge.Window(this=value, partition_by=group_by)
 
-    if not window.bounds:
+    if not window.bounds and not include_framing_clauses:
         return sge.Window(this=value, partition_by=group_by, order=order)
 
     kind = (
-        "ROWS" if isinstance(window.bounds, window_spec.RowsWindowBounds) else "RANGE"
+        "RANGE" if isinstance(window.bounds, window_spec.RangeWindowBounds) else "ROWS"
     )
 
     start: typing.Union[int, float, None] = None
@@ -125,7 +126,7 @@ def get_window_order_by(
                         nulls_first=nulls_first,
                     )
                 )
-            elif not nulls_first and not desc:
+            elif (not nulls_first) and (not desc):
                 order_by.append(
                     sge.Ordered(
                         this=is_null_expr,

bigframes/core/indexes/base.py

@@ -383,9 +383,16 @@ def to_series(
 
         name = self.name if name is None else name
         if index is None:
-            return bigframes.series.Series(data=self, index=self, name=name)
+            return bigframes.series.Series(
+                data=self, index=self, name=name, session=self._session
+            )
         else:
-            return bigframes.series.Series(data=self, index=Index(index), name=name)
+            return bigframes.series.Series(
+                data=self,
+                index=Index(index, session=self._session),
+                name=name,
+                session=self._session,
+            )
 
     def get_level_values(self, level) -> Index:
         level_n = level if isinstance(level, int) else self.names.index(level)

bigframes/core/indexes/multi.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import cast, Hashable, Iterable, Sequence
+from typing import cast, Hashable, Iterable, Optional, Sequence, TYPE_CHECKING
 
 import bigframes_vendored.pandas.core.indexes.multi as vendored_pandas_multindex
 import pandas
@@ -23,6 +23,9 @@
 from bigframes.core import expression as ex
 from bigframes.core.indexes.base import Index
 
+if TYPE_CHECKING:
+    import bigframes.session
+
 
 class MultiIndex(Index, vendored_pandas_multindex.MultiIndex):
     __doc__ = vendored_pandas_multindex.MultiIndex.__doc__
@@ -33,21 +36,25 @@ def from_tuples(
         tuples: Iterable[tuple[Hashable, ...]],
         sortorder: int | None = None,
         names: Sequence[Hashable] | Hashable | None = None,
+        *,
+        session: Optional[bigframes.session.Session] = None,
     ) -> MultiIndex:
         pd_index = pandas.MultiIndex.from_tuples(tuples, sortorder, names)
         # Index.__new__ should detect multiple levels and properly create a multiindex
-        return cast(MultiIndex, Index(pd_index))
+        return cast(MultiIndex, Index(pd_index, session=session))
 
     @classmethod
     def from_arrays(
         cls,
         arrays,
         sortorder: int | None = None,
         names=None,
+        *,
+        session: Optional[bigframes.session.Session] = None,
     ) -> MultiIndex:
         pd_index = pandas.MultiIndex.from_arrays(arrays, sortorder, names)
         # Index.__new__ should detect multiple levels and properly create a multiindex
-        return cast(MultiIndex, Index(pd_index))
+        return cast(MultiIndex, Index(pd_index, session=session))
 
     def __eq__(self, other) -> Index:  # type: ignore
         import bigframes.operations as ops
@@ -71,3 +78,38 @@ def __eq__(self, other) -> Index:  # type: ignore
                 index_labels=[None],
             )
         )
+
+
+class MultiIndexAccessor:
+    """Proxy to MultiIndex constructors to allow a session to be passed in."""
+
+    def __init__(self, session: bigframes.session.Session):
+        self._session = session
+
+    def __call__(self, *args, **kwargs) -> MultiIndex:
+        """Construct a MultiIndex using the associated Session.
+
+        See :class:`bigframes.pandas.MultiIndex`.
+        """
+        return MultiIndex(*args, session=self._session, **kwargs)
+
+    def from_arrays(self, *args, **kwargs) -> MultiIndex:
+        """Construct a MultiIndex using the associated Session.
+
+        See :func:`bigframes.pandas.MultiIndex.from_arrays`.
+        """
+        return MultiIndex.from_arrays(*args, session=self._session, **kwargs)
+
+    def from_frame(self, *args, **kwargs) -> MultiIndex:
+        """Construct a MultiIndex using the associated Session.
+
+        See :func:`bigframes.pandas.MultiIndex.from_frame`.
+        """
+        return cast(MultiIndex, MultiIndex.from_frame(*args, **kwargs))
+
+    def from_tuples(self, *args, **kwargs) -> MultiIndex:
+        """Construct a MultiIndex using the associated Session.
+
+        See :func:`bigframes.pandas.MultiIndex.from_tuples`.
+        """
+        return MultiIndex.from_tuples(*args, session=self._session, **kwargs)

bigframes/core/log_adapter.py

@@ -155,7 +155,9 @@ def method_logger(method=None, /, *, custom_base_name: Optional[str] = None):
     def outer_wrapper(method):
         @functools.wraps(method)
         def wrapper(*args, **kwargs):
-            api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__)
+            api_method_name = getattr(
+                method, LOG_OVERRIDE_NAME, method.__name__
+            ).lower()
             if custom_base_name is None:
                 qualname_parts = getattr(method, "__qualname__", method.__name__).split(
                     "."

bigframes/core/reshape/tile.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import typing
+from typing import Optional, TYPE_CHECKING
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.reshape.tile as vendored_pandas_tile
@@ -31,6 +32,9 @@
 import bigframes.operations.aggregations as agg_ops
 import bigframes.series
 
+if TYPE_CHECKING:
+    import bigframes.session
+
 
 def cut(
     x,
@@ -42,6 +46,7 @@ def cut(
     *,
     right: typing.Optional[bool] = True,
     labels: typing.Union[typing.Iterable[str], bool, None] = None,
+    session: Optional[bigframes.session.Session] = None,
 ) -> bigframes.series.Series:
     if (
         labels is not None
@@ -65,7 +70,7 @@ def cut(
         raise ValueError("Cannot cut empty array.")
 
     if not isinstance(x, bigframes.series.Series):
-        x = bigframes.series.Series(x)
+        x = bigframes.series.Series(x, session=session)
 
     if isinstance(bins, int):
         if bins <= 0:

bigframes/core/tools/datetimes.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from collections.abc import Mapping
 from datetime import date, datetime
-from typing import Optional, Union
+from typing import Optional, TYPE_CHECKING, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.tools.datetimes as vendored_pandas_datetimes
@@ -25,6 +27,9 @@
 import bigframes.operations as ops
 import bigframes.series
 
+if TYPE_CHECKING:
+    import bigframes.session
+
 
 def to_datetime(
     arg: Union[
@@ -37,6 +42,7 @@ def to_datetime(
     utc: bool = False,
     format: Optional[str] = None,
     unit: Optional[str] = None,
+    session: Optional[bigframes.session.Session] = None,
 ) -> Union[pd.Timestamp, datetime, bigframes.series.Series]:
     if isinstance(arg, (int, float, str, datetime, date)):
         return pd.to_datetime(
@@ -52,7 +58,7 @@ def to_datetime(
             f"to datetime is not implemented. {constants.FEEDBACK_LINK}"
         )
 
-    arg = bigframes.series.Series(arg)
+    arg = bigframes.series.Series(arg, session=session)
 
     if format and unit and arg.dtype in (bigframes.dtypes.INT_DTYPE, bigframes.dtypes.FLOAT_DTYPE):  # type: ignore
         raise ValueError("cannot specify both format and unit")

bigframes/formatting_helpers.py

@@ -105,8 +105,14 @@ def progress_callback(
     """Displays a progress bar while the query is running"""
     global current_display, current_display_id, previous_display_html
 
-    import bigframes._config
-    import bigframes.core.events
+    try:
+        import bigframes._config
+        import bigframes.core.events
+    except ImportError:
+        # Since this gets called from __del__, skip if the import fails to avoid
+        # ImportError: sys.meta_path is None, Python is likely shutting down.
+        # This will allow cleanup to continue.
+        return
 
     progress_bar = bigframes._config.options.display.progress_bar
 

bigframes/operations/blob.py

@@ -804,35 +804,29 @@ def audio_transcribe(
             raise ValueError("Must specify the engine, supported value is 'bigquery'.")
 
         import bigframes.bigquery as bbq
-        import bigframes.ml.llm as llm
         import bigframes.pandas as bpd
 
         # col name doesn't matter here. Rename to avoid column name conflicts
         audio_series = bigframes.series.Series(self._block)
 
         prompt_text = "**Task:** Transcribe the provided audio. **Instructions:** - Your response must contain only the verbatim transcription of the audio. - Do not include any introductory text, summaries, or conversational filler in your response. The output should begin directly with the first word of the audio."
 
-        llm_model = llm.GeminiTextGenerator(
-            model_name=model_name,
-            session=self._block.session,
-            connection_name=connection,
-        )
+        # Convert the audio series to the runtime representation required by the model.
+        audio_runtime = audio_series.blob._get_runtime("R", with_metadata=True)
 
-        # transcribe audio using ML.GENERATE_TEXT
-        transcribed_results = llm_model.predict(
-            X=audio_series,
-            prompt=[prompt_text, audio_series],
-            temperature=0.0,
+        transcribed_results = bbq.ai.generate(
+            prompt=(prompt_text, audio_runtime),
+            connection_id=connection,
+            endpoint=model_name,
+            model_params={"generationConfig": {"temperature": 0.0}},
         )
 
-        transcribed_content_series = cast(
-            bpd.Series, transcribed_results["ml_generate_text_llm_result"]
-        ).rename("transcribed_content")
+        transcribed_content_series = transcribed_results.struct.field("result").rename(
+            "transcribed_content"
+        )
 
         if verbose:
-            transcribed_status_series = cast(
-                bpd.Series, transcribed_results["ml_generate_text_status"]
-            )
+            transcribed_status_series = transcribed_results.struct.field("status")
             results_df = bpd.DataFrame(
                 {
                     "status": transcribed_status_series,