Skip to content

Commit 8622cae

Browse files
authored
Merge branch 'main' into df-where-callable
2 parents ae42e8c + 41dda88 commit 8622cae

36 files changed

+1617
-343
lines changed

bigframes/_importing.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import importlib
1515
from types import ModuleType
1616

17+
import numpy
1718
from packaging import version
1819

1920
# Keep this in sync with setup.py
@@ -22,9 +23,13 @@
2223

2324
def import_polars() -> ModuleType:
2425
polars_module = importlib.import_module("polars")
25-
imported_version = version.Version(polars_module.build_info()["version"])
26-
if imported_version < POLARS_MIN_VERSION:
26+
# Check for necessary methods instead of the version number because we
27+
# can't trust the polars version until
28+
# https://github.com/pola-rs/polars/issues/23940 is fixed.
29+
try:
30+
polars_module.lit(numpy.int64(100), dtype=polars_module.Int64())
31+
except TypeError:
2732
raise ImportError(
28-
f"Imported polars version: {imported_version} is below the minimum version: {POLARS_MIN_VERSION}"
33+
f"Imported polars version is likely below the minimum version: {POLARS_MIN_VERSION}"
2934
)
3035
return polars_module

bigframes/core/compile/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@
1414
from __future__ import annotations
1515

1616
from bigframes.core.compile.api import test_only_ibis_inferred_schema
17-
from bigframes.core.compile.compiler import compile_sql
1817
from bigframes.core.compile.configs import CompileRequest, CompileResult
18+
from bigframes.core.compile.ibis_compiler.ibis_compiler import compile_sql
1919

2020
__all__ = [
2121
"test_only_ibis_inferred_schema",

bigframes/core/compile/api.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from typing import TYPE_CHECKING
1717

1818
from bigframes.core import rewrite
19-
from bigframes.core.compile import compiler
19+
from bigframes.core.compile.ibis_compiler import ibis_compiler
2020

2121
if TYPE_CHECKING:
2222
import bigframes.core.nodes
@@ -26,9 +26,9 @@ def test_only_ibis_inferred_schema(node: bigframes.core.nodes.BigFrameNode):
2626
"""Use only for testing paths to ensure ibis inferred schema does not diverge from bigframes inferred schema."""
2727
import bigframes.core.schema
2828

29-
node = compiler._replace_unsupported_ops(node)
29+
node = ibis_compiler._replace_unsupported_ops(node)
3030
node = rewrite.bake_order(node)
31-
ir = compiler.compile_node(node)
31+
ir = ibis_compiler.compile_node(node)
3232
items = tuple(
3333
bigframes.core.schema.SchemaItem(name, ir.get_column_type(ibis_id))
3434
for name, ibis_id in zip(node.schema.names, ir.column_ids)

bigframes/core/compile/compiled.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,10 @@
3030
import pyarrow as pa
3131

3232
from bigframes.core import utils
33-
import bigframes.core.compile.aggregate_compiler as agg_compiler
3433
import bigframes.core.compile.googlesql
34+
import bigframes.core.compile.ibis_compiler.aggregate_compiler as agg_compiler
35+
import bigframes.core.compile.ibis_compiler.scalar_op_compiler as op_compilers
3536
import bigframes.core.compile.ibis_types
36-
import bigframes.core.compile.scalar_op_compiler as op_compilers
37-
import bigframes.core.compile.scalar_op_compiler as scalar_op_compiler
3837
import bigframes.core.expression as ex
3938
from bigframes.core.ordering import OrderingExpression
4039
import bigframes.core.sql
@@ -460,7 +459,7 @@ def project_window_op(
460459
for column in inputs:
461460
clauses.append((column.isnull(), ibis_types.null()))
462461
if window_spec.min_periods and len(inputs) > 0:
463-
if expression.op.skips_nulls:
462+
if not expression.op.nulls_count_for_min_values:
464463
# Most operations do not count NULL values towards min_periods
465464
per_col_does_count = (column.notnull() for column in inputs)
466465
# All inputs must be non-null for observation to count
@@ -679,13 +678,15 @@ def _join_condition(
679678

680679

681680
def _as_groupable(value: ibis_types.Value):
681+
from bigframes.core.compile.ibis_compiler import scalar_op_registry
682+
682683
# Some types need to be converted to another type to enable groupby
683684
if value.type().is_float64():
684685
return value.cast(ibis_dtypes.str)
685686
elif value.type().is_geospatial():
686687
return typing.cast(ibis_types.GeoSpatialColumn, value).as_binary()
687688
elif value.type().is_json():
688-
return scalar_op_compiler.to_json_string(value)
689+
return scalar_op_registry.to_json_string(value)
689690
else:
690691
return value
691692

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,14 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14+
15+
"""Compiler for BigFrames expression to Ibis expression.
16+
17+
Make sure to import all ibis_compiler implementations here so that they get
18+
registered.
19+
"""
20+
21+
from __future__ import annotations
22+
23+
import bigframes.core.compile.ibis_compiler.operations.generic_ops # noqa: F401
24+
import bigframes.core.compile.ibis_compiler.scalar_op_registry # noqa: F401

bigframes/core/compile/aggregate_compiler.py renamed to bigframes/core/compile/ibis_compiler/aggregate_compiler.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@
2727
import pandas as pd
2828

2929
from bigframes.core.compile import constants as compiler_constants
30+
import bigframes.core.compile.ibis_compiler.scalar_op_compiler as scalar_compilers
3031
import bigframes.core.compile.ibis_types as compile_ibis_types
31-
import bigframes.core.compile.scalar_op_compiler as scalar_compilers
3232
import bigframes.core.expression as ex
3333
import bigframes.core.window_spec as window_spec
3434
import bigframes.operations.aggregations as agg_ops

bigframes/core/compile/compiler.py renamed to bigframes/core/compile/ibis_compiler/ibis_compiler.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
import bigframes.core.compile.concat as concat_impl
3030
import bigframes.core.compile.configs as configs
3131
import bigframes.core.compile.explode
32-
import bigframes.core.compile.scalar_op_compiler as compile_scalar
3332
import bigframes.core.nodes as nodes
3433
import bigframes.core.ordering as bf_ordering
3534
import bigframes.core.rewrite as rewrites
@@ -178,6 +177,8 @@ def compile_readlocal(node: nodes.ReadLocalNode, *args):
178177

179178
@_compile_node.register
180179
def compile_readtable(node: nodes.ReadTableNode, *args):
180+
from bigframes.core.compile.ibis_compiler import scalar_op_registry
181+
181182
ibis_table = _table_to_ibis(
182183
node.source, scan_cols=[col.source_id for col in node.scan_list.items]
183184
)
@@ -188,7 +189,7 @@ def compile_readtable(node: nodes.ReadTableNode, *args):
188189
scan_item.dtype == dtypes.JSON_DTYPE
189190
and ibis_table[scan_item.source_id].type() == ibis_dtypes.string
190191
):
191-
json_column = compile_scalar.parse_json(
192+
json_column = scalar_op_registry.parse_json(
192193
ibis_table[scan_item.source_id]
193194
).name(scan_item.source_id)
194195
ibis_table = ibis_table.mutate(json_column)

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,11 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14+
15+
"""Operation implementations for the Ibis-based compiler.
16+
17+
This directory structure should reflect the same layout as the
18+
`bigframes/operations` directory where the operations are defined.
19+
20+
Prefer a few ops per file to keep file sizes manageable for text editors and LLMs.
21+
"""
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""
16+
BigFrames -> Ibis compilation for the operations in bigframes.operations.generic_ops.
17+
18+
Please keep implementations in sequential order by op name.
19+
"""
20+
21+
from __future__ import annotations
22+
23+
from bigframes_vendored.ibis.expr import types as ibis_types
24+
25+
from bigframes.core.compile.ibis_compiler import scalar_op_compiler
26+
from bigframes.operations import generic_ops
27+
28+
register_unary_op = scalar_op_compiler.scalar_op_compiler.register_unary_op
29+
30+
31+
@register_unary_op(generic_ops.notnull_op)
32+
def notnull_op_impl(x: ibis_types.Value):
33+
return x.notnull()
34+
35+
36+
@register_unary_op(generic_ops.isnull_op)
37+
def isnull_op_impl(x: ibis_types.Value):
38+
return x.isnull()

0 commit comments

Comments
 (0)