perf: Reduce dry runs from read_gbq with table (#1129)

* perf: Reduce dry runs from read_gbq with table

* add check for table creation time to not travel before that

* handle table.created is None

Author: TrevorBergeron

bigframes/session/_io/bigquery/read_gbq_table.py

@@ -77,6 +77,11 @@ def get_table_metadata(
         return cached_table
 
     table = bqclient.get_table(table_ref)
+    # local time will lag a little bit do to network latency
+    # make sure it is at least table creation time.
+    # This is relevant if the table was created immediately before loading it here.
+    if (table.created is not None) and (table.created > bq_time):
+        bq_time = table.created
 
     cached_table = (bq_time, table)
     cache[table_ref] = cached_table
@@ -85,64 +90,66 @@ def get_table_metadata(
 
 def validate_table(
     bqclient: bigquery.Client,
-    table_ref: bigquery.table.TableReference,
+    table: bigquery.table.Table,
     columns: Optional[Sequence[str]],
     snapshot_time: datetime.datetime,
-    table_type: str,
     filter_str: Optional[str] = None,
 ) -> bool:
     """Validates that the table can be read, returns True iff snapshot is supported."""
-    # First run without snapshot to verify table can be read
-    sql = bigframes.session._io.bigquery.to_query(
-        query_or_table=f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}",
-        columns=columns or (),
-        sql_predicate=filter_str,
-    )
-    dry_run_config = bigquery.QueryJobConfig()
-    dry_run_config.dry_run = True
-    try:
-        bqclient.query_and_wait(sql, job_config=dry_run_config)
-    except google.api_core.exceptions.Forbidden as ex:
-        if "Drive credentials" in ex.message:
-            ex.message += "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions."
-        raise
 
+    time_travel_not_found = False
     # Anonymous dataset, does not support snapshot ever
-    if table_ref.dataset_id.startswith("_"):
-        return False
-
-    # Materialized views，does not support snapshot
-    if table_type == "MATERIALIZED_VIEW":
-        warnings.warn(
-            "Materialized views do not support FOR SYSTEM_TIME AS OF queries. "
-            "Attempting query without time travel. Be aware that as materialized views "
-            "are updated periodically, modifications to the underlying data in the view may "
-            "result in errors or unexpected behavior.",
-            category=bigframes.exceptions.TimeTravelDisabledWarning,
+    if table.dataset_id.startswith("_"):
+        pass
+    # Only true tables support time travel
+    elif table.table_type != "TABLE":
+        if table.table_type == "MATERIALIZED_VIEW":
+            warnings.warn(
+                "Materialized views do not support FOR SYSTEM_TIME AS OF queries. "
+                "Attempting query without time travel. Be aware that as materialized views "
+                "are updated periodically, modifications to the underlying data in the view may "
+                "result in errors or unexpected behavior.",
+                category=bigframes.exceptions.TimeTravelDisabledWarning,
+            )
+    else:
+        # table might support time travel, lets do a dry-run query with time travel
+        snapshot_sql = bigframes.session._io.bigquery.to_query(
+            query_or_table=f"{table.reference.project}.{table.reference.dataset_id}.{table.reference.table_id}",
+            columns=columns or (),
+            sql_predicate=filter_str,
+            time_travel_timestamp=snapshot_time,
         )
-        return False
+        try:
+            # If this succeeds, we don't need to query without time travel, that would surely succeed
+            bqclient.query_and_wait(
+                snapshot_sql, job_config=bigquery.QueryJobConfig(dry_run=True)
+            )
+            return True
+        except google.api_core.exceptions.NotFound:
+            # note that a notfound caused by a simple typo will be
+            # caught above when the metadata is fetched, not here
+            time_travel_not_found = True
 
-    # Second, try with snapshot to verify table supports this feature
+    # At this point, time travel is known to fail, but can we query without time travel?
     snapshot_sql = bigframes.session._io.bigquery.to_query(
-        query_or_table=f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}",
+        query_or_table=f"{table.reference.project}.{table.reference.dataset_id}.{table.reference.table_id}",
         columns=columns or (),
         sql_predicate=filter_str,
-        time_travel_timestamp=snapshot_time,
+        time_travel_timestamp=None,
     )
-    try:
-        bqclient.query_and_wait(snapshot_sql, job_config=dry_run_config)
-        return True
-    except google.api_core.exceptions.NotFound:
-        # note that a notfound caused by a simple typo will be
-        # caught above when the metadata is fetched, not here
+    # Any erorrs here should just be raised to user
+    bqclient.query_and_wait(
+        snapshot_sql, job_config=bigquery.QueryJobConfig(dry_run=True)
+    )
+    if time_travel_not_found:
         warnings.warn(
             "NotFound error when reading table with time travel."
             " Attempting query without time travel. Warning: Without"
             " time travel, modifications to the underlying table may"
             " result in errors or unexpected behavior.",
             category=bigframes.exceptions.TimeTravelDisabledWarning,
         )
-        return False
+    return False
 
 
 def are_index_cols_unique(

bigframes/session/loader.py

@@ -343,14 +343,18 @@ def read_gbq_table(
             else (*columns, *[col for col in index_cols if col not in columns])
         )
 
-        enable_snapshot = enable_snapshot and bf_read_gbq_table.validate_table(
-            self._bqclient,
-            table_ref,
-            all_columns,
-            time_travel_timestamp,
-            table.table_type,
-            filter_str,
-        )
+        try:
+            enable_snapshot = enable_snapshot and bf_read_gbq_table.validate_table(
+                self._bqclient,
+                table,
+                all_columns,
+                time_travel_timestamp,
+                filter_str,
+            )
+        except google.api_core.exceptions.Forbidden as ex:
+            if "Drive credentials" in ex.message:
+                ex.message += "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions."
+            raise
 
         # ----------------------------
         # Create ordering and validate

tests/unit/session/test_session.py

@@ -181,6 +181,7 @@ def test_read_gbq_cached_table():
     table._properties["location"] = session._location
     table._properties["numRows"] = "1000000000"
     table._properties["location"] = session._location
+    table._properties["type"] = "TABLE"
     session._loader._df_snapshot[table_ref] = (
         datetime.datetime(1999, 1, 2, 3, 4, 5, 678901, tzinfo=datetime.timezone.utc),
         table,