Merge branch 'main' into join_suffix

Author: Genesis929

bigframes/_importing.py

@@ -14,6 +14,7 @@
 import importlib
 from types import ModuleType
 
+import numpy
 from packaging import version
 
 # Keep this in sync with setup.py
@@ -22,9 +23,13 @@
 
 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:
+    # Check for necessary methods instead of the version number because we
+    # can't trust the polars version until
+    # https://github.com/pola-rs/polars/issues/23940 is fixed.
+    try:
+        polars_module.lit(numpy.int64(100), dtype=polars_module.Int64())
+    except TypeError:
         raise ImportError(
-            f"Imported polars version: {imported_version} is below the minimum version: {POLARS_MIN_VERSION}"
+            f"Imported polars version is likely below the minimum version: {POLARS_MIN_VERSION}"
         )
     return polars_module

bigframes/core/block_transforms.py

@@ -355,24 +355,28 @@ def value_counts(
     normalize: bool = False,
     sort: bool = True,
     ascending: bool = False,
-    dropna: bool = True,
+    drop_na: bool = True,
+    grouping_keys: typing.Sequence[str] = (),
 ):
-    block, dummy = block.create_constant(1)
+    if grouping_keys and drop_na:
+        # only need this if grouping_keys is involved, otherwise the drop_na in the aggregation will handle it for us
+        block = dropna(block, columns, how="any")
     block, agg_ids = block.aggregate(
-        by_column_ids=columns,
-        aggregations=[ex.UnaryAggregation(agg_ops.count_op, ex.deref(dummy))],
-        dropna=dropna,
+        by_column_ids=(*grouping_keys, *columns),
+        aggregations=[ex.NullaryAggregation(agg_ops.size_op)],
+        dropna=drop_na and not grouping_keys,
     )
     count_id = agg_ids[0]
     if normalize:
-        unbound_window = windows.unbound()
+        unbound_window = windows.unbound(grouping_keys=tuple(grouping_keys))
         block, total_count_id = block.apply_window_op(
             count_id, agg_ops.sum_op, unbound_window
         )
         block, count_id = block.apply_binary_op(count_id, total_count_id, ops.div_op)
 
     if sort:
-        block = block.order_by(
+        order_parts = [ordering.ascending_over(id) for id in grouping_keys]
+        order_parts.extend(
             [
                 ordering.OrderingExpression(
                     ex.deref(count_id),
@@ -382,6 +386,7 @@ def value_counts(
                 )
             ]
         )
+        block = block.order_by(order_parts)
     return block.select_column(count_id).with_column_labels(
         ["proportion" if normalize else "count"]
     )

bigframes/core/compile/__init__.py

@@ -14,8 +14,8 @@
 from __future__ import annotations
 
 from bigframes.core.compile.api import test_only_ibis_inferred_schema
-from bigframes.core.compile.compiler import compile_sql
 from bigframes.core.compile.configs import CompileRequest, CompileResult
+from bigframes.core.compile.ibis_compiler.ibis_compiler import compile_sql
 
 __all__ = [
     "test_only_ibis_inferred_schema",

bigframes/core/compile/api.py

@@ -16,7 +16,7 @@
 from typing import TYPE_CHECKING
 
 from bigframes.core import rewrite
-from bigframes.core.compile import compiler
+from bigframes.core.compile.ibis_compiler import ibis_compiler
 
 if TYPE_CHECKING:
     import bigframes.core.nodes
@@ -26,9 +26,9 @@ def test_only_ibis_inferred_schema(node: bigframes.core.nodes.BigFrameNode):
     """Use only for testing paths to ensure ibis inferred schema does not diverge from bigframes inferred schema."""
     import bigframes.core.schema
 
-    node = compiler._replace_unsupported_ops(node)
+    node = ibis_compiler._replace_unsupported_ops(node)
     node = rewrite.bake_order(node)
-    ir = compiler.compile_node(node)
+    ir = ibis_compiler.compile_node(node)
     items = tuple(
         bigframes.core.schema.SchemaItem(name, ir.get_column_type(ibis_id))
         for name, ibis_id in zip(node.schema.names, ir.column_ids)

bigframes/core/compile/compiled.py

@@ -30,11 +30,10 @@
 import pyarrow as pa
 
 from bigframes.core import utils
-import bigframes.core.compile.aggregate_compiler as agg_compiler
 import bigframes.core.compile.googlesql
+import bigframes.core.compile.ibis_compiler.aggregate_compiler as agg_compiler
+import bigframes.core.compile.ibis_compiler.scalar_op_compiler as op_compilers
 import bigframes.core.compile.ibis_types
-import bigframes.core.compile.scalar_op_compiler as op_compilers
-import bigframes.core.compile.scalar_op_compiler as scalar_op_compiler
 import bigframes.core.expression as ex
 from bigframes.core.ordering import OrderingExpression
 import bigframes.core.sql
@@ -460,7 +459,7 @@ def project_window_op(
             for column in inputs:
                 clauses.append((column.isnull(), ibis_types.null()))
         if window_spec.min_periods and len(inputs) > 0:
-            if expression.op.skips_nulls:
+            if not expression.op.nulls_count_for_min_values:
                 # Most operations do not count NULL values towards min_periods
                 per_col_does_count = (column.notnull() for column in inputs)
                 # All inputs must be non-null for observation to count
@@ -679,13 +678,15 @@ def _join_condition(
 
 
 def _as_groupable(value: ibis_types.Value):
+    from bigframes.core.compile.ibis_compiler import scalar_op_registry
+
     # Some types need to be converted to another type to enable groupby
     if value.type().is_float64():
         return value.cast(ibis_dtypes.str)
     elif value.type().is_geospatial():
         return typing.cast(ibis_types.GeoSpatialColumn, value).as_binary()
     elif value.type().is_json():
-        return scalar_op_compiler.to_json_string(value)
+        return scalar_op_registry.to_json_string(value)
     else:
         return value
 

tests/system/small/pandas/io/__init__.py …s/core/compile/ibis_compiler/__init__.pytests/system/small/pandas/io/__init__.py renamed to bigframes/core/compile/ibis_compiler/__init__.py tests/system/small/pandas/io/__init__.py renamed to bigframes/core/compile/ibis_compiler/__init__.py

@@ -11,3 +11,14 @@
 # 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.
+
+"""Compiler for BigFrames expression to Ibis expression.
+
+Make sure to import all ibis_compiler implementations here so that they get
+registered.
+"""
+
+from __future__ import annotations
+
+import bigframes.core.compile.ibis_compiler.operations.generic_ops  # noqa: F401
+import bigframes.core.compile.ibis_compiler.scalar_op_registry  # noqa: F401

…rames/core/compile/aggregate_compiler.py …pile/ibis_compiler/aggregate_compiler.pybigframes/core/compile/aggregate_compiler.py renamed to bigframes/core/compile/ibis_compiler/aggregate_compiler.py bigframes/core/compile/aggregate_compiler.py renamed to bigframes/core/compile/ibis_compiler/aggregate_compiler.py

@@ -27,8 +27,8 @@
 import pandas as pd
 
 from bigframes.core.compile import constants as compiler_constants
+import bigframes.core.compile.ibis_compiler.scalar_op_compiler as scalar_compilers
 import bigframes.core.compile.ibis_types as compile_ibis_types
-import bigframes.core.compile.scalar_op_compiler as scalar_compilers
 import bigframes.core.expression as ex
 import bigframes.core.window_spec as window_spec
 import bigframes.operations.aggregations as agg_ops

bigframes/core/compile/compiler.py …e/compile/ibis_compiler/ibis_compiler.pybigframes/core/compile/compiler.py renamed to bigframes/core/compile/ibis_compiler/ibis_compiler.py bigframes/core/compile/compiler.py renamed to bigframes/core/compile/ibis_compiler/ibis_compiler.py

@@ -29,7 +29,6 @@
 import bigframes.core.compile.concat as concat_impl
 import bigframes.core.compile.configs as configs
 import bigframes.core.compile.explode
-import bigframes.core.compile.scalar_op_compiler as compile_scalar
 import bigframes.core.nodes as nodes
 import bigframes.core.ordering as bf_ordering
 import bigframes.core.rewrite as rewrites
@@ -178,6 +177,8 @@ def compile_readlocal(node: nodes.ReadLocalNode, *args):
 
 @_compile_node.register
 def compile_readtable(node: nodes.ReadTableNode, *args):
+    from bigframes.core.compile.ibis_compiler import scalar_op_registry
+
     ibis_table = _table_to_ibis(
         node.source, scan_cols=[col.source_id for col in node.scan_list.items]
     )
@@ -188,7 +189,7 @@ def compile_readtable(node: nodes.ReadTableNode, *args):
             scan_item.dtype == dtypes.JSON_DTYPE
             and ibis_table[scan_item.source_id].type() == ibis_dtypes.string
         ):
-            json_column = compile_scalar.parse_json(
+            json_column = scalar_op_registry.parse_json(
                 ibis_table[scan_item.source_id]
             ).name(scan_item.source_id)
             ibis_table = ibis_table.mutate(json_column)

…s/system/small/pandas/io/api/__init__.py …ile/ibis_compiler/operations/__init__.pytests/system/small/pandas/io/api/__init__.py renamed to bigframes/core/compile/ibis_compiler/operations/__init__.py tests/system/small/pandas/io/api/__init__.py renamed to bigframes/core/compile/ibis_compiler/operations/__init__.py

@@ -11,3 +11,11 @@
 # 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.
+
+"""Operation implementations for the Ibis-based compiler.
+
+This directory structure should reflect the same layout as the
+`bigframes/operations` directory where the operations are defined.
+
+Prefer a few ops per file to keep file sizes manageable for text editors and LLMs.
+"""

bigframes/core/compile/ibis_compiler/operations/generic_ops.py

@@ -0,0 +1,38 @@
+# 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.
+
+"""
+BigFrames -> Ibis compilation for the operations in bigframes.operations.generic_ops.
+
+Please keep implementations in sequential order by op name.
+"""
+
+from __future__ import annotations
+
+from bigframes_vendored.ibis.expr import types as ibis_types
+
+from bigframes.core.compile.ibis_compiler import scalar_op_compiler
+from bigframes.operations import generic_ops
+
+register_unary_op = scalar_op_compiler.scalar_op_compiler.register_unary_op
+
+
+@register_unary_op(generic_ops.notnull_op)
+def notnull_op_impl(x: ibis_types.Value):
+    return x.notnull()
+
+
+@register_unary_op(generic_ops.isnull_op)
+def isnull_op_impl(x: ibis_types.Value):
+    return x.isnull()