Merge branch 'main' into shuowei-anywidget-complete-status-update

Author: shuoweil

bigframes/_config/experiment_options.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Literal, Optional
 import warnings
 
 import bigframes
@@ -27,6 +27,7 @@ class ExperimentOptions:
     def __init__(self):
         self._semantic_operators: bool = False
         self._ai_operators: bool = False
+        self._sql_compiler: Literal["legacy", "stable", "experimental"] = "stable"
 
     @property
     def semantic_operators(self) -> bool:
@@ -55,6 +56,24 @@ def ai_operators(self, value: bool):
             warnings.warn(msg, category=bfe.PreviewWarning)
         self._ai_operators = value
 
+    @property
+    def sql_compiler(self) -> Literal["legacy", "stable", "experimental"]:
+        return self._sql_compiler
+
+    @sql_compiler.setter
+    def sql_compiler(self, value: Literal["legacy", "stable", "experimental"]):
+        if value not in ["legacy", "stable", "experimental"]:
+            raise ValueError(
+                "sql_compiler must be one of 'legacy', 'stable', or 'experimental'"
+            )
+        if value == "experimental":
+            msg = bfe.format_message(
+                "The experimental SQL compiler is still under experiments, and is subject "
+                "to change in the future."
+            )
+            warnings.warn(msg, category=FutureWarning)
+        self._sql_compiler = value
+
     @property
     def blob(self) -> bool:
         msg = bfe.format_message(

bigframes/bigquery/__init__.py

@@ -60,7 +60,7 @@
 from bigframes.bigquery._operations.search import create_vector_index, vector_search
 from bigframes.bigquery._operations.sql import sql_scalar
 from bigframes.bigquery._operations.struct import struct
-from bigframes.bigquery.table import create_external_table
+from bigframes.bigquery._operations.table import create_external_table
 from bigframes.core.logging import log_adapter
 
 _functions = [

bigframes/bigquery/_operations/ml.py

@@ -520,3 +520,63 @@ def generate_text(
         return bpd.read_gbq_query(sql)
     else:
         return session.read_gbq_query(sql)
+
+
+@log_adapter.method_logger(custom_base_name="bigquery_ml")
+def generate_embedding(
+    model: Union[bigframes.ml.base.BaseEstimator, str, pd.Series],
+    input_: Union[pd.DataFrame, dataframe.DataFrame, str],
+    *,
+    flatten_json_output: Optional[bool] = None,
+    task_type: Optional[str] = None,
+    output_dimensionality: Optional[int] = None,
+) -> dataframe.DataFrame:
+    """
+    Generates text embedding using a BigQuery ML model.
+
+    See the `BigQuery ML GENERATE_EMBEDDING function syntax
+    <https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-generate-embedding>`_
+    for additional reference.
+
+    Args:
+        model (bigframes.ml.base.BaseEstimator or str):
+            The model to use for text embedding.
+        input_ (Union[bigframes.pandas.DataFrame, str]):
+            The DataFrame or query to use for text embedding.
+        flatten_json_output (bool, optional):
+            A BOOL value that determines the content of the generated JSON column.
+        task_type (str, optional):
+            A STRING value that specifies the intended downstream application task.
+            Supported values are:
+            - `RETRIEVAL_QUERY`
+            - `RETRIEVAL_DOCUMENT`
+            - `SEMANTIC_SIMILARITY`
+            - `CLASSIFICATION`
+            - `CLUSTERING`
+            - `QUESTION_ANSWERING`
+            - `FACT_VERIFICATION`
+            - `CODE_RETRIEVAL_QUERY`
+        output_dimensionality (int, optional):
+            An INT64 value that specifies the size of the output embedding.
+
+    Returns:
+        bigframes.pandas.DataFrame:
+            The generated text embedding.
+    """
+    import bigframes.pandas as bpd
+
+    model_name, session = _get_model_name_and_session(model, input_)
+    table_sql = _to_sql(input_)
+
+    sql = bigframes.core.sql.ml.generate_embedding(
+        model_name=model_name,
+        table=table_sql,
+        flatten_json_output=flatten_json_output,
+        task_type=task_type,
+        output_dimensionality=output_dimensionality,
+    )
+
+    if session is None:
+        return bpd.read_gbq_query(sql)
+    else:
+        return session.read_gbq_query(sql)

bigframes/bigquery/table.py bigframes/bigquery/_operations/table.pybigframes/bigquery/table.py renamed to bigframes/bigquery/_operations/table.py

@@ -16,7 +16,6 @@
 
 from typing import Mapping, Optional, Union
 
-import bigframes_vendored.constants
 import google.cloud.bigquery
 import pandas as pd
 
@@ -94,9 +93,6 @@ def create_external_table(
     if session is None:
         bpd.read_gbq_query(sql)
         session = bpd.get_global_session()
-        assert (
-            session is not None
-        ), f"Missing connection to BigQuery. Please report how you encountered this error at {bigframes_vendored.constants.FEEDBACK_LINK}."
     else:
         session.read_gbq_query(sql)
 

bigframes/bigquery/ml.py

@@ -23,6 +23,7 @@
     create_model,
     evaluate,
     explain_predict,
+    generate_embedding,
     generate_text,
     global_explain,
     predict,
@@ -37,4 +38,5 @@
     "global_explain",
     "transform",
     "generate_text",
+    "generate_embedding",
 ]

bigframes/core/compile/__init__.py

@@ -13,13 +13,28 @@
 # limitations under the License.
 from __future__ import annotations
 
+from typing import Any
+
+from bigframes import options
 from bigframes.core.compile.api import test_only_ibis_inferred_schema
 from bigframes.core.compile.configs import CompileRequest, CompileResult
-from bigframes.core.compile.ibis_compiler.ibis_compiler import compile_sql
+
+
+def compiler() -> Any:
+    """Returns the appropriate compiler module based on session options."""
+    if options.experiments.sql_compiler == "experimental":
+        import bigframes.core.compile.sqlglot.compiler as sqlglot_compiler
+
+        return sqlglot_compiler
+    else:
+        import bigframes.core.compile.ibis_compiler.ibis_compiler as ibis_compiler
+
+        return ibis_compiler
+
 
 __all__ = [
     "test_only_ibis_inferred_schema",
-    "compile_sql",
     "CompileRequest",
     "CompileResult",
+    "compiler",
 ]

bigframes/core/sql/ml.py

@@ -296,3 +296,31 @@ def generate_text(
     sql += _build_struct_sql(struct_options)
     sql += ")\n"
     return sql
+
+
+def generate_embedding(
+    model_name: str,
+    table: str,
+    *,
+    flatten_json_output: Optional[bool] = None,
+    task_type: Optional[str] = None,
+    output_dimensionality: Optional[int] = None,
+) -> str:
+    """Encode the ML.GENERATE_EMBEDDING statement.
+    See https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-generate-embedding for reference.
+    """
+    struct_options: Dict[
+        str,
+        Union[str, int, float, bool, Mapping[str, str], List[str], Mapping[str, Any]],
+    ] = {}
+    if flatten_json_output is not None:
+        struct_options["flatten_json_output"] = flatten_json_output
+    if task_type is not None:
+        struct_options["task_type"] = task_type
+    if output_dimensionality is not None:
+        struct_options["output_dimensionality"] = output_dimensionality
+
+    sql = f"SELECT * FROM ML.GENERATE_EMBEDDING(MODEL {googlesql.identifier(model_name)}, ({table})"
+    sql += _build_struct_sql(struct_options)
+    sql += ")\n"
+    return sql

bigframes/formatting_helpers.py

@@ -25,10 +25,10 @@
 import google.api_core.exceptions as api_core_exceptions
 import google.cloud.bigquery as bigquery
 import humanize
-import IPython
-import IPython.display as display
 
 if TYPE_CHECKING:
+    from IPython import display
+
     import bigframes.core.events
 
 GenericJob = Union[
@@ -160,6 +160,8 @@ def progress_callback(
         progress_bar = "notebook" if in_ipython() else "terminal"
 
     if progress_bar == "notebook":
+        import IPython.display as display
+
         if (
             isinstance(event, bigframes.core.events.ExecutionStarted)
             or current_display is None
@@ -245,6 +247,8 @@ def wait_for_job(job: GenericJob, progress_bar: Optional[str] = None):
 
     try:
         if progress_bar == "notebook":
+            import IPython.display as display
+
             display_id = str(random.random())
             loading_bar = display.HTML(get_base_job_loading_html(job))
             display.display(loading_bar, display_id=display_id)
@@ -613,4 +617,8 @@ def get_bytes_processed_string(val: Any):
 
 def in_ipython():
     """Return True iff we're in a colab-like IPython."""
+    try:
+        import IPython
+    except (ImportError, NameError):
+        return False
     return hasattr(IPython.get_ipython(), "kernel")

bigframes/operations/blob.py

@@ -18,7 +18,6 @@
 from typing import cast, Literal, Optional, Union
 import warnings
 
-import IPython.display as ipy_display
 import pandas as pd
 import requests
 
@@ -241,6 +240,8 @@ def display(
             width (int or None, default None): width in pixels that the image/video are constrained to. If unset, use the global setting in bigframes.options.display.blob_display_width, otherwise image/video's original size or ratio is used. No-op for other content types.
             height (int or None, default None): height in pixels that the image/video are constrained to. If unset, use the global setting in bigframes.options.display.blob_display_height, otherwise image/video's original size or ratio is used. No-op for other content types.
         """
+        import IPython.display as ipy_display
+
         width = width or bigframes.options.display.blob_display_width
         height = height or bigframes.options.display.blob_display_height
 

bigframes/session/bq_caching_executor.py

@@ -174,7 +174,9 @@ def to_sql(
             else array_value.node
         )
         node = self._substitute_large_local_sources(node)
-        compiled = compile.compile_sql(compile.CompileRequest(node, sort_rows=ordered))
+        compiled = compile.compiler().compile_sql(
+            compile.CompileRequest(node, sort_rows=ordered)
+        )
         return compiled.sql
 
     def execute(
@@ -290,7 +292,9 @@ def _export_gbq(
         # validate destination table
         existing_table = self._maybe_find_existing_table(spec)
 
-        compiled = compile.compile_sql(compile.CompileRequest(plan, sort_rows=False))
+        compiled = compile.compiler().compile_sql(
+            compile.CompileRequest(plan, sort_rows=False)
+        )
         sql = compiled.sql
 
         if (existing_table is not None) and _if_schema_match(
@@ -641,7 +645,7 @@ def _execute_plan_gbq(
                 ]
                 cluster_cols = cluster_cols[:_MAX_CLUSTER_COLUMNS]
 
-        compiled = compile.compile_sql(
+        compiled = compile.compiler().compile_sql(
             compile.CompileRequest(
                 plan,
                 sort_rows=ordered,