feat: Add experimental polars execution (#1747)

Author: TrevorBergeron

bigframes/_config/bigquery_options.py

@@ -22,6 +22,7 @@
 import google.auth.credentials
 import requests.adapters
 
+import bigframes._importing
 import bigframes.enums
 import bigframes.exceptions as bfe
 
@@ -94,6 +95,7 @@ def __init__(
         requests_transport_adapters: Sequence[
             Tuple[str, requests.adapters.BaseAdapter]
         ] = (),
+        enable_polars_execution: bool = False,
     ):
         self._credentials = credentials
         self._project = project
@@ -113,6 +115,9 @@ def __init__(
             client_endpoints_override = {}
 
         self._client_endpoints_override = client_endpoints_override
+        if enable_polars_execution:
+            bigframes._importing.import_polars()
+        self._enable_polars_execution = enable_polars_execution
 
     @property
     def application_name(self) -> Optional[str]:
@@ -424,3 +429,22 @@ def requests_transport_adapters(
                 SESSION_STARTED_MESSAGE.format(attribute="requests_transport_adapters")
             )
         self._requests_transport_adapters = value
+
+    @property
+    def enable_polars_execution(self) -> bool:
+        """If True, will use polars to execute some simple query plans locally."""
+        return self._enable_polars_execution
+
+    @enable_polars_execution.setter
+    def enable_polars_execution(self, value: bool):
+        if self._session_started and self._enable_polars_execution != value:
+            raise ValueError(
+                SESSION_STARTED_MESSAGE.format(attribute="enable_polars_execution")
+            )
+        if value is True:
+            msg = bfe.format_message(
+                "Polars execution is an experimental feature, and may not be stable. Must have polars installed."
+            )
+            warnings.warn(msg, category=bfe.PreviewWarning)
+            bigframes._importing.import_polars()
+        self._enable_polars_execution = value

bigframes/_importing.py

@@ -0,0 +1,30 @@
+# Copyright 2025 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.
+import importlib
+from types import ModuleType
+
+from packaging import version
+
+# Keep this in sync with setup.py
+POLARS_MIN_VERSION = version.Version("1.7.0")
+
+
+def import_polars() -> ModuleType:
+    polars_module = importlib.import_module("polars")
+    imported_version = version.Version(polars_module.build_info()["version"])
+    if imported_version < POLARS_MIN_VERSION:
+        raise ImportError(
+            f"Imported polars version: {imported_version} is below the minimum version: {POLARS_MIN_VERSION}"
+        )
+    return polars_module

bigframes/core/compile/polars/compiler.py

@@ -393,15 +393,15 @@ class PolarsCompiler:
     expr_compiler = PolarsExpressionCompiler()
     agg_compiler = PolarsAggregateCompiler()
 
-    def compile(self, array_value: bigframes.core.ArrayValue) -> pl.LazyFrame:
+    def compile(self, plan: nodes.BigFrameNode) -> pl.LazyFrame:
         if not polars_installed:
             raise ValueError(
                 "Polars is not installed, cannot compile to polars engine."
             )
 
         # TODO: Create standard way to configure BFET -> BFET rewrites
         # Polars has incomplete slice support in lazy mode
-        node = array_value.node
+        node = plan
         node = bigframes.core.rewrite.column_pruning(node)
         node = nodes.bottom_up(node, bigframes.core.rewrite.rewrite_slice)
         node = bigframes.core.rewrite.pull_out_window_order(node)

bigframes/session/__init__.py

@@ -255,6 +255,7 @@ def __init__(
             storage_manager=self._temp_storage_manager,
             strictly_ordered=self._strictly_ordered,
             metrics=self._metrics,
+            enable_polars_execution=context.enable_polars_execution,
         )
 
     def __del__(self):

bigframes/session/bq_caching_executor.py

@@ -41,7 +41,13 @@
 import bigframes.core.tree_properties as tree_properties
 import bigframes.dtypes
 import bigframes.features
-from bigframes.session import executor, loader, local_scan_executor, read_api_execution
+from bigframes.session import (
+    executor,
+    loader,
+    local_scan_executor,
+    read_api_execution,
+    semi_executor,
+)
 import bigframes.session._io.bigquery as bq_io
 import bigframes.session.metrics
 import bigframes.session.planner
@@ -147,6 +153,7 @@ def __init__(
         *,
         strictly_ordered: bool = True,
         metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
+        enable_polars_execution: bool = False,
     ):
         self.bqclient = bqclient
         self.storage_manager = storage_manager
@@ -155,14 +162,21 @@ def __init__(
         self.metrics = metrics
         self.loader = loader
         self.bqstoragereadclient = bqstoragereadclient
-        # Simple left-to-right precedence for now
-        self._semi_executors = (
+        self._enable_polars_execution = enable_polars_execution
+        self._semi_executors: Sequence[semi_executor.SemiExecutor] = (
             read_api_execution.ReadApiSemiExecutor(
                 bqstoragereadclient=bqstoragereadclient,
                 project=self.bqclient.project,
             ),
             local_scan_executor.LocalScanExecutor(),
         )
+        if enable_polars_execution:
+            from bigframes.session import polars_executor
+
+            self._semi_executors = (
+                *self._semi_executors,
+                polars_executor.PolarsExecutor(),
+            )
         self._upload_lock = threading.Lock()
 
     def to_sql(
@@ -637,8 +651,8 @@ def _execute_plan(
         """Just execute whatever plan as is, without further caching or decomposition."""
         # First try to execute fast-paths
         if not output_spec.require_bq_table:
-            for semi_executor in self._semi_executors:
-                maybe_result = semi_executor.execute(plan, ordered=ordered, peek=peek)
+            for exec in self._semi_executors:
+                maybe_result = exec.execute(plan, ordered=ordered, peek=peek)
                 if maybe_result:
                     return maybe_result
 

bigframes/session/polars_executor.py

@@ -73,7 +73,7 @@ def execute(
         # Note: Ignoring ordered flag, as just executing totally ordered is fine.
         try:
             lazy_frame: pl.LazyFrame = self._compiler.compile(
-                array_value.ArrayValue(plan)
+                array_value.ArrayValue(plan).node
             )
         except Exception:
             return None

bigframes/testing/polars_session.py

@@ -41,7 +41,7 @@ def peek(
         """
         A 'peek' efficiently accesses a small number of rows in the dataframe.
         """
-        lazy_frame: polars.LazyFrame = self.compiler.compile(array_value)
+        lazy_frame: polars.LazyFrame = self.compiler.compile(array_value.node)
         pa_table = lazy_frame.collect().limit(n_rows).to_arrow()
         # Currently, pyarrow types might not quite be exactly the ones in the bigframes schema.
         # Nullability may be different, and might use large versions of list, string datatypes.
@@ -64,7 +64,7 @@ def execute(
         """
         Execute the ArrayValue, storing the result to a temporary session-owned table.
         """
-        lazy_frame: polars.LazyFrame = self.compiler.compile(array_value)
+        lazy_frame: polars.LazyFrame = self.compiler.compile(array_value.node)
         pa_table = lazy_frame.collect().to_arrow()
         # Currently, pyarrow types might not quite be exactly the ones in the bigframes schema.
         # Nullability may be different, and might use large versions of list, string datatypes.

noxfile.py

@@ -108,7 +108,7 @@
 SYSTEM_TEST_EXTRAS: List[str] = []
 SYSTEM_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = {
     "3.9": ["tests", "anywidget"],
-    "3.10": ["tests"],
+    "3.10": ["tests", "polars"],
     "3.12": ["tests", "scikit-learn", "polars", "anywidget"],
     "3.13": ["tests", "polars"],
 }

testing/constraints-3.10.txt

@@ -15,3 +15,4 @@ matplotlib==3.7.1
 psutil==5.9.5
 seaborn==0.13.1
 traitlets==5.7.1
+polars==1.7.0

tests/system/small/test_polars_execution.py

@@ -0,0 +1,76 @@
+# Copyright 2025 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.
+import pytest
+
+import bigframes
+from bigframes.testing.utils import assert_pandas_df_equal
+
+polars = pytest.importorskip("polars", reason="polars is required for this test")
+
+
+@pytest.fixture(scope="module")
+def session_w_polars():
+    context = bigframes.BigQueryOptions(location="US", enable_polars_execution=True)
+    session = bigframes.Session(context=context)
+    yield session
+    session.close()  # close generated session at cleanup time
+
+
+def test_polar_execution_sorted(session_w_polars, scalars_pandas_df_index):
+    execution_count_before = session_w_polars._metrics.execution_count
+    bf_df = session_w_polars.read_pandas(scalars_pandas_df_index)
+
+    pd_result = scalars_pandas_df_index.sort_index(ascending=False)[
+        ["int64_too", "bool_col"]
+    ]
+    bf_result = bf_df.sort_index(ascending=False)[["int64_too", "bool_col"]].to_pandas()
+
+    assert session_w_polars._metrics.execution_count == execution_count_before
+    assert_pandas_df_equal(bf_result, pd_result)
+
+
+def test_polar_execution_sorted_filtered(session_w_polars, scalars_pandas_df_index):
+    execution_count_before = session_w_polars._metrics.execution_count
+    bf_df = session_w_polars.read_pandas(scalars_pandas_df_index)
+
+    pd_result = scalars_pandas_df_index.sort_index(ascending=False).dropna(
+        subset=["int64_col", "string_col"]
+    )
+    bf_result = (
+        bf_df.sort_index(ascending=False)
+        .dropna(subset=["int64_col", "string_col"])
+        .to_pandas()
+    )
+
+    # Filter and isnull not supported by polar engine yet, so falls back to bq execution
+    assert session_w_polars._metrics.execution_count == (execution_count_before + 1)
+    assert_pandas_df_equal(bf_result, pd_result)
+
+
+def test_polar_execution_unsupported_sql_fallback(
+    session_w_polars, scalars_pandas_df_index
+):
+    execution_count_before = session_w_polars._metrics.execution_count
+    bf_df = session_w_polars.read_pandas(scalars_pandas_df_index)
+
+    pd_df = scalars_pandas_df_index.copy()
+    pd_df["str_len_col"] = pd_df.string_col.str.len()
+    pd_result = pd_df
+
+    bf_df["str_len_col"] = bf_df.string_col.str.len()
+    bf_result = bf_df.to_pandas()
+
+    # str len not supported by polar engine yet, so falls back to bq execution
+    assert session_w_polars._metrics.execution_count == (execution_count_before + 1)
+    assert_pandas_df_equal(bf_result, pd_result)