This repository was archived by the owner on May 7, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 68
Expand file tree
/
Copy pathbase_ops.py
More file actions
208 lines (166 loc) · 6.07 KB
/
Copy pathbase_ops.py
File metadata and controls
208 lines (166 loc) · 6.07 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
# 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.
from __future__ import annotations
import dataclasses
import typing
from bigframes import dtypes
import bigframes.operations.type as op_typing
if typing.TYPE_CHECKING:
# Avoids circular dependency
import bigframes.core.expression
class RowOp(typing.Protocol):
@property
def name(self) -> str:
...
def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
...
@property
def is_monotonic(self) -> bool:
"""Whether the row operation preserves total ordering. Can be pruned from ordering expressions."""
...
@property
def is_bijective(self) -> bool:
"""Whether the operation has a 1:1 mapping between inputs and outputs"""
...
@property
def deterministic(self) -> bool:
"""Whether the operation is deterministic" (given deterministic inputs)"""
...
@property
def expensive(self) -> bool:
"""Whether the operation is expensive to calculate. Such ops shouldn't be inlined if referenced multiple places."""
...
@dataclasses.dataclass(frozen=True)
class ScalarOp:
@property
def name(self) -> str:
raise NotImplementedError("RowOp abstract base class has no implementation")
def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
raise NotImplementedError("Abstract operation has no output type")
@property
def is_monotonic(self) -> bool:
"""Whether the row operation preserves total ordering. Can be pruned from ordering expressions."""
return False
@property
def is_bijective(self) -> bool:
"""Whether the operation has a 1:1 mapping between inputs and outputs"""
return False
@property
def deterministic(self) -> bool:
"""Whether the operation is deterministic" (given deterministic inputs)"""
return True
@property
def expensive(self) -> bool:
return False
@dataclasses.dataclass(frozen=True)
class NaryOp(ScalarOp):
def as_expr(
self,
*exprs: typing.Union[str, bigframes.core.expression.Expression],
) -> bigframes.core.expression.Expression:
import bigframes.core.expression
# Keep this in sync with output_type and compilers
inputs: list[bigframes.core.expression.Expression] = []
for expr in exprs:
inputs.append(_convert_expr_input(expr))
return bigframes.core.expression.OpExpression(
self,
tuple(inputs),
)
# These classes can be used to create simple ops that don't take local parameters
# All is needed is a unique name, and to register an implementation in ibis_mappings.py
@dataclasses.dataclass(frozen=True)
class UnaryOp(ScalarOp):
@property
def arguments(self) -> int:
return 1
def as_expr(
self, input_id: typing.Union[str, bigframes.core.expression.Expression] = "arg"
) -> bigframes.core.expression.Expression:
import bigframes.core.expression
return bigframes.core.expression.OpExpression(
self, (_convert_expr_input(input_id),)
)
@dataclasses.dataclass(frozen=True)
class BinaryOp(ScalarOp):
@property
def arguments(self) -> int:
return 2
def as_expr(
self,
left_input: typing.Union[str, bigframes.core.expression.Expression] = "arg1",
right_input: typing.Union[str, bigframes.core.expression.Expression] = "arg2",
) -> bigframes.core.expression.Expression:
import bigframes.core.expression
return bigframes.core.expression.OpExpression(
self,
(
_convert_expr_input(left_input),
_convert_expr_input(right_input),
),
)
@dataclasses.dataclass(frozen=True)
class TernaryOp(ScalarOp):
@property
def arguments(self) -> int:
return 3
def as_expr(
self,
input1: typing.Union[str, bigframes.core.expression.Expression] = "arg1",
input2: typing.Union[str, bigframes.core.expression.Expression] = "arg2",
input3: typing.Union[str, bigframes.core.expression.Expression] = "arg3",
) -> bigframes.core.expression.Expression:
import bigframes.core.expression
return bigframes.core.expression.OpExpression(
self,
(
_convert_expr_input(input1),
_convert_expr_input(input2),
_convert_expr_input(input3),
),
)
def _convert_expr_input(
input: typing.Union[str, bigframes.core.expression.Expression]
) -> bigframes.core.expression.Expression:
"""Allows creating column references with just a string"""
import bigframes.core.expression
if isinstance(input, str):
return bigframes.core.expression.deref(input)
else:
return input
# Operation Factories
def create_unary_op(
name: str, type_signature: op_typing.UnaryTypeSignature
) -> type[UnaryOp]:
return dataclasses.make_dataclass(
name,
[
("name", typing.ClassVar[str], name),
("output_type", typing.ClassVar[typing.Callable], type_signature.as_method),
],
bases=(UnaryOp,),
frozen=True,
)
def create_binary_op(
name: str, type_signature: op_typing.BinaryTypeSignature
) -> type[BinaryOp]:
return dataclasses.make_dataclass(
name,
[
("name", typing.ClassVar[str], name),
("output_type", typing.ClassVar[typing.Callable], type_signature.as_method),
],
bases=(BinaryOp,),
frozen=True,
)