Skip to content

Commit 9cd1fde

Browse files
tswasttswast
committed
limit scope to just splitting large files
1 parent 360edb3 commit 9cd1fde

File tree

21 files changed

+270
-60
lines changed

21 files changed

+270
-60
lines changed

‎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/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/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: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,8 @@
3232
from bigframes.core import utils
3333
import bigframes.core.compile.aggregate_compiler as agg_compiler
3434
import bigframes.core.compile.googlesql
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
3737
import bigframes.core.expression as ex
3838
from bigframes.core.ordering import OrderingExpression
3939
import bigframes.core.sql
@@ -47,7 +47,7 @@
4747
# This must be the last import. Currently depending on side-effects.
4848
# TODO(tswast): Refactor all ops to register in the same file as where they are
4949
# defined so we don't need this.
50-
import bigframes.core.compile.scalar_op_registry # noqa: F401,E402
50+
import bigframes.core.compile.ibis_compiler.scalar_op_registry # noqa: F401,E402
5151

5252

5353
# Ibis Implementations
@@ -684,7 +684,7 @@ def _join_condition(
684684

685685

686686
def _as_groupable(value: ibis_types.Value):
687-
from bigframes.core.compile import scalar_op_registry
687+
from bigframes.core.compile.ibis_compiler import scalar_op_registry
688688

689689
# Some types need to be converted to another type to enable groupby
690690
if value.type().is_float64():

‎bigframes/core/compile/ibis_compiler/__init__.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
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+
"""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.isnull_op # noqa: F401
24+
import bigframes.core.compile.ibis_compiler.operations.generic_ops.notnull_op # noqa: F401
25+
import bigframes.core.compile.ibis_compiler.scalar_op_registry # noqa: F401

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ def compile_readlocal(node: nodes.ReadLocalNode, *args):
177177

178178
@_compile_node.register
179179
def compile_readtable(node: nodes.ReadTableNode, *args):
180-
from bigframes.core.compile import scalar_op_registry
180+
from bigframes.core.compile.ibis_compiler import scalar_op_registry
181181

182182
ibis_table = _table_to_ibis(
183183
node.source, scan_cols=[col.source_id for col in node.scan_list.items]

‎bigframes/core/compile/ibis_compiler/operations/__init__.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
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+
"""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 one file per op to keep file sizes manageable for text editors and LLMs.
21+
"""

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
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.

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
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+
from __future__ import annotations
16+
17+
from bigframes_vendored.ibis.expr import types as ibis_types
18+
19+
from bigframes.core.compile.ibis_compiler import scalar_op_compiler
20+
from bigframes.operations.generic_ops import isnull_op
21+
22+
23+
@scalar_op_compiler.scalar_op_compiler.register_unary_op(isnull_op.isnull_op)
24+
def _ibis_isnull_op_impl(x: ibis_types.Value):
25+
return x.isnull()

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
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+
from __future__ import annotations
16+
17+
import dataclasses
18+
from typing import ClassVar
19+
20+
# Imports for Ibis compilation
21+
from bigframes_vendored.ibis.expr import types as ibis_types
22+
23+
# Direct imports from bigframes
24+
from bigframes import dtypes
25+
from bigframes.core.compile.ibis_compiler import scalar_op_compiler
26+
from bigframes.operations import base_ops
27+
28+
29+
@dataclasses.dataclass(frozen=True)
30+
class NotNullOp(base_ops.UnaryOp):
31+
name: ClassVar[str] = "notnull"
32+
33+
def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
34+
return dtypes.BOOL_DTYPE
35+
36+
37+
notnull_op = NotNullOp()
38+
39+
40+
def _ibis_notnull_op_impl(x: ibis_types.Value):
41+
return x.notnull()
42+
43+
44+
scalar_op_compiler.scalar_op_compiler.register_unary_op(notnull_op)(
45+
_ibis_notnull_op_impl
46+
)
47+
48+
49+
__all__ = [
50+
"NotNullOp",
51+
"notnull_op",
52+
]

0 commit comments

Comments
 (0)