Merge branch 'main' into shuowei-feat-persist-obj-ref

Author: shuoweil

.github/workflows/docs.yml

@@ -36,6 +36,3 @@ jobs:
       run: |
         python -m pip install --upgrade setuptools pip wheel
         python -m pip install nox
-    - name: Run docfx
-      run: |
-        nox -s docfx

bigframes/__init__.py

@@ -32,6 +32,9 @@
 )
 import bigframes.enums as enums  # noqa: E402
 import bigframes.exceptions as exceptions  # noqa: E402
+
+# Register pandas extensions
+import bigframes.extensions.pandas.dataframe_accessor  # noqa: F401, E402
 from bigframes.session import connect, Session  # noqa: E402
 from bigframes.version import __version__  # noqa: E402
 

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/bigquery/_operations/io.py

@@ -19,8 +19,8 @@
 import pandas as pd
 
 from bigframes.bigquery._operations.table import _get_table_metadata
+import bigframes.core.compile.sqlglot.sql as sql
 import bigframes.core.logging.log_adapter as log_adapter
-import bigframes.core.sql.io
 import bigframes.session
 
 
@@ -73,7 +73,7 @@ def load_data(
     """
     import bigframes.pandas as bpd
 
-    sql = bigframes.core.sql.io.load_data_ddl(
+    load_data_expr = sql.load_data(
         table_name=table_name,
         write_disposition=write_disposition,
         columns=columns,
@@ -84,11 +84,12 @@ def load_data(
         with_partition_columns=with_partition_columns,
         connection_name=connection_name,
     )
+    sql_text = sql.to_sql(load_data_expr)
 
     if session is None:
-        bpd.read_gbq_query(sql)
+        bpd.read_gbq_query(sql_text)
         session = bpd.get_global_session()
     else:
-        session.read_gbq_query(sql)
+        session.read_gbq_query(sql_text)
 
     return _get_table_metadata(bqclient=session.bqclient, table_name=table_name)

bigframes/bigquery/_operations/sql.py

@@ -16,19 +16,31 @@
 
 from __future__ import annotations
 
-from typing import Sequence
+from typing import cast, Optional, Sequence, Union
 
 import google.cloud.bigquery
 
 from bigframes.core.compile.sqlglot import sql
+import bigframes.dataframe
 import bigframes.dtypes
 import bigframes.operations
 import bigframes.series
 
 
+def _format_names(sql_template: str, dataframe: bigframes.dataframe.DataFrame):
+    """Turn sql_template from a template that uses names to one that uses
+    numbers.
+    """
+    names_to_numbers = {name: f"{{{i}}}" for i, name in enumerate(dataframe.columns)}
+    numbers = [f"{{{i}}}" for i in range(len(dataframe.columns))]
+    return sql_template.format(*numbers, **names_to_numbers)
+
+
 def sql_scalar(
     sql_template: str,
-    columns: Sequence[bigframes.series.Series],
+    columns: Union[bigframes.dataframe.DataFrame, Sequence[bigframes.series.Series]],
+    *,
+    output_dtype: Optional[bigframes.dtypes.Dtype] = None,
 ) -> bigframes.series.Series:
     """Create a Series from a SQL template.
 
@@ -37,6 +49,9 @@ def sql_scalar(
         >>> import bigframes.pandas as bpd
         >>> import bigframes.bigquery as bbq
 
+    Either pass in a sequence of series, in which case use  integers in the
+    format strings.
+
         >>> s = bpd.Series(["1.5", "2.5", "3.5"])
         >>> s = s.astype(pd.ArrowDtype(pa.decimal128(38, 9)))
         >>> bbq.sql_scalar("ROUND({0}, 0, 'ROUND_HALF_EVEN')", [s])
@@ -45,13 +60,29 @@ def sql_scalar(
         2    4.000000000
         dtype: decimal128(38, 9)[pyarrow]
 
+    Or pass in a DataFrame, in which case use the column names in the format
+    strings.
+
+        >>> df = bpd.DataFrame({"a": ["1.5", "2.5", "3.5"]})
+        >>> df = df.astype({"a": pd.ArrowDtype(pa.decimal128(38, 9))})
+        >>> bbq.sql_scalar("ROUND({a}, 0, 'ROUND_HALF_EVEN')", df)
+        0    2.000000000
+        1    2.000000000
+        2    4.000000000
+        dtype: decimal128(38, 9)[pyarrow]
+
     Args:
         sql_template (str):
             A SQL format string with Python-style {0} placeholders for each of
             the Series objects in ``columns``.
-        columns (Sequence[bigframes.pandas.Series]):
+        columns (
+            Sequence[bigframes.pandas.Series] | bigframes.pandas.DataFrame
+        ):
             Series objects representing the column inputs to the
             ``sql_template``. Must contain at least one Series.
+        output_dtype (a BigQuery DataFrames compatible dtype, optional):
+            If provided, BigQuery DataFrames uses this to determine the output
+            of the returned Series. This avoids a dry run query.
 
     Returns:
         bigframes.pandas.Series:
@@ -60,28 +91,38 @@ def sql_scalar(
     Raises:
         ValueError: If ``columns`` is empty.
     """
+    if isinstance(columns, bigframes.dataframe.DataFrame):
+        sql_template = _format_names(sql_template, columns)
+        columns = [
+            cast(bigframes.series.Series, columns[column]) for column in columns.columns
+        ]
+
     if len(columns) == 0:
         raise ValueError("Must provide at least one column in columns")
 
+    base_series = columns[0]
+
     # To integrate this into our expression trees, we need to get the output
     # type, so we do some manual compilation and a dry run query to get that.
     # Another benefit of this is that if there is a syntax error in the SQL
     # template, then this will fail with an error earlier in the process,
     # aiding users in debugging.
-    literals_sql = [sql.to_sql(sql.literal(None, column.dtype)) for column in columns]
-    select_sql = sql_template.format(*literals_sql)
-    dry_run_sql = f"SELECT {select_sql}"
-
-    # Use the executor directly, because we want the original column IDs, not
-    # the user-friendly column names that block.to_sql_query() would produce.
-    base_series = columns[0]
-    bqclient = base_series._session.bqclient
-    job = bqclient.query(
-        dry_run_sql, job_config=google.cloud.bigquery.QueryJobConfig(dry_run=True)
-    )
-    _, output_type = bigframes.dtypes.convert_schema_field(job.schema[0])
+    if output_dtype is None:
+        literals_sql = [
+            sql.to_sql(sql.literal(None, column.dtype)) for column in columns
+        ]
+        select_sql = sql_template.format(*literals_sql)
+        dry_run_sql = f"SELECT {select_sql}"
+
+        # Use the executor directly, because we want the original column IDs, not
+        # the user-friendly column names that block.to_sql_query() would produce.
+        bqclient = base_series._session.bqclient
+        job = bqclient.query(
+            dry_run_sql, job_config=google.cloud.bigquery.QueryJobConfig(dry_run=True)
+        )
+        _, output_dtype = bigframes.dtypes.convert_schema_field(job.schema[0])
 
     op = bigframes.operations.SqlScalarOp(
-        _output_type=output_type, sql_template=sql_template
+        _output_type=output_dtype, sql_template=sql_template
     )
     return base_series._apply_nary_op(op, columns[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/core/compile/sqlglot/expressions/datetime_ops.py

@@ -371,7 +371,12 @@ def _(expr: TypedExpr, op: ops.ToDatetimeOp) -> sge.Expression:
         )
         return sge.Cast(this=result, to="DATETIME")
 
-    if expr.dtype in (dtypes.STRING_DTYPE, dtypes.TIMESTAMP_DTYPE):
+    if expr.dtype in (
+        dtypes.STRING_DTYPE,
+        dtypes.TIMESTAMP_DTYPE,
+        dtypes.DATETIME_DTYPE,
+        dtypes.DATE_DTYPE,
+    ):
         return sge.TryCast(this=expr.expr, to="DATETIME")
 
     value = expr.expr
@@ -396,7 +401,12 @@ def _(expr: TypedExpr, op: ops.ToTimestampOp) -> sge.Expression:
             "PARSE_TIMESTAMP", sge.convert(op.format), expr.expr, sge.convert("UTC")
         )
 
-    if expr.dtype in (dtypes.STRING_DTYPE, dtypes.DATETIME_DTYPE):
+    if expr.dtype in (
+        dtypes.STRING_DTYPE,
+        dtypes.DATETIME_DTYPE,
+        dtypes.TIMESTAMP_DTYPE,
+        dtypes.DATE_DTYPE,
+    ):
         return sge.func("TIMESTAMP", expr.expr)
 
     value = expr.expr

bigframes/core/compile/sqlglot/sql/__init__.py

@@ -22,6 +22,7 @@
     table,
     to_sql,
 )
+from bigframes.core.compile.sqlglot.sql.ddl import load_data
 from bigframes.core.compile.sqlglot.sql.dml import insert, replace
 
 __all__ = [
@@ -33,6 +34,8 @@
     "literal",
     "table",
     "to_sql",
+    # From ddl.py
+    "load_data",
     # From dml.py
     "insert",
     "replace",