test: enables assorted tests (#16611)

chalmerlowe · parthea · jskeet · web-flow · commit 3cc859c8510c · 2026-04-14T13:05:42.000-04:00
This PR seeks to enable a number of tests to help ensure full validation. See #16489 for additional details. --------- Co-authored-by: Anthonios Partheniou <partheniou@google.com> Co-authored-by: Jon Skeet <jonskeet@google.com> Co-authored-by: Daniel Sanche <sanche@google.com> Co-authored-by: Cody Oss <6331106+codyoss@users.noreply.github.com> Co-authored-by: TrevorBergeron <tbergeron@google.com>
diff --git a/.kokoro/continuous/continuous-bigframes.cfg b/.kokoro/continuous/continuous-bigframes.cfg
@@ -0,0 +1,7 @@
+# Format: //devtools/kokoro/config/proto/build.proto
+
+# Only run these nox sessions.
+env_vars: {
+    key: "NOX_SESSION"
+    value: "e2e load system_prerelease notebook system_noextras"
+}
diff --git a/.kokoro/presubmit/presubmit-doctest-bigframes.cfg b/.kokoro/presubmit/presubmit-doctest-bigframes.cfg
@@ -0,0 +1,12 @@
+# Format: //devtools/kokoro/config/proto/build.proto
+
+# Only run this nox session.
+env_vars: {
+    key: "NOX_SESSION"
+    value: "cleanup doctest"
+}
+
+env_vars: {
+    key: "GOOGLE_CLOUD_PROJECT"
+    value: "bigframes-testing"
+}
diff --git a/.kokoro/system.sh b/.kokoro/system.sh
@@ -43,7 +43,8 @@ run_package_test() {
   local PROJECT_ID
   local GOOGLE_APPLICATION_CREDENTIALS
   local NOX_FILE
-  local NOX_SESSION
+  # Inherit NOX_SESSION from environment to allow configs (like prerelease.cfg) to pass it in
+  local NOX_SESSION="${NOX_SESSION}"
 
   echo "------------------------------------------------------------"
   echo "Configuring environment for: ${package_name}"
@@ -66,7 +67,8 @@ run_package_test() {
       PROJECT_ID=$(cat "${KOKORO_GFILE_DIR}/project-id.json")
       GOOGLE_APPLICATION_CREDENTIALS="${KOKORO_GFILE_DIR}/service-account.json"
       NOX_FILE="noxfile.py"
-      NOX_SESSION="system-3.12"
+      # Use inherited NOX_SESSION if set, otherwise fallback to system-3.12
+      NOX_SESSION="${NOX_SESSION:-system-3.12}"
       ;;
   esac
 
diff --git a/packages/bigframes/bigframes/core/blocks.py b/packages/bigframes/bigframes/core/blocks.py
@@ -2796,6 +2796,7 @@ def _is_monotonic(
                 )
                 block = block.drop_columns([equal_monotonic_id, strict_monotonic_id])
 
+        assert last_result_id is not None
         block, monotonic_result_id = block.apply_binary_op(
             last_result_id,
             last_notna_id,
diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/aggregate_compiler.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/aggregate_compiler.py
@@ -528,8 +528,9 @@ def _(
     column: ibis_types.Column,
     window=None,
 ) -> ibis_types.Value:
+    # Ibis FirstNonNullValue expects Value[Any, Columnar], Mypy struggles to see Column as compatible.
     return _apply_window_if_present(
-        ibis_ops.FirstNonNullValue(column).to_expr(),
+        ibis_ops.FirstNonNullValue(column).to_expr(),  # type: ignore[arg-type]
         window,  # type: ignore
     )
 
@@ -549,8 +550,9 @@ def _(
     column: ibis_types.Column,
     window=None,
 ) -> ibis_types.Value:
+    # Ibis LastNonNullValue expects Value[Any, Columnar], Mypy struggles to see Column as compatible.
     return _apply_window_if_present(
-        ibis_ops.LastNonNullValue(column).to_expr(),
+        ibis_ops.LastNonNullValue(column).to_expr(),  # type: ignore[arg-type]
         window,  # type: ignore
     )
 
@@ -803,8 +805,9 @@ def _to_ibis_boundary(
 ) -> Optional[ibis_expr_window.WindowBoundary]:
     if boundary is None:
         return None
+    # WindowBoundary expects Value[Any, Any], ibis_types.literal returns Scalar which Mypy doesn't see as compatible.
     return ibis_expr_window.WindowBoundary(
-        abs(boundary),
+        ibis_types.literal(boundary if boundary >= 0 else -boundary),  # type: ignore[arg-type]
         preceding=boundary <= 0,  # type:ignore
     )
 
diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/ibis_compiler.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/ibis_compiler.py
@@ -49,8 +49,8 @@ def compile_sql(request: configs.CompileRequest) -> configs.CompileResult:
         # Can only pullup slice if we are doing ORDER BY in outermost SELECT
         # Need to do this before replacing unsupported ops, as that will rewrite slice ops
         result_node = rewrites.pull_up_limits(result_node)
-    result_node = _replace_unsupported_ops(result_node)
-    result_node = result_node.bottom_up(rewrites.simplify_join)
+    result_node = cast(nodes.ResultNode, _replace_unsupported_ops(result_node))
+    result_node = cast(nodes.ResultNode, result_node.bottom_up(rewrites.simplify_join))
     # prune before pulling up order to avoid unnnecessary row_number() ops
     result_node = cast(nodes.ResultNode, rewrites.column_pruning(result_node))
     result_node = rewrites.defer_order(
diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/operations/geo_ops.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/operations/geo_ops.py
@@ -182,7 +182,9 @@ def st_buffer(
 
 @ibis_udf.scalar.builtin
 def st_distance(
-    a: ibis_dtypes.geography, b: ibis_dtypes.geography, use_spheroid: bool
+    a: ibis_dtypes.geography,  # type: ignore
+    b: ibis_dtypes.geography,  # type: ignore
+    use_spheroid: bool,  # type: ignore
 ) -> ibis_dtypes.float:  # type: ignore
     """Convert string to geography."""
 
diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py
@@ -2168,9 +2168,12 @@ def obj_make_ref_json(objectref_json: ibis_dtypes.JSON) -> _OBJ_REF_IBIS_DTYPE:
 
 
 @ibis_udf.scalar.builtin(name="OBJ.GET_ACCESS_URL")
-def obj_get_access_url(
-    obj_ref: _OBJ_REF_IBIS_DTYPE, mode: ibis_dtypes.String
-) -> ibis_dtypes.JSON:  # type: ignore
+# Stub for BigQuery UDF, empty body is intentional.
+# _OBJ_REF_IBIS_DTYPE is a variable holding a type, Mypy complains about it being used as type hint.
+def obj_get_access_url(  # type: ignore[empty-body]
+    obj_ref: _OBJ_REF_IBIS_DTYPE,  # type: ignore[valid-type]
+    mode: ibis_dtypes.String,
+) -> ibis_dtypes.JSON:
     """Get access url (as ObjectRefRumtime JSON) from ObjectRef."""
 
 
diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/compiler.py b/packages/bigframes/bigframes/core/compile/sqlglot/compiler.py
@@ -53,8 +53,10 @@ def compile_sql(request: configs.CompileRequest) -> configs.CompileResult:
         # Can only pullup slice if we are doing ORDER BY in outermost SELECT
         # Need to do this before replacing unsupported ops, as that will rewrite slice ops
         result_node = rewrite.pull_up_limits(result_node)
-    result_node = _replace_unsupported_ops(result_node)
-    result_node = result_node.bottom_up(rewrite.simplify_join)
+    result_node = typing.cast(nodes.ResultNode, _replace_unsupported_ops(result_node))
+    result_node = typing.cast(
+        nodes.ResultNode, result_node.bottom_up(rewrite.simplify_join)
+    )
     # prune before pulling up order to avoid unnnecessary row_number() ops
     result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node))
     result_node = rewrite.defer_order(
diff --git a/packages/bigframes/bigframes/core/expression_factoring.py b/packages/bigframes/bigframes/core/expression_factoring.py
@@ -243,7 +243,12 @@ def factor_aggregation(root: nodes.ColumnDef) -> FactoredAggregation:
     }
 
     root_scalar_expr = nodes.ColumnDef(
-        sub_expressions(root.expression, agg_outputs_dict),
+        sub_expressions(
+            root.expression,
+            cast(
+                Mapping[expression.Expression, expression.Expression], agg_outputs_dict
+            ),
+        ),
         root.id,  # type: ignore
     )
 
diff --git a/packages/bigframes/bigframes/core/local_data.py b/packages/bigframes/bigframes/core/local_data.py
@@ -33,7 +33,7 @@
 
 import bigframes.core.schema as schemata
 import bigframes.dtypes
-from bigframes.core import pyarrow_utils
+from bigframes.core import identifiers, pyarrow_utils
 
 
 @dataclasses.dataclass(frozen=True)
@@ -155,7 +155,7 @@ def to_arrow(
             return schema, batches
 
     def is_nullable(self, column_id: identifiers.ColumnId) -> bool:
-        return self.data.column(column_id).null_count > 0
+        return self.data.column(column_id.name).null_count > 0
 
     def to_pyarrow_table(
         self,
diff --git a/packages/bigframes/bigframes/core/nodes.py b/packages/bigframes/bigframes/core/nodes.py
@@ -674,7 +674,9 @@ def fields(self) -> Sequence[Field]:
             Field(
                 col_id,
                 self.local_data_source.schema.get_type(source_id),
-                nullable=self.local_data_source.is_nullable(source_id),
+                nullable=self.local_data_source.is_nullable(
+                    identifiers.ColumnId(source_id)
+                ),
             )
             for col_id, source_id in self.scan_list.items
         )
diff --git a/packages/bigframes/bigframes/core/rewrite/as_sql.py b/packages/bigframes/bigframes/core/rewrite/as_sql.py
@@ -291,8 +291,9 @@ def _extract_ctes_to_with_expr(
         root.top_down(lambda x: mapping.get(x, x)),
         cte_names,
         tuple(
-            cte_node.child.top_down(lambda x: mapping.get(x, x))
-            for cte_node in topological_ctes  # type: ignore
+            # Mypy loses context that cte_node is a CteNode with a child attribute, despite the isinstance filter above.
+            cte_node.child.top_down(lambda x: mapping.get(x, x))  # type: ignore[attr-defined]
+            for cte_node in topological_ctes
         ),
     )
 
diff --git a/packages/bigframes/bigframes/dataframe.py b/packages/bigframes/bigframes/dataframe.py
@@ -3926,12 +3926,13 @@ def round(self, decimals: Union[int, dict[Hashable, int]] = 0) -> DataFrame:
                 bigframes.dtypes.BOOL_DTYPE
             }:
                 if is_mapping:
-                    if label in decimals:  # type: ignore
+                    decimals_dict = typing.cast(dict[typing.Hashable, int], decimals)
+                    if label in decimals_dict:
                         exprs.append(
                             ops.round_op.as_expr(
                                 col_id,
                                 ex.const(
-                                    decimals[label],
+                                    decimals_dict[label],
                                     dtype=bigframes.dtypes.INT_DTYPE,  # type: ignore
                                 ),
                             )
@@ -4447,8 +4448,8 @@ def to_latex(
     ) -> str | None:
         return self.to_pandas(allow_large_results=allow_large_results).to_latex(
             buf,
-            columns=columns,
-            header=header,
+            columns=typing.cast(typing.Optional[list[str]], columns),
+            header=typing.cast(typing.Union[bool, list[str]], header),
             index=index,
             **kwargs,  # type: ignore
         )
diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py
@@ -2308,9 +2308,10 @@ def to_json(
             )
         else:
             pd_series = self.to_pandas(allow_large_results=allow_large_results)
+            # Pandas Series.to_json only supports a subset of orients, but bigframes Series.to_json allows all of them.
             return pd_series.to_json(
                 path_or_buf=path_or_buf,
-                orient=orient,
+                orient=orient,  # type: ignore[arg-type]
                 lines=lines,
                 index=index,  # type: ignore
             )
diff --git a/packages/bigframes/bigframes/session/iceberg.py b/packages/bigframes/bigframes/session/iceberg.py
@@ -98,9 +98,10 @@ def _extract_location_from_catalog_extension_data(data):
 
 
 class SchemaVisitor(pyiceberg.schema.SchemaVisitorPerPrimitiveType[bq.SchemaField]):
-    def schema(
+    # Override returns a tuple of fields instead of a single field, violating supertype signature but intentional for this visitor.
+    def schema(  # type: ignore[override]
         self, schema: pyiceberg.schema.Schema, struct_result: bq.SchemaField
-    ) -> tuple[bq.SchemaField, ...]:  # type: ignore
+    ) -> tuple[bq.SchemaField, ...]:
         return tuple(f for f in struct_result.fields)
 
     def struct(
diff --git a/packages/bigframes/noxfile.py b/packages/bigframes/noxfile.py
diff --git a/packages/bigframes/tests/system/large/test_session.py b/packages/bigframes/tests/system/large/test_session.py
diff --git a/packages/bigframes/tests/system/small/test_dataframe.py b/packages/bigframes/tests/system/small/test_dataframe.py
diff --git a/packages/bigframes/tests/system/small/test_session.py b/packages/bigframes/tests/system/small/test_session.py
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/tpch/conftest.py b/packages/bigframes/tests/unit/core/compile/sqlglot/tpch/conftest.py
diff --git a/packages/bigframes/tests/unit/test_col.py b/packages/bigframes/tests/unit/test_col.py
diff --git a/packages/bigframes/tests/unit/test_local_engine.py b/packages/bigframes/tests/unit/test_local_engine.py

Original file line number	Diff line number	Diff line change
`@@ -2796,6 +2796,7 @@ def _is_monotonic(`
`2796`	`2796`	`)`
`2797`	`2797`	`block = block.drop_columns([equal_monotonic_id, strict_monotonic_id])`
`2798`	`2798`
	`2799`	`+ assert last_result_id is not None`
`2799`	`2800`	`block, monotonic_result_id = block.apply_binary_op(`
`2800`	`2801`	`last_result_id,`
`2801`	`2802`	`last_notna_id,`
Original file line number	Diff line number	Diff line change
`@@ -243,7 +243,12 @@ def factor_aggregation(root: nodes.ColumnDef) -> FactoredAggregation:`
`243`	`243`	`}`
`244`	`244`
`245`	`245`	`root_scalar_expr = nodes.ColumnDef(`
`246`		`- sub_expressions(root.expression, agg_outputs_dict),`
	`246`	`+ sub_expressions(`
	`247`	`+ root.expression,`
	`248`	`+ cast(`
	`249`	`+ Mapping[expression.Expression, expression.Expression], agg_outputs_dict`
	`250`	`+ ),`
	`251`	`+ ),`
`247`	`252`	`root.id, # type: ignore`
`248`	`253`	`)`
`249`	`254`