chore: fix mypy errors, exclude experimental notebook from testing (#1083)

* chore: fix mypy errors

* use quoted annotation to avoid error in non type-checking environment

* exclude experimental semantic operators notebook from testing

Author: shobsi

bigframes/dataframe.py

@@ -81,6 +81,8 @@
 import bigframes.session._io.bigquery
 
 if typing.TYPE_CHECKING:
+    from _typeshed import SupportsRichComparison
+
     import bigframes.session
 
     SingleItemValue = Union[bigframes.series.Series, int, float, Callable]
@@ -2464,7 +2466,7 @@ def pivot_table(
             values = [values]
 
         # Unlike pivot, pivot_table has values always ordered.
-        values.sort()
+        values.sort(key=lambda val: typing.cast("SupportsRichComparison", val))
 
         keys = index + columns
         agged = self.groupby(keys, dropna=True)[values].agg(aggfunc)

bigframes/session/loader.py

@@ -300,7 +300,9 @@ def read_gbq_table(
         ):
             # TODO(b/338111344): If we are running a query anyway, we might as
             # well generate ROW_NUMBER() at the same time.
-            all_columns = itertools.chain(index_cols, columns) if columns else ()
+            all_columns: Iterable[str] = (
+                itertools.chain(index_cols, columns) if columns else ()
+            )
             query = bf_io_bigquery.to_query(
                 query,
                 columns=all_columns,

bigframes/session/metrics.py

@@ -44,7 +44,7 @@ def count_job_stats(self, query_job: bq_job.QueryJob):
 
 def get_performance_stats(
     query_job: bigquery.QueryJob,
-) -> Optional[Tuple[int, int, float]]:
+) -> Optional[Tuple[int, int, Optional[float]]]:
     """Parse the query job for performance stats.
 
     Return None if the stats do not reflect real work done in bigquery.
@@ -73,7 +73,9 @@ def get_performance_stats(
     return bytes_processed, slot_millis, exec_seconds
 
 
-def write_stats_to_disk(bytes_processed: int, slot_millis: int, exec_seconds: float):
+def write_stats_to_disk(
+    bytes_processed: int, slot_millis: int, exec_seconds: Optional[float]
+):
     """For pytest runs only, log information about the query job
     to a file in order to create a performance report.
     """

noxfile.py

@@ -744,6 +744,7 @@ def notebook(session: nox.Session):
         # The experimental notebooks imagine features that don't yet
         # exist or only exist as temporary prototypes.
         "notebooks/experimental/longer_ml_demo.ipynb",
+        "notebooks/experimental/semantic_operators.ipynb",
         # The notebooks that are added for more use cases, such as backing a
         # blog post, which may take longer to execute and need not be
         # continuously tested.