diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_array_ops/test_to_array_with_subquery_expression/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_array_ops/test_to_array_with_subquery_expression/out.sql new file mode 100644 index 000000000000..63dfcec026b2 --- /dev/null +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_array_ops/test_to_array_with_subquery_expression/out.sql @@ -0,0 +1,12 @@ +SELECT + [ + COALESCE( + ( + SELECT + COALESCE(SUM(bf_arr_reduce_uid), 0) + FROM UNNEST(`float_list_col`) AS bf_arr_reduce_uid + ), + 0.0 + ) + ] AS `arr_subquery_coalesce` +FROM `bigframes-dev`.`sqlglot_test`.`repeated_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_array_ops.py b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_array_ops.py index e0bb1fd7377e..1b358b3a3b11 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_array_ops.py +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_array_ops.py @@ -104,3 +104,14 @@ def test_to_array_op(scalar_types_df: bpd.DataFrame, snapshot): sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys())) snapshot.assert_match(sql, "out.sql") + + +def test_to_array_with_subquery_expression(repeated_types_df: bpd.DataFrame, snapshot): + reduced = ops.ArrayReduceOp(agg_ops.SumOp()).as_expr("float_list_col") + coalesced_reduced = ops.coalesce_op.as_expr(reduced, expression.const(0.0)) + array_expr = ops.ToArrayOp().as_expr(coalesced_reduced) + + sql = utils._apply_ops_to_sql( + repeated_types_df, [array_expr], ["arr_subquery_coalesce"] + ) + snapshot.assert_match(sql, "out.sql") diff --git a/packages/bigframes/third_party/bigframes_vendored/sqlglot/dialects/dialect.py b/packages/bigframes/third_party/bigframes_vendored/sqlglot/dialects/dialect.py index 64d26775c86f..8e26b777abdc 100644 --- a/packages/bigframes/third_party/bigframes_vendored/sqlglot/dialects/dialect.py +++ b/packages/bigframes/third_party/bigframes_vendored/sqlglot/dialects/dialect.py @@ -1236,7 +1236,14 @@ def inline_array_sql(self: Generator, expression: exp.Expression) -> str: def inline_array_unless_query(self: Generator, expression: exp.Expression) -> str: elem = seq_get(expression.expressions, 0) - if isinstance(elem, exp.Expression) and elem.find(exp.Query): + if ( + len(expression.expressions) == 1 + and isinstance(elem, exp.Expression) + and ( + isinstance(elem, exp.Query) + or (isinstance(elem, exp.Subquery) and isinstance(elem.this, exp.Query)) + ) + ): return self.func("ARRAY", elem) return inline_array_sql(self, expression)