Skip to content

Commit 7eba6ee

Browse files
authored
feat: add bigframes.pandas.options.experiments.sql_compiler for switching the backend compiler (#2417)
This change adds `bigframes.pandas.options.experiments.sql_compiler` to allow switching the backend compiler. Currently, the default remains set to 'legacy' (ibis), but users can now optionally switch to the 'experimental' (sqlglot) compiler. Fixes internal issue 479912001🦕
1 parent cfa1c8e commit 7eba6ee

File tree

5 files changed

+64
-8
lines changed

5 files changed

+64
-8
lines changed

bigframes/_config/experiment_options.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from typing import Optional
15+
from typing import Literal, Optional
1616
import warnings
1717

1818
import bigframes
@@ -27,6 +27,7 @@ class ExperimentOptions:
2727
def __init__(self):
2828
self._semantic_operators: bool = False
2929
self._ai_operators: bool = False
30+
self._sql_compiler: Literal["legacy", "stable", "experimental"] = "stable"
3031

3132
@property
3233
def semantic_operators(self) -> bool:
@@ -55,6 +56,24 @@ def ai_operators(self, value: bool):
5556
warnings.warn(msg, category=bfe.PreviewWarning)
5657
self._ai_operators = value
5758

59+
@property
60+
def sql_compiler(self) -> Literal["legacy", "stable", "experimental"]:
61+
return self._sql_compiler
62+
63+
@sql_compiler.setter
64+
def sql_compiler(self, value: Literal["legacy", "stable", "experimental"]):
65+
if value not in ["legacy", "stable", "experimental"]:
66+
raise ValueError(
67+
"sql_compiler must be one of 'legacy', 'stable', or 'experimental'"
68+
)
69+
if value == "experimental":
70+
msg = bfe.format_message(
71+
"The experimental SQL compiler is still under experiments, and is subject "
72+
"to change in the future."
73+
)
74+
warnings.warn(msg, category=FutureWarning)
75+
self._sql_compiler = value
76+
5877
@property
5978
def blob(self) -> bool:
6079
msg = bfe.format_message(

bigframes/core/compile/__init__.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,28 @@
1313
# limitations under the License.
1414
from __future__ import annotations
1515

16+
from typing import Any
17+
18+
from bigframes import options
1619
from bigframes.core.compile.api import test_only_ibis_inferred_schema
1720
from bigframes.core.compile.configs import CompileRequest, CompileResult
18-
from bigframes.core.compile.ibis_compiler.ibis_compiler import compile_sql
21+
22+
23+
def compiler() -> Any:
24+
"""Returns the appropriate compiler module based on session options."""
25+
if options.experiments.sql_compiler == "experimental":
26+
import bigframes.core.compile.sqlglot.compiler as sqlglot_compiler
27+
28+
return sqlglot_compiler
29+
else:
30+
import bigframes.core.compile.ibis_compiler.ibis_compiler as ibis_compiler
31+
32+
return ibis_compiler
33+
1934

2035
__all__ = [
2136
"test_only_ibis_inferred_schema",
22-
"compile_sql",
2337
"CompileRequest",
2438
"CompileResult",
39+
"compiler",
2540
]

bigframes/session/bq_caching_executor.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,9 @@ def to_sql(
174174
else array_value.node
175175
)
176176
node = self._substitute_large_local_sources(node)
177-
compiled = compile.compile_sql(compile.CompileRequest(node, sort_rows=ordered))
177+
compiled = compile.compiler().compile_sql(
178+
compile.CompileRequest(node, sort_rows=ordered)
179+
)
178180
return compiled.sql
179181

180182
def execute(
@@ -290,7 +292,9 @@ def _export_gbq(
290292
# validate destination table
291293
existing_table = self._maybe_find_existing_table(spec)
292294

293-
compiled = compile.compile_sql(compile.CompileRequest(plan, sort_rows=False))
295+
compiled = compile.compiler().compile_sql(
296+
compile.CompileRequest(plan, sort_rows=False)
297+
)
294298
sql = compiled.sql
295299

296300
if (existing_table is not None) and _if_schema_match(
@@ -641,7 +645,7 @@ def _execute_plan_gbq(
641645
]
642646
cluster_cols = cluster_cols[:_MAX_CLUSTER_COLUMNS]
643647

644-
compiled = compile.compile_sql(
648+
compiled = compile.compiler().compile_sql(
645649
compile.CompileRequest(
646650
plan,
647651
sort_rows=ordered,

bigframes/session/direct_gbq_execution.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@
2020
import google.cloud.bigquery.table as bq_table
2121

2222
from bigframes.core import compile, nodes
23-
from bigframes.core.compile import sqlglot
23+
import bigframes.core.compile.ibis_compiler.ibis_compiler as ibis_compiler
24+
import bigframes.core.compile.sqlglot.compiler as sqlglot_compiler
2425
import bigframes.core.events
2526
from bigframes.session import executor, semi_executor
2627
import bigframes.session._io.bigquery as bq_io
@@ -40,7 +41,9 @@ def __init__(
4041
):
4142
self.bqclient = bqclient
4243
self._compile_fn = (
43-
compile.compile_sql if compiler == "ibis" else sqlglot.compile_sql
44+
ibis_compiler.compile_sql
45+
if compiler == "ibis"
46+
else sqlglot_compiler.compile_sql
4447
)
4548
self._publisher = publisher
4649

tests/unit/_config/test_experiment_options.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,3 +46,18 @@ def test_ai_operators_set_true_shows_warning():
4646
options.ai_operators = True
4747

4848
assert options.ai_operators is True
49+
50+
51+
def test_sql_compiler_default_stable():
52+
options = experiment_options.ExperimentOptions()
53+
54+
assert options.sql_compiler == "stable"
55+
56+
57+
def test_sql_compiler_set_experimental_shows_warning():
58+
options = experiment_options.ExperimentOptions()
59+
60+
with pytest.warns(FutureWarning):
61+
options.sql_compiler = "experimental"
62+
63+
assert options.sql_compiler == "experimental"

0 commit comments

Comments
 (0)