Merge branch 'main' into main_chelsealin_compilejoin

Author: chelsea-lin

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/array_value.py

@@ -330,12 +330,27 @@ def create_constant(
 
         return self.project_to_id(ex.const(value, dtype))
 
-    def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue:
+    def select_columns(
+        self, column_ids: typing.Sequence[str], allow_renames: bool = False
+    ) -> ArrayValue:
         # This basically just drops and reorders columns - logically a no-op except as a final step
-        selections = (
-            bigframes.core.nodes.AliasedRef.identity(ids.ColumnId(col_id))
-            for col_id in column_ids
-        )
+        selections = []
+        seen = set()
+
+        for id in column_ids:
+            if id not in seen:
+                ref = nodes.AliasedRef.identity(ids.ColumnId(id))
+            elif allow_renames:
+                ref = nodes.AliasedRef(
+                    ex.deref(id), ids.ColumnId(bigframes.core.guid.generate_guid())
+                )
+            else:
+                raise ValueError(
+                    "Must set allow_renames=True to select columns repeatedly"
+                )
+            selections.append(ref)
+            seen.add(id)
+
         return ArrayValue(
             nodes.SelectionNode(
                 child=self.node,

bigframes/core/blocks.py

@@ -50,6 +50,7 @@
 import bigframes.core.identifiers
 import bigframes.core.join_def as join_defs
 import bigframes.core.ordering as ordering
+import bigframes.core.pyarrow_utils as pyarrow_utils
 import bigframes.core.schema as bf_schema
 import bigframes.core.sql as sql
 import bigframes.core.utils as utils
@@ -156,6 +157,36 @@ def __init__(
         self._view_ref: Optional[bigquery.TableReference] = None
         self._view_ref_dry_run: Optional[bigquery.TableReference] = None
 
+    @classmethod
+    def from_pyarrow(
+        cls,
+        data: pa.Table,
+        session: bigframes.Session,
+    ) -> Block:
+        column_labels = data.column_names
+
+        # TODO(tswast): Use array_value.promote_offsets() instead once that node is
+        # supported by the local engine.
+        offsets_col = bigframes.core.guid.generate_guid()
+        index_ids = [offsets_col]
+        index_labels = [None]
+
+        # TODO(https://github.com/googleapis/python-bigquery-dataframes/issues/859):
+        # Allow users to specify the "total ordering" column(s) or allow multiple
+        # such columns.
+        data = pyarrow_utils.append_offsets(data, offsets_col=offsets_col)
+
+        # from_pyarrow will normalize the types for us.
+        managed_data = local_data.ManagedArrowTable.from_pyarrow(data)
+        array_value = core.ArrayValue.from_managed(managed_data, session=session)
+        block = cls(
+            array_value,
+            column_labels=column_labels,
+            index_columns=index_ids,
+            index_labels=index_labels,
+        )
+        return block
+
     @classmethod
     def from_local(
         cls,
@@ -1210,7 +1241,10 @@ def select_column(self, id: str) -> Block:
         return self.select_columns([id])
 
     def select_columns(self, ids: typing.Sequence[str]) -> Block:
-        expr = self._expr.select_columns([*self.index_columns, *ids])
+        # Allow renames as may end up selecting same columns multiple times
+        expr = self._expr.select_columns(
+            [*self.index_columns, *ids], allow_renames=True
+        )
         col_labels = self._get_labels_for_columns(ids)
         return Block(expr, self.index_columns, col_labels, self.index.names)
 
@@ -1996,7 +2030,7 @@ def _generate_resample_label(
         return block.set_index([resample_label_id])
 
     def _create_stack_column(self, col_label: typing.Tuple, stack_labels: pd.Index):
-        dtype = None
+        input_dtypes = []
         input_columns: list[Optional[str]] = []
         for uvalue in utils.index_as_tuples(stack_labels):
             label_to_match = (*col_label, *uvalue)
@@ -2006,15 +2040,18 @@ def _create_stack_column(self, col_label: typing.Tuple, stack_labels: pd.Index):
             matching_ids = self.label_to_col_id.get(label_to_match, [])
             input_id = matching_ids[0] if len(matching_ids) > 0 else None
             if input_id:
-                if dtype and dtype != self._column_type(input_id):
-                    raise NotImplementedError(
-                        "Cannot stack columns with non-matching dtypes."
-                    )
-                else:
-                    dtype = self._column_type(input_id)
+                input_dtypes.append(self._column_type(input_id))
             input_columns.append(input_id)
             # Input column i is the first one that
-        return tuple(input_columns), dtype or pd.Float64Dtype()
+        if len(input_dtypes) > 0:
+            output_dtype = bigframes.dtypes.lcd_type(*input_dtypes)
+            if output_dtype is None:
+                raise NotImplementedError(
+                    "Cannot stack columns with non-matching dtypes."
+                )
+        else:
+            output_dtype = pd.Float64Dtype()
+        return tuple(input_columns), output_dtype
 
     def _column_type(self, col_id: str) -> bigframes.dtypes.Dtype:
         col_offset = self.value_columns.index(col_id)

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/core/compile/polars/lowering.py

@@ -12,15 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import dataclasses
+
 from bigframes import dtypes
 from bigframes.core import bigframe_node, expression
 from bigframes.core.rewrite import op_lowering
-from bigframes.operations import numeric_ops
+from bigframes.operations import comparison_ops, numeric_ops
 import bigframes.operations as ops
 
 # TODO: Would be more precise to actually have separate op set for polars ops (where they diverge from the original ops)
 
 
+@dataclasses.dataclass
+class CoerceArgsRule(op_lowering.OpLoweringRule):
+    op_type: type[ops.BinaryOp]
+
+    @property
+    def op(self) -> type[ops.ScalarOp]:
+        return self.op_type
+
+    def lower(self, expr: expression.OpExpression) -> expression.Expression:
+        assert isinstance(expr.op, self.op_type)
+        larg, rarg = _coerce_comparables(expr.children[0], expr.children[1])
+        return expr.op.as_expr(larg, rarg)
+
+
 class LowerFloorDivRule(op_lowering.OpLoweringRule):
     @property
     def op(self) -> type[ops.ScalarOp]:
@@ -40,7 +56,42 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
         return ops.where_op.as_expr(zero_result, divisor_is_zero, expr)
 
 
-POLARS_LOWERING_RULES = (LowerFloorDivRule(),)
+def _coerce_comparables(expr1: expression.Expression, expr2: expression.Expression):
+
+    target_type = dtypes.coerce_to_common(expr1.output_type, expr2.output_type)
+    if expr1.output_type != target_type:
+        expr1 = _lower_cast(ops.AsTypeOp(target_type), expr1)
+    if expr2.output_type != target_type:
+        expr2 = _lower_cast(ops.AsTypeOp(target_type), expr2)
+    return expr1, expr2
+
+
+# TODO: Need to handle bool->string cast to get capitalization correct
+def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
+    if arg.output_type == dtypes.BOOL_DTYPE and dtypes.is_numeric(cast_op.to_type):
+        # bool -> decimal needs two-step cast
+        new_arg = ops.AsTypeOp(to_type=dtypes.INT_DTYPE).as_expr(arg)
+        return cast_op.as_expr(new_arg)
+    return cast_op.as_expr(arg)
+
+
+LOWER_COMPARISONS = tuple(
+    CoerceArgsRule(op)
+    for op in (
+        comparison_ops.EqOp,
+        comparison_ops.EqNullsMatchOp,
+        comparison_ops.NeOp,
+        comparison_ops.LtOp,
+        comparison_ops.GtOp,
+        comparison_ops.LeOp,
+        comparison_ops.GeOp,
+    )
+)
+
+POLARS_LOWERING_RULES = (
+    *LOWER_COMPARISONS,
+    LowerFloorDivRule(),
+)
 
 
 def lower_ops_to_polars(root: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode:

bigframes/core/compile/scalar_op_compiler.py

@@ -1498,6 +1498,7 @@ def eq_op(
     x: ibis_types.Value,
     y: ibis_types.Value,
 ):
+    x, y = _coerce_comparables(x, y)
     return x == y
 
 
@@ -1507,6 +1508,7 @@ def eq_nulls_match_op(
     y: ibis_types.Value,
 ):
     """Variant of eq_op where nulls match each other. Only use where dtypes are known to be same."""
+    x, y = _coerce_comparables(x, y)
     literal = ibis_types.literal("$NULL_SENTINEL$")
     if hasattr(x, "fill_null"):
         left = x.cast(ibis_dtypes.str).fill_null(literal)
@@ -1523,6 +1525,7 @@ def ne_op(
     x: ibis_types.Value,
     y: ibis_types.Value,
 ):
+    x, y = _coerce_comparables(x, y)
     return x != y
 
 
@@ -1534,6 +1537,17 @@ def _null_or_value(value: ibis_types.Value, where_value: ibis_types.BooleanValue
     )
 
 
+def _coerce_comparables(
+    x: ibis_types.Value,
+    y: ibis_types.Value,
+):
+    if x.type().is_boolean() and not y.type().is_boolean():
+        x = x.cast(ibis_dtypes.int64)
+    elif y.type().is_boolean() and not x.type().is_boolean():
+        y = y.cast(ibis_dtypes.int64)
+    return x, y
+
+
 @scalar_op_compiler.register_binary_op(ops.and_op)
 def and_op(
     x: ibis_types.Value,
@@ -1735,6 +1749,7 @@ def lt_op(
     x: ibis_types.Value,
     y: ibis_types.Value,
 ):
+    x, y = _coerce_comparables(x, y)
     return x < y
 
 
@@ -1744,6 +1759,7 @@ def le_op(
     x: ibis_types.Value,
     y: ibis_types.Value,
 ):
+    x, y = _coerce_comparables(x, y)
     return x <= y
 
 
@@ -1753,6 +1769,7 @@ def gt_op(
     x: ibis_types.Value,
     y: ibis_types.Value,
 ):
+    x, y = _coerce_comparables(x, y)
     return x > y
 
 
@@ -1762,6 +1779,7 @@ def ge_op(
     x: ibis_types.Value,
     y: ibis_types.Value,
 ):
+    x, y = _coerce_comparables(x, y)
     return x >= y
 
 

bigframes/core/indexes/base.py

@@ -174,6 +174,11 @@ def dtypes(self) -> pandas.Series:
             index=typing.cast(typing.Tuple, self._block.index.names),
         )
 
+    def __setitem__(self, key, value) -> None:
+        """Index objects are immutable. Use Index constructor to create
+        modified Index."""
+        raise TypeError("Index does not support mutable operations")
+
     @property
     def size(self) -> int:
         return self.shape[0]