Merge branch 'main' into sycai_type_usage

Author: sycai

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/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/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(
@@ -643,7 +647,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,

bigframes/session/direct_gbq_execution.py

@@ -20,7 +20,8 @@
 import google.cloud.bigquery.table as bq_table
 
 from bigframes.core import compile, nodes
-from bigframes.core.compile import sqlglot
+import bigframes.core.compile.ibis_compiler.ibis_compiler as ibis_compiler
+import bigframes.core.compile.sqlglot.compiler as sqlglot_compiler
 import bigframes.core.events
 from bigframes.session import executor, semi_executor
 import bigframes.session._io.bigquery as bq_io
@@ -40,7 +41,9 @@ def __init__(
     ):
         self.bqclient = bqclient
         self._compile_fn = (
-            compile.compile_sql if compiler == "ibis" else sqlglot.compile_sql
+            ibis_compiler.compile_sql
+            if compiler == "ibis"
+            else sqlglot_compiler.compile_sql
         )
         self._publisher = publisher
 

tests/unit/_config/test_experiment_options.py

@@ -46,3 +46,18 @@ def test_ai_operators_set_true_shows_warning():
         options.ai_operators = True
 
     assert options.ai_operators is True
+
+
+def test_sql_compiler_default_stable():
+    options = experiment_options.ExperimentOptions()
+
+    assert options.sql_compiler == "stable"
+
+
+def test_sql_compiler_set_experimental_shows_warning():
+    options = experiment_options.ExperimentOptions()
+
+    with pytest.warns(FutureWarning):
+        options.sql_compiler = "experimental"
+
+    assert options.sql_compiler == "experimental"