Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit 98147d0

Browse files
TrevorBergeronTrevorBergeron
committed
Merge remote-tracking branch 'github/main' into udf_return_change
2 parents 3f180c3 + b589de9 commit 98147d0

File tree

16 files changed

+799
-26
lines changed

16 files changed

+799
-26
lines changed

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

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
from bigframes import dtypes
2020
from bigframes import operations as ops
21+
from bigframes.core.compile.sqlglot import sql
2122
import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
2223
from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
2324

@@ -29,10 +30,10 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
2930
# For AND, when we encounter a NULL value, we only know when the result is FALSE,
3031
# otherwise the result is unknown (NULL). See: truth table at
3132
# https://en.wikibooks.org/wiki/Structured_Query_Language/NULLs_and_the_Three_Valued_Logic#AND,_OR
32-
if left.expr == sge.null():
33+
if sql.is_null_literal(left.expr):
3334
condition = sge.EQ(this=right.expr, expression=sge.convert(False))
3435
return sge.If(this=condition, true=right.expr, false=sge.null())
35-
if right.expr == sge.null():
36+
if sql.is_null_literal(right.expr):
3637
condition = sge.EQ(this=left.expr, expression=sge.convert(False))
3738
return sge.If(this=condition, true=left.expr, false=sge.null())
3839

@@ -46,10 +47,10 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
4647
# For OR, when we encounter a NULL value, we only know when the result is TRUE,
4748
# otherwise the result is unknown (NULL). See: truth table at
4849
# https://en.wikibooks.org/wiki/Structured_Query_Language/NULLs_and_the_Three_Valued_Logic#AND,_OR
49-
if left.expr == sge.null():
50+
if sql.is_null_literal(left.expr):
5051
condition = sge.EQ(this=right.expr, expression=sge.convert(True))
5152
return sge.If(this=condition, true=right.expr, false=sge.null())
52-
if right.expr == sge.null():
53+
if sql.is_null_literal(right.expr):
5354
condition = sge.EQ(this=left.expr, expression=sge.convert(True))
5455
return sge.If(this=condition, true=left.expr, false=sge.null())
5556

@@ -64,12 +65,12 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
6465
# maintains the boolean data type.
6566
left_expr = left.expr
6667
left_dtype = left.dtype
67-
if left_expr == sge.null():
68+
if sql.is_null_literal(left_expr):
6869
left_expr = sge.Cast(this=sge.convert(None), to="BOOLEAN")
6970
left_dtype = dtypes.BOOL_DTYPE
7071
right_expr = right.expr
7172
right_dtype = right.dtype
72-
if right_expr == sge.null():
73+
if sql.is_null_literal(right_expr):
7374
right_expr = sge.Cast(this=sge.convert(None), to="BOOLEAN")
7475
right_dtype = dtypes.BOOL_DTYPE
7576

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
102102

103103
@register_binary_op(ops.ge_op)
104104
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
105-
if left.expr == sge.null() or right.expr == sge.null():
105+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
106106
return sge.null()
107107

108108
left_expr = _coerce_bool_to_int(left)
@@ -112,7 +112,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
112112

113113
@register_binary_op(ops.gt_op)
114114
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
115-
if left.expr == sge.null() or right.expr == sge.null():
115+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
116116
return sge.null()
117117

118118
left_expr = _coerce_bool_to_int(left)
@@ -122,7 +122,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
122122

123123
@register_binary_op(ops.lt_op)
124124
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
125-
if left.expr == sge.null() or right.expr == sge.null():
125+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
126126
return sge.null()
127127

128128
left_expr = _coerce_bool_to_int(left)
@@ -132,7 +132,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
132132

133133
@register_binary_op(ops.le_op)
134134
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
135-
if left.expr == sge.null() or right.expr == sge.null():
135+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
136136
return sge.null()
137137

138138
left_expr = _coerce_bool_to_int(left)

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

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
from bigframes import dtypes
2121
from bigframes import operations as ops
22+
from bigframes.core.compile.sqlglot import sql
2223
import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
2324
from bigframes.core.compile.sqlglot.expressions.common import round_towards_zero
2425
import bigframes.core.compile.sqlglot.expressions.constants as constants
@@ -260,6 +261,9 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
260261
def _int_pow_op(
261262
left_expr: sge.Expression, right_expr: sge.Expression
262263
) -> sge.Expression:
264+
if sql.is_null_literal(left_expr) or sql.is_null_literal(right_expr):
265+
return sge.null()
266+
263267
overflow_cond = sge.and_(
264268
sge.NEQ(this=left_expr, expression=sge.convert(0)),
265269
sge.GT(
@@ -292,6 +296,9 @@ def _int_pow_op(
292296
def _float_pow_op(
293297
left_expr: sge.Expression, right_expr: sge.Expression
294298
) -> sge.Expression:
299+
if sql.is_null_literal(left_expr) or sql.is_null_literal(right_expr):
300+
return sge.null()
301+
295302
# Most conditions here seek to prevent calling BQ POW with inputs that would generate errors.
296303
# See: https://cloud.google.com/bigquery/docs/reference/standard-sql/mathematical_functions#pow
297304
overflow_cond = sge.and_(
@@ -425,7 +432,7 @@ def _(expr: TypedExpr) -> sge.Expression:
425432

426433
@register_binary_op(ops.add_op)
427434
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
428-
if left.expr == sge.null() or right.expr == sge.null():
435+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
429436
return sge.null()
430437

431438
if left.dtype == dtypes.STRING_DTYPE and right.dtype == dtypes.STRING_DTYPE:
@@ -463,6 +470,9 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
463470

464471
@register_binary_op(ops.div_op)
465472
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
473+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
474+
return sge.null()
475+
466476
left_expr = _coerce_bool_to_int(left)
467477
right_expr = _coerce_bool_to_int(right)
468478

@@ -482,7 +492,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
482492

483493
@register_binary_op(ops.floordiv_op)
484494
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
485-
if left.expr == sge.null() or right.expr == sge.null():
495+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
486496
return sge.null()
487497

488498
left_expr = _coerce_bool_to_int(left)
@@ -525,6 +535,9 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
525535

526536
@register_binary_op(ops.mod_op)
527537
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
538+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
539+
return sge.null()
540+
528541
# 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)
529542
left_expr = _coerce_bool_to_int(left)
530543
right_expr = _coerce_bool_to_int(right)
@@ -568,7 +581,7 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
568581

569582
@register_binary_op(ops.mul_op)
570583
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
571-
if left.expr == sge.null() or right.expr == sge.null():
584+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
572585
return sge.null()
573586

574587
left_expr = _coerce_bool_to_int(left)
@@ -594,7 +607,7 @@ def _(expr: TypedExpr, n_digits: TypedExpr) -> sge.Expression:
594607

595608
@register_binary_op(ops.sub_op)
596609
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
597-
if left.expr == sge.null() or right.expr == sge.null():
610+
if sql.is_null_literal(left.expr) or sql.is_null_literal(right.expr):
598611
return sge.null()
599612

600613
if dtypes.is_numeric(left.dtype) and dtypes.is_numeric(right.dtype):

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -366,10 +366,16 @@ def string_slice(
366366
column_length + sge.convert(start + 1),
367367
]
368368
)
369-
length_expr = sge.convert(op_end) - sge.Greatest(
369+
length_expr = sge.Greatest(
370370
expressions=[
371371
sge.convert(0),
372-
column_length + sge.convert(start),
372+
sge.convert(op_end)
373+
- sge.Greatest(
374+
expressions=[
375+
sge.convert(0),
376+
column_length + sge.convert(start),
377+
]
378+
),
373379
]
374380
)
375381
else:

‎bigframes/operations/ai.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -612,8 +612,11 @@ def sim_join(
612612
>>> df1 = bpd.DataFrame({'animal': ['monkey', 'spider']})
613613
>>> df2 = bpd.DataFrame({'animal': ['scorpion', 'baboon']})
614614
615-
>>> df1.ai.sim_join(df2, left_on='animal', right_on='animal', model=model, top_k=1)
616-
animal animal_1
615+
>>> res = df1.ai.sim_join(df2, left_on='animal', right_on='animal', model=model, top_k=1)
616+
>>> print("---"); print(res) # doctest: +ELLIPSIS
617+
---
618+
...
619+
animal animal_1
617620
0 monkey baboon
618621
1 spider scorpion
619622
<BLANKLINE>

‎bigframes/session/__init__.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -664,9 +664,11 @@ def read_gbq_query(
664664
... WHERE year = 2016
665665
... GROUP BY pitcherFirstName, pitcherLastName
666666
... ''', index_col="rowindex")
667-
>>> df.head(2)
667+
>>> print("START_OF_OUTPUT"); df.head(2) # doctest: +ELLIPSIS,+NORMALIZE_WHITESPACE
668+
START_OF_OUTPUT
669+
...
668670
pitcherFirstName pitcherLastName averagePitchSpeed
669-
rowindex
671+
...
670672
1 Albertin Chapman 96.514113
671673
2 Zachary Britton 94.591039
672674
<BLANKLINE>

‎docs/user_guide/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ User Guide
4444
:maxdepth: 1
4545

4646
AI Functions <../notebooks/generative_ai/ai_functions.ipynb>
47+
AI Functions for Poster Analysis <../notebooks/generative_ai/ai_movie_poster.ipynb>
4748
AI Forecast <../notebooks/generative_ai/bq_dataframes_ai_forecast.ipynb>
4849
LLM Code Generation <../notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb>
4950
LLM KMeans <../notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb>

0 commit comments

Comments
 (0)