Skip to content

Commit 37c75ec

Browse files
chelsea-linchelsea-lin
committed
refactor: support ops.mod_op for the sqlglot compiler
1 parent 9130a61 commit 37c75ec

File tree

4 files changed

+308
-1
lines changed

4 files changed

+308
-1
lines changed

‎bigframes/core/compile/sqlglot/expressions/numeric_ops.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -323,6 +323,49 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
323323
return result
324324

325325

326+
@register_binary_op(ops.mod_op)
327+
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
328+
# In BigQuery returned value has the same sign as X. In pandas, the sign of y is used, so we need to flip the result if sign(x) != sign(y)
329+
left_expr = _coerce_bool_to_int(left)
330+
right_expr = _coerce_bool_to_int(right)
331+
332+
# BigQuery MOD function doesn't support float types, so cast to BIGNUMERIC
333+
if left.dtype == dtypes.FLOAT_DTYPE or right.dtype == dtypes.FLOAT_DTYPE:
334+
left_expr = sge.Cast(this=left_expr, to="BIGNUMERIC")
335+
right_expr = sge.Cast(this=right_expr, to="BIGNUMERIC")
336+
337+
# MOD(N, 0) will error in bigquery, but needs to return null
338+
bq_mod = sge.Mod(this=left_expr, expression=right_expr)
339+
zero_result = (
340+
constants._NAN
341+
if (left.dtype == dtypes.FLOAT_DTYPE or right.dtype == dtypes.FLOAT_DTYPE)
342+
else constants._ZERO
343+
)
344+
return sge.Case(
345+
ifs=[
346+
sge.If(
347+
this=sge.EQ(this=right_expr, expression=constants._ZERO),
348+
true=zero_result * left_expr,
349+
),
350+
sge.If(
351+
this=sge.and_(
352+
right_expr < constants._ZERO,
353+
bq_mod > constants._ZERO,
354+
),
355+
true=right_expr + bq_mod,
356+
),
357+
sge.If(
358+
this=sge.and_(
359+
right_expr > constants._ZERO,
360+
bq_mod < constants._ZERO,
361+
),
362+
true=right_expr + bq_mod,
363+
),
364+
],
365+
default=bq_mod,
366+
)
367+
368+
326369
@register_binary_op(ops.mul_op)
327370
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
328371
left_expr = _coerce_bool_to_int(left)

‎tests/system/small/engines/test_numeric_ops.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ def test_engines_project_floordiv_durations(
161161
assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
162162

163163

164-
@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
164+
@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
165165
def test_engines_project_mod(
166166
scalars_array_value: array_value.ArrayValue,
167167
engine,

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mod_numeric/out.sql‎

Lines changed: 248 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`int64_col` AS `bfcol_0`,
4+
`float64_col` AS `bfcol_1`,
5+
`rowindex` AS `bfcol_2`
6+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
7+
), `bfcte_1` AS (
8+
SELECT
9+
*,
10+
`bfcol_2` AS `bfcol_6`,
11+
`bfcol_0` AS `bfcol_7`,
12+
`bfcol_1` AS `bfcol_8`,
13+
CASE
14+
WHEN `bfcol_0` = CAST(0 AS INT64)
15+
THEN NULL
16+
WHEN `bfcol_0` < CAST(0 AS INT64)
17+
AND (
18+
MOD(`bfcol_0`, `bfcol_0`)
19+
) > CAST(0 AS INT64)
20+
THEN `bfcol_0` + (
21+
MOD(`bfcol_0`, `bfcol_0`)
22+
)
23+
WHEN `bfcol_0` > CAST(0 AS INT64)
24+
AND (
25+
MOD(`bfcol_0`, `bfcol_0`)
26+
) < CAST(0 AS INT64)
27+
THEN `bfcol_0` + (
28+
MOD(`bfcol_0`, `bfcol_0`)
29+
)
30+
ELSE MOD(`bfcol_0`, `bfcol_0`)
31+
END AS `bfcol_9`
32+
FROM `bfcte_0`
33+
), `bfcte_2` AS (
34+
SELECT
35+
*,
36+
`bfcol_6` AS `bfcol_14`,
37+
`bfcol_7` AS `bfcol_15`,
38+
`bfcol_8` AS `bfcol_16`,
39+
`bfcol_9` AS `bfcol_17`,
40+
CASE
41+
WHEN -`bfcol_7` = CAST(0 AS INT64)
42+
THEN NULL
43+
WHEN -`bfcol_7` < CAST(0 AS INT64)
44+
AND (
45+
MOD(`bfcol_7`, -`bfcol_7`)
46+
) > CAST(0 AS INT64)
47+
THEN -`bfcol_7` + (
48+
MOD(`bfcol_7`, -`bfcol_7`)
49+
)
50+
WHEN -`bfcol_7` > CAST(0 AS INT64)
51+
AND (
52+
MOD(`bfcol_7`, -`bfcol_7`)
53+
) < CAST(0 AS INT64)
54+
THEN -`bfcol_7` + (
55+
MOD(`bfcol_7`, -`bfcol_7`)
56+
)
57+
ELSE MOD(`bfcol_7`, -`bfcol_7`)
58+
END AS `bfcol_18`
59+
FROM `bfcte_1`
60+
), `bfcte_3` AS (
61+
SELECT
62+
*,
63+
`bfcol_14` AS `bfcol_24`,
64+
`bfcol_15` AS `bfcol_25`,
65+
`bfcol_16` AS `bfcol_26`,
66+
`bfcol_17` AS `bfcol_27`,
67+
`bfcol_18` AS `bfcol_28`,
68+
CASE
69+
WHEN 1 = CAST(0 AS INT64)
70+
THEN NULL
71+
WHEN 1 < CAST(0 AS INT64) AND (
72+
MOD(`bfcol_15`, 1)
73+
) > CAST(0 AS INT64)
74+
THEN 1 + (
75+
MOD(`bfcol_15`, 1)
76+
)
77+
WHEN 1 > CAST(0 AS INT64) AND (
78+
MOD(`bfcol_15`, 1)
79+
) < CAST(0 AS INT64)
80+
THEN 1 + (
81+
MOD(`bfcol_15`, 1)
82+
)
83+
ELSE MOD(`bfcol_15`, 1)
84+
END AS `bfcol_29`
85+
FROM `bfcte_2`
86+
), `bfcte_4` AS (
87+
SELECT
88+
*,
89+
`bfcol_24` AS `bfcol_36`,
90+
`bfcol_25` AS `bfcol_37`,
91+
`bfcol_26` AS `bfcol_38`,
92+
`bfcol_27` AS `bfcol_39`,
93+
`bfcol_28` AS `bfcol_40`,
94+
`bfcol_29` AS `bfcol_41`,
95+
CASE
96+
WHEN 0 = CAST(0 AS INT64)
97+
THEN NULL
98+
WHEN 0 < CAST(0 AS INT64) AND (
99+
MOD(`bfcol_25`, 0)
100+
) > CAST(0 AS INT64)
101+
THEN 0 + (
102+
MOD(`bfcol_25`, 0)
103+
)
104+
WHEN 0 > CAST(0 AS INT64) AND (
105+
MOD(`bfcol_25`, 0)
106+
) < CAST(0 AS INT64)
107+
THEN 0 + (
108+
MOD(`bfcol_25`, 0)
109+
)
110+
ELSE MOD(`bfcol_25`, 0)
111+
END AS `bfcol_42`
112+
FROM `bfcte_3`
113+
), `bfcte_5` AS (
114+
SELECT
115+
*,
116+
`bfcol_36` AS `bfcol_50`,
117+
`bfcol_37` AS `bfcol_51`,
118+
`bfcol_38` AS `bfcol_52`,
119+
`bfcol_39` AS `bfcol_53`,
120+
`bfcol_40` AS `bfcol_54`,
121+
`bfcol_41` AS `bfcol_55`,
122+
`bfcol_42` AS `bfcol_56`,
123+
CASE
124+
WHEN `bfcol_38` = CAST(0 AS INT64)
125+
THEN NULL
126+
WHEN `bfcol_38` < CAST(0 AS INT64)
127+
AND (
128+
MOD(`bfcol_38`, `bfcol_38`)
129+
) > CAST(0 AS INT64)
130+
THEN `bfcol_38` + (
131+
MOD(`bfcol_38`, `bfcol_38`)
132+
)
133+
WHEN `bfcol_38` > CAST(0 AS INT64)
134+
AND (
135+
MOD(`bfcol_38`, `bfcol_38`)
136+
) < CAST(0 AS INT64)
137+
THEN `bfcol_38` + (
138+
MOD(`bfcol_38`, `bfcol_38`)
139+
)
140+
ELSE MOD(`bfcol_38`, `bfcol_38`)
141+
END AS `bfcol_57`
142+
FROM `bfcte_4`
143+
), `bfcte_6` AS (
144+
SELECT
145+
*,
146+
`bfcol_50` AS `bfcol_66`,
147+
`bfcol_51` AS `bfcol_67`,
148+
`bfcol_52` AS `bfcol_68`,
149+
`bfcol_53` AS `bfcol_69`,
150+
`bfcol_54` AS `bfcol_70`,
151+
`bfcol_55` AS `bfcol_71`,
152+
`bfcol_56` AS `bfcol_72`,
153+
`bfcol_57` AS `bfcol_73`,
154+
CASE
155+
WHEN -`bfcol_52` = CAST(0 AS INT64)
156+
THEN NULL
157+
WHEN -`bfcol_52` < CAST(0 AS INT64)
158+
AND (
159+
MOD(`bfcol_52`, -`bfcol_52`)
160+
) > CAST(0 AS INT64)
161+
THEN -`bfcol_52` + (
162+
MOD(`bfcol_52`, -`bfcol_52`)
163+
)
164+
WHEN -`bfcol_52` > CAST(0 AS INT64)
165+
AND (
166+
MOD(`bfcol_52`, -`bfcol_52`)
167+
) < CAST(0 AS INT64)
168+
THEN -`bfcol_52` + (
169+
MOD(`bfcol_52`, -`bfcol_52`)
170+
)
171+
ELSE MOD(`bfcol_52`, -`bfcol_52`)
172+
END AS `bfcol_74`
173+
FROM `bfcte_5`
174+
), `bfcte_7` AS (
175+
SELECT
176+
*,
177+
`bfcol_66` AS `bfcol_84`,
178+
`bfcol_67` AS `bfcol_85`,
179+
`bfcol_68` AS `bfcol_86`,
180+
`bfcol_69` AS `bfcol_87`,
181+
`bfcol_70` AS `bfcol_88`,
182+
`bfcol_71` AS `bfcol_89`,
183+
`bfcol_72` AS `bfcol_90`,
184+
`bfcol_73` AS `bfcol_91`,
185+
`bfcol_74` AS `bfcol_92`,
186+
CASE
187+
WHEN 1 = CAST(0 AS INT64)
188+
THEN NULL
189+
WHEN 1 < CAST(0 AS INT64) AND (
190+
MOD(`bfcol_68`, 1)
191+
) > CAST(0 AS INT64)
192+
THEN 1 + (
193+
MOD(`bfcol_68`, 1)
194+
)
195+
WHEN 1 > CAST(0 AS INT64) AND (
196+
MOD(`bfcol_68`, 1)
197+
) < CAST(0 AS INT64)
198+
THEN 1 + (
199+
MOD(`bfcol_68`, 1)
200+
)
201+
ELSE MOD(`bfcol_68`, 1)
202+
END AS `bfcol_93`
203+
FROM `bfcte_6`
204+
), `bfcte_8` AS (
205+
SELECT
206+
*,
207+
`bfcol_84` AS `bfcol_104`,
208+
`bfcol_85` AS `bfcol_105`,
209+
`bfcol_86` AS `bfcol_106`,
210+
`bfcol_87` AS `bfcol_107`,
211+
`bfcol_88` AS `bfcol_108`,
212+
`bfcol_89` AS `bfcol_109`,
213+
`bfcol_90` AS `bfcol_110`,
214+
`bfcol_91` AS `bfcol_111`,
215+
`bfcol_92` AS `bfcol_112`,
216+
`bfcol_93` AS `bfcol_113`,
217+
CASE
218+
WHEN 0 = CAST(0 AS INT64)
219+
THEN NULL
220+
WHEN 0 < CAST(0 AS INT64) AND (
221+
MOD(`bfcol_86`, 0)
222+
) > CAST(0 AS INT64)
223+
THEN 0 + (
224+
MOD(`bfcol_86`, 0)
225+
)
226+
WHEN 0 > CAST(0 AS INT64) AND (
227+
MOD(`bfcol_86`, 0)
228+
) < CAST(0 AS INT64)
229+
THEN 0 + (
230+
MOD(`bfcol_86`, 0)
231+
)
232+
ELSE MOD(`bfcol_86`, 0)
233+
END AS `bfcol_114`
234+
FROM `bfcte_7`
235+
)
236+
SELECT
237+
`bfcol_104` AS `rowindex`,
238+
`bfcol_105` AS `int64_col`,
239+
`bfcol_106` AS `float64_col`,
240+
`bfcol_107` AS `int_mod_int`,
241+
`bfcol_108` AS `int_mod_int_neg`,
242+
`bfcol_109` AS `int_mod_1`,
243+
`bfcol_110` AS `int_mod_0`,
244+
`bfcol_111` AS `float_mod_float`,
245+
`bfcol_112` AS `float_mod_float_neg`,
246+
`bfcol_113` AS `float_mod_1`,
247+
`bfcol_114` AS `float_mod_0`
248+
FROM `bfcte_8`

‎tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,22 @@ def test_mul_numeric(scalar_types_df: bpd.DataFrame, snapshot):
287287
snapshot.assert_match(bf_df.sql, "out.sql")
288288

289289

290+
def test_mod_numeric(scalar_types_df: bpd.DataFrame, snapshot):
291+
bf_df = scalar_types_df[["int64_col", "float64_col"]]
292+
293+
bf_df["int_mod_int"] = bf_df["int64_col"] % bf_df["int64_col"]
294+
bf_df["int_mod_int_neg"] = bf_df["int64_col"] % -bf_df["int64_col"]
295+
bf_df["int_mod_1"] = bf_df["int64_col"] % 1
296+
bf_df["int_mod_0"] = bf_df["int64_col"] % 0
297+
298+
bf_df["float_mod_float"] = bf_df["float64_col"] % bf_df["float64_col"]
299+
bf_df["float_mod_float_neg"] = bf_df["float64_col"] % -bf_df["float64_col"]
300+
bf_df["float_mod_1"] = bf_df["float64_col"] % 1
301+
bf_df["float_mod_0"] = bf_df["float64_col"] % 0
302+
303+
snapshot.assert_match(bf_df.sql, "out.sql")
304+
305+
290306
def test_sub_numeric(scalar_types_df: bpd.DataFrame, snapshot):
291307
bf_df = scalar_types_df[["int64_col", "bool_col"]]
292308

0 commit comments

Comments
 (0)