Merge branch 'main' into shuowei-fix-compiler-syntax-guards

Author: shuoweil

bigframes/bigquery/_operations/ai.py

@@ -745,7 +745,7 @@ def if_(
             or pandas Series.
         connection_id (str, optional):
             Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
-            If not provided, the connection from the current session will be used.
+            If not provided, the query uses your end-user credential.
 
     Returns:
         bigframes.series.Series: A new series of bools.
@@ -756,7 +756,7 @@ def if_(
 
     operator = ai_ops.AIIf(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
     )
 
     return series_list[0]._apply_nary_op(operator, series_list[1:])
@@ -800,7 +800,7 @@ def classify(
             Categories to classify the input into.
         connection_id (str, optional):
             Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
-            If not provided, the connection from the current session will be used.
+            If not provided, the query uses your end-user credential.
 
     Returns:
         bigframes.series.Series: A new series of strings.
@@ -812,7 +812,7 @@ def classify(
     operator = ai_ops.AIClassify(
         prompt_context=tuple(prompt_context),
         categories=tuple(categories),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
     )
 
     return series_list[0]._apply_nary_op(operator, series_list[1:])
@@ -853,7 +853,7 @@ def score(
             or pandas Series.
         connection_id (str, optional):
             Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
-            If not provided, the connection from the current session will be used.
+            If not provided, the query uses your end-user credential.
 
     Returns:
         bigframes.series.Series: A new series of double (float) values.
@@ -864,7 +864,7 @@ def score(
 
     operator = ai_ops.AIScore(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
     )
 
     return series_list[0]._apply_nary_op(operator, series_list[1:])

bigframes/core/bigframe_node.py

@@ -330,12 +330,32 @@ def top_down(
         """
         Perform a top-down transformation of the BigFrameNode tree.
         """
+        results: Dict[BigFrameNode, BigFrameNode] = {}
+        # Each stack entry is (node, t_node). t_node is None until transform(node) is called.
+        stack: list[tuple[BigFrameNode, typing.Optional[BigFrameNode]]] = [(self, None)]
 
-        @functools.cache
-        def recursive_transform(node: BigFrameNode) -> BigFrameNode:
-            return transform(node).transform_children(recursive_transform)
+        while stack:
+            node, t_node = stack[-1]
+
+            if t_node is None:
+                if node in results:
+                    stack.pop()
+                    continue
+                t_node = transform(node)
+                stack[-1] = (node, t_node)
+
+            all_done = True
+            for child in reversed(t_node.child_nodes):
+                if child not in results:
+                    stack.append((child, None))
+                    all_done = False
+                    break
+
+            if all_done:
+                results[node] = t_node.transform_children(lambda x: results[x])
+                stack.pop()
 
-        return recursive_transform(self)
+        return results[self]
 
     def bottom_up(
         self: BigFrameNode,

bigframes/core/compile/sqlglot/expressions/ai_ops.py

@@ -113,9 +113,9 @@ def _construct_named_args(op: ops.NaryOp) -> list[sge.Kwarg]:
             )
         )
 
-    endpoit = op_args.get("endpoint", None)
-    if endpoit is not None:
-        args.append(sge.Kwarg(this="endpoint", expression=sge.Literal.string(endpoit)))
+    endpoint = op_args.get("endpoint", None)
+    if endpoint is not None:
+        args.append(sge.Kwarg(this="endpoint", expression=sge.Literal.string(endpoint)))
 
     request_type = op_args.get("request_type", None)
     if request_type is not None:

bigframes/operations/ai_ops.py

@@ -123,7 +123,7 @@ class AIIf(base_ops.NaryOp):
     name: ClassVar[str] = "ai_if"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.BOOL_DTYPE
@@ -135,7 +135,7 @@ class AIClassify(base_ops.NaryOp):
 
     prompt_context: Tuple[str | None, ...]
     categories: tuple[str, ...]
-    connection_id: str
+    connection_id: str | None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.STRING_DTYPE
@@ -146,7 +146,7 @@ class AIScore(base_ops.NaryOp):
     name: ClassVar[str] = "ai_score"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.FLOAT_DTYPE

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/None/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  AI.CLASSIFY(input => (`string_col`), categories => ['greeting', 'rejection']) AS `result`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

…ots/test_ai_ops/test_ai_classify/out.sql ….us.bigframes-default-connection/out.sqltests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  AI.IF(prompt => (`string_col`, ' is the same as ', `string_col`)) AS `result`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/None/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  AI.SCORE(prompt => (`string_col`, ' is the same as ', `string_col`)) AS `result`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`