merge main

Author: shuoweil

bigframes/bigquery/_operations/ai.py

@@ -88,7 +88,7 @@ def generate(
             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.
         endpoint (str, optional):
             Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
             generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
@@ -131,7 +131,7 @@ def generate(
 
     operator = ai_ops.AIGenerate(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
         endpoint=endpoint,
         request_type=request_type,
         model_params=json.dumps(model_params) if model_params else None,
@@ -186,7 +186,7 @@ def generate_bool(
             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.
         endpoint (str, optional):
             Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
             generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
@@ -216,7 +216,7 @@ def generate_bool(
 
     operator = ai_ops.AIGenerateBool(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
         endpoint=endpoint,
         request_type=request_type,
         model_params=json.dumps(model_params) if model_params else None,
@@ -267,7 +267,7 @@ def generate_int(
             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.
         endpoint (str, optional):
             Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
             generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
@@ -297,7 +297,7 @@ def generate_int(
 
     operator = ai_ops.AIGenerateInt(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
         endpoint=endpoint,
         request_type=request_type,
         model_params=json.dumps(model_params) if model_params else None,
@@ -348,7 +348,7 @@ def generate_double(
             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.
         endpoint (str, optional):
             Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
             generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
@@ -378,7 +378,7 @@ def generate_double(
 
     operator = ai_ops.AIGenerateDouble(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
         endpoint=endpoint,
         request_type=request_type,
         model_params=json.dumps(model_params) if model_params else None,

bigframes/core/compile/sqlglot/aggregations/ordered_unary_compiler.py

@@ -27,7 +27,7 @@ def compile(
     op: agg_ops.WindowOp,
     column: typed_expr.TypedExpr,
     *,
-    order_by: tuple[sge.Expression, ...],
+    order_by: tuple[sge.Expression, ...] = (),
 ) -> sge.Expression:
     return ORDERED_UNARY_OP_REGISTRATION[op](op, column, order_by=order_by)
 

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -49,6 +49,19 @@ def _(
     return sge.func("IFNULL", result, sge.true())
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.AnyOp)
+def _(
+    op: agg_ops.AnyOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    expr = column.expr
+    expr = apply_window_if_present(sge.func("LOGICAL_OR", expr), window)
+
+    # BQ will return null for empty column, result would be false in pandas.
+    return sge.func("COALESCE", expr, sge.convert(False))
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.ApproxQuartilesOp)
 def _(
     op: agg_ops.ApproxQuartilesOp,

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

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

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

@@ -16,19 +16,16 @@
 
 import typing
 
-import sqlglot
+import sqlglot as sg
 import sqlglot.expressions as sge
 
 from bigframes import operations as ops
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+import bigframes.dtypes as dtypes
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
-
-
-@register_unary_op(ops.ArrayToStringOp, pass_op=True)
-def _(expr: TypedExpr, op: ops.ArrayToStringOp) -> sge.Expression:
-    return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
+register_nary_op = scalar_compiler.scalar_op_compiler.register_nary_op
 
 
 @register_unary_op(ops.ArrayIndexOp, pass_op=True)
@@ -41,17 +38,45 @@ def _(expr: TypedExpr, op: ops.ArrayIndexOp) -> sge.Expression:
     )
 
 
+@register_unary_op(ops.ArrayReduceOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.ArrayReduceOp) -> sge.Expression:
+    sub_expr = sg.to_identifier("bf_arr_reduce_uid")
+    sub_type = dtypes.get_array_inner_type(expr.dtype)
+
+    if op.aggregation.order_independent:
+        from bigframes.core.compile.sqlglot.aggregations import unary_compiler
+
+        agg_expr = unary_compiler.compile(op.aggregation, TypedExpr(sub_expr, sub_type))
+    else:
+        from bigframes.core.compile.sqlglot.aggregations import ordered_unary_compiler
+
+        agg_expr = ordered_unary_compiler.compile(
+            op.aggregation, TypedExpr(sub_expr, sub_type)
+        )
+
+    return (
+        sge.select(agg_expr)
+        .from_(
+            sge.Unnest(
+                expressions=[expr.expr],
+                alias=sge.TableAlias(columns=[sub_expr]),
+            )
+        )
+        .subquery()
+    )
+
+
 @register_unary_op(ops.ArraySliceOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:
-    slice_idx = sqlglot.to_identifier("slice_idx")
+    slice_idx = sg.to_identifier("slice_idx")
 
     conditions: typing.List[sge.Predicate] = [slice_idx >= op.start]
 
     if op.stop is not None:
         conditions.append(slice_idx < op.stop)
 
     # local name for each element in the array
-    el = sqlglot.to_identifier("el")
+    el = sg.to_identifier("el")
 
     selected_elements = (
         sge.select(el)
@@ -66,3 +91,27 @@ def _(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:
     )
 
     return sge.array(selected_elements)
+
+
+@register_unary_op(ops.ArrayToStringOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.ArrayToStringOp) -> sge.Expression:
+    return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
+
+
+@register_nary_op(ops.ToArrayOp)
+def _(*exprs: TypedExpr) -> sge.Expression:
+    do_upcast_bool = any(
+        dtypes.is_numeric(expr.dtype, include_bool=False) for expr in exprs
+    )
+    if do_upcast_bool:
+        sg_exprs = [_coerce_bool_to_int(expr) for expr in exprs]
+    else:
+        sg_exprs = [expr.expr for expr in exprs]
+    return sge.Array(expressions=sg_exprs)
+
+
+def _coerce_bool_to_int(typed_expr: TypedExpr) -> sge.Expression:
+    """Coerce boolean expression to integer."""
+    if typed_expr.dtype == dtypes.BOOL_DTYPE:
+        return sge.Cast(this=typed_expr.expr, to="INT64")
+    return typed_expr.expr

bigframes/display/anywidget.py

@@ -56,7 +56,11 @@ class TableWidget(WIDGET_BASE):
 
     page = traitlets.Int(0).tag(sync=True)
     page_size = traitlets.Int(0).tag(sync=True)
-    row_count = traitlets.Int(0).tag(sync=True)
+    row_count = traitlets.Union(
+        [traitlets.Int(), traitlets.Instance(type(None))],
+        default_value=None,
+        allow_none=True,
+    ).tag(sync=True)
     table_html = traitlets.Unicode().tag(sync=True)
     sort_column = traitlets.Unicode("").tag(sync=True)
     sort_ascending = traitlets.Bool(True).tag(sync=True)
@@ -103,12 +107,17 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # obtain the row counts
         # TODO(b/428238610): Start iterating over the result of `to_pandas_batches()`
         # before we get here so that the count might already be cached.
-        # TODO(b/452747934): Allow row_count to be None and check to see if
-        # there are multiple pages and show "page 1 of many" in this case
         self._reset_batches_for_new_page_size()
-        if self._batches is None or self._batches.total_rows is None:
-            self._error_message = "Could not determine total row count. Data might be unavailable or an error occurred."
-            self.row_count = 0
+
+        if self._batches is None:
+            self._error_message = "Could not retrieve data batches. Data might be unavailable or an error occurred."
+            self.row_count = None
+        elif self._batches.total_rows is None:
+            # Total rows is unknown, this is an expected state.
+            # TODO(b/461536343): Cheaply discover if we have exactly 1 page.
+            # There are cases where total rows is not set, but there are no additional
+            # pages. We could disable the "next" button in these cases.
+            self.row_count = None
         else:
             self.row_count = self._batches.total_rows
 
@@ -141,11 +150,22 @@ def _validate_page(self, proposal: Dict[str, Any]) -> int:
         Returns:
             The validated and clamped page number as an integer.
         """
-
         value = proposal["value"]
+
+        if value < 0:
+            raise ValueError("Page number cannot be negative.")
+
+        # If truly empty or invalid page size, stay on page 0.
+        # This handles cases where row_count is 0 or page_size is 0, preventing
+        # division by zero or nonsensical pagination, regardless of row_count being None.
         if self.row_count == 0 or self.page_size == 0:
             return 0
 
+        # If row count is unknown, allow any non-negative page. The previous check
+        # ensures that invalid page_size (0) is already handled.
+        if self.row_count is None:
+            return value
+
         # Calculate the zero-indexed maximum page number.
         max_page = max(0, math.ceil(self.row_count / self.page_size) - 1)
 
@@ -260,6 +280,23 @@ def _set_table_html(self) -> None:
         # Get the data for the current page
         page_data = cached_data.iloc[start:end]
 
+        # Handle case where user navigated beyond available data with unknown row count
+        is_unknown_count = self.row_count is None
+        is_beyond_data = self._all_data_loaded and len(page_data) == 0 and self.page > 0
+        if is_unknown_count and is_beyond_data:
+            # Calculate the last valid page (zero-indexed)
+            total_rows = len(cached_data)
+            if total_rows > 0:
+                last_valid_page = max(0, math.ceil(total_rows / self.page_size) - 1)
+                # Navigate back to the last valid page
+                self.page = last_valid_page
+                # Recursively call to display the correct page
+                return self._set_table_html()
+            else:
+                # If no data at all, stay on page 0 with empty display
+                self.page = 0
+                return self._set_table_html()
+
         # Generate HTML table
         self.table_html = bigframes.display.html.render_html(
             dataframe=page_data,

bigframes/display/table_widget.js

@@ -99,14 +99,21 @@ function render({ model, el }) {
 		const rowCount = model.get(ModelProperty.ROW_COUNT);
 		const pageSize = model.get(ModelProperty.PAGE_SIZE);
 		const currentPage = model.get(ModelProperty.PAGE);
-		const totalPages = Math.ceil(rowCount / pageSize);
-
-		rowCountLabel.textContent = `${rowCount.toLocaleString()} total rows`;
-		paginationLabel.textContent = `Page ${(
-			currentPage + 1
-		).toLocaleString()} of ${(totalPages || 1).toLocaleString()}`;
-		prevPage.disabled = currentPage === 0;
-		nextPage.disabled = currentPage >= totalPages - 1;
+
+		if (rowCount === null) {
+			// Unknown total rows
+			rowCountLabel.textContent = "Total rows unknown";
+			paginationLabel.textContent = `Page ${(currentPage + 1).toLocaleString()} of many`;
+			prevPage.disabled = currentPage === 0;
+			nextPage.disabled = false; // Allow navigation until we hit the end
+		} else {
+			// Known total rows
+			const totalPages = Math.ceil(rowCount / pageSize);
+			rowCountLabel.textContent = `${rowCount.toLocaleString()} total rows`;
+			paginationLabel.textContent = `Page ${(currentPage + 1).toLocaleString()} of ${rowCount.toLocaleString()}`;
+			prevPage.disabled = currentPage === 0;
+			nextPage.disabled = currentPage >= totalPages - 1;
+		}
 		pageSizeSelect.value = pageSize;
 	}
 

bigframes/operations/ai_ops.py

@@ -29,7 +29,7 @@ class AIGenerate(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]
     model_params: str | None
@@ -57,7 +57,7 @@ class AIGenerateBool(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_bool"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]
     model_params: str | None
@@ -79,7 +79,7 @@ class AIGenerateInt(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_int"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]
     model_params: str | None
@@ -101,7 +101,7 @@ class AIGenerateDouble(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_double"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]
     model_params: str | None