Merge branch 'main' into window_null_skip_refactor

Author: TrevorBergeron

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/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/display/anywidget.py

@@ -55,7 +55,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)
     _initial_load_complete = traitlets.Bool(False).tag(sync=True)
     _batches: Optional[blocks.PandasBatches] = None
@@ -94,12 +98,17 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # SELECT COUNT(*) query. It is a must have however.
         # 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
 
@@ -131,11 +140,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)
 
@@ -229,6 +249,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

@@ -85,14 +85,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