fix: use table clone instead of system time for read_gbq_table

bigframes/constants.py

@@ -26,4 +26,4 @@
 
 ABSTRACT_METHOD_ERROR_MESSAGE = f"Abstract method. You have likely encountered a bug. Please share this stacktrace and how you reached it with the BigQuery DataFrames team. {FEEDBACK_LINK}"
 
-DEFAULT_EXPIRATION = datetime.timedelta(days=1)
+DEFAULT_EXPIRATION = datetime.timedelta(days=7)

bigframes/session/__init__.py

@@ -430,7 +430,9 @@ def _read_gbq_query(
             index_cols = list(index_col)
 
         destination, query_job = self._query_to_destination(
-            query, index_cols, api_name="read_gbq_query"
+            query,
+            index_cols,
+            api_name=api_name,
         )
 
         # If there was no destination table, that means the query must have
@@ -508,25 +510,29 @@ def _read_gbq_table_to_ibis_with_total_ordering(
         If we can get a total ordering from the table, such as via primary key
         column(s), then return those too so that ordering generation can be
         avoided.
-        """
-        if table_ref.dataset_id.upper() == "_SESSION":
-            # _SESSION tables aren't supported by the tables.get REST API.
-            return (
-                self.ibis_client.sql(
-                    f"SELECT * FROM `_SESSION`.`{table_ref.table_id}`"
-                ),
-                None,
-            )
 
+        For tables that aren't already read-only, this creates Create a table
+        clone so that any changes to the underlying table don't affect the
+        DataFrame and break our assumptions, especially with regards to unique
+        index and ordering. See:
+        https://cloud.google.com/bigquery/docs/table-clones-create
+        """
+        destination = bigframes_io.create_table_clone(
+            table_ref,
+            self._anonymous_dataset,
+            constants.DEFAULT_EXPIRATION,
+            self,
+            api_name,
+        )
         table_expression = self.ibis_client.table(
-            table_ref.table_id,
-            database=f"{table_ref.project}.{table_ref.dataset_id}",
+            destination.table_id,
+            database=f"{destination.project}.{destination.dataset_id}",
         )
 
         # If there are primary keys defined, the query engine assumes these
         # columns are unique, even if the constraint is not enforced. We make
         # the same assumption and use these columns as the total ordering keys.
-        table = self.bqclient.get_table(table_ref)
+        table = self.bqclient.get_table(destination)
 
         # TODO(b/305264153): Use public properties to fetch primary keys once
         # added to google-cloud-bigquery.
@@ -535,23 +541,7 @@ def _read_gbq_table_to_ibis_with_total_ordering(
             .get("primaryKey", {})
             .get("columns")
         )
-
-        if not primary_keys:
-            return table_expression, None
-        else:
-            # Read from a snapshot since we won't have to copy the table data to create a total ordering.
-            job_config = bigquery.QueryJobConfig()
-            job_config.labels["bigframes-api"] = api_name
-            current_timestamp = list(
-                self.bqclient.query(
-                    "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`",
-                    job_config=job_config,
-                ).result()
-            )[0][0]
-            table_expression = self.ibis_client.sql(
-                bigframes_io.create_snapshot_sql(table_ref, current_timestamp)
-            )
-            return table_expression, primary_keys
+        return table_expression, primary_keys
 
     def _read_gbq_table(
         self,
@@ -672,9 +662,6 @@ def _read_gbq_table(
 
                 # The job finished, so we should have a start time.
                 assert current_timestamp is not None
-                table_expression = self.ibis_client.sql(
-                    bigframes_io.create_snapshot_sql(table_ref, current_timestamp)
-                )
             else:
                 # Make sure when we generate an ordering, the row_number()
                 # coresponds to the index columns.

bigframes/session/_io/bigquery.py

@@ -14,14 +14,21 @@
 
 """Private module: Helpers for I/O operations."""
 
+from __future__ import annotations
+
 import datetime
 import textwrap
 import types
+import typing
 from typing import Dict, Iterable, Union
 import uuid
 
 import google.cloud.bigquery as bigquery
 
+if typing.TYPE_CHECKING:
+    import bigframes.session
+
+
 IO_ORDERING_ID = "bqdf_row_nums"
 TEMP_TABLE_PREFIX = "bqdf{date}_{random_id}"
 
@@ -69,27 +76,52 @@ def create_export_data_statement(
     )
 
 
-def create_snapshot_sql(
-    table_ref: bigquery.TableReference, current_timestamp: datetime.datetime
-) -> str:
-    """Query a table via 'time travel' for consistent reads."""
+def random_table(dataset: bigquery.DatasetReference) -> bigquery.TableReference:
+    now = datetime.datetime.now(datetime.timezone.utc)
+    random_id = uuid.uuid4().hex
+    table_id = TEMP_TABLE_PREFIX.format(
+        date=now.strftime("%Y%m%d"), random_id=random_id
+    )
+    return dataset.table(table_id)
 
-    # If we have a _SESSION table, assume that it's already a copy. Nothing to do here.
-    if table_ref.dataset_id.upper() == "_SESSION":
-        return f"SELECT * FROM `_SESSION`.`{table_ref.table_id}`"
 
-    # If we have an anonymous query results table, it can't be modified and
-    # there isn't any BigQuery time travel.
-    if table_ref.dataset_id.startswith("_"):
-        return f"SELECT * FROM `{table_ref.project}`.`{table_ref.dataset_id}`.`{table_ref.table_id}`"
+def table_ref_to_sql(table: bigquery.TableReference) -> str:
+    return f"`{table.project}`.`{table.dataset_id}`.`{table.table_id}`"
 
-    return textwrap.dedent(
+
+def create_table_clone(
+    source: bigquery.TableReference,
+    dataset: bigquery.DatasetReference,
+    expiration: datetime.timedelta,
+    session: bigframes.session.Session,
+    api_name: str,
+) -> bigquery.TableReference:
+    """Create a table clone for consistent reads."""
+    now = datetime.datetime.now(datetime.timezone.utc)
+    expiration_timestamp = now + expiration
+    fully_qualified_source_id = table_ref_to_sql(source)
+    destination = random_table(dataset)
+    fully_qualified_destination_id = table_ref_to_sql(destination)
+
+    # Include a label so that Dataplex Lineage can identify temporary
+    # tables that BigQuery DataFrames creates. Googlers: See internal issue
+    # 296779699.
+    ddl = textwrap.dedent(
         f"""
-        SELECT *
-        FROM `{table_ref.project}`.`{table_ref.dataset_id}`.`{table_ref.table_id}`
-        FOR SYSTEM_TIME AS OF TIMESTAMP({repr(current_timestamp.isoformat())})
+        CREATE OR REPLACE TABLE
+        {fully_qualified_destination_id}
+        CLONE {fully_qualified_source_id}
+        OPTIONS(
+            expiration_timestamp=TIMESTAMP "{expiration_timestamp.isoformat()}",
+            labels=[
+                ("source", "bigquery-dataframes-temp"),
+                ("bigframes-api", {repr(api_name)})
+            ]
+        )
         """
     )
+    session._start_query(ddl)
+    return destination
 
 
 def create_temp_table(
@@ -98,13 +130,9 @@ def create_temp_table(
     expiration: datetime.timedelta,
 ) -> str:
     """Create an empty table with an expiration in the desired dataset."""
-    now = datetime.datetime.now(datetime.timezone.utc)
-    random_id = uuid.uuid4().hex
-    table_id = TEMP_TABLE_PREFIX.format(
-        date=now.strftime("%Y%m%d"), random_id=random_id
-    )
-    table_ref = dataset.table(table_id)
+    table_ref = random_table(dataset)
     destination = bigquery.Table(table_ref)
+    now = datetime.datetime.now(datetime.timezone.utc)
     destination.expires = now + expiration
     bqclient.create_table(destination)
     return f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}"

tests/unit/session/test_io_bigquery.py

@@ -22,12 +22,12 @@
 import bigframes.session._io.bigquery
 
 
-def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets():
+def test_create_table_clone_doesnt_timetravel_anonymous_datasets():
     table_ref = bigquery.TableReference.from_string(
         "my-test-project._e8166e0cdb.anonbb92cd"
     )
 
-    sql = bigframes.session._io.bigquery.create_snapshot_sql(
+    sql = bigframes.core.io.create_table_clone(
         table_ref, datetime.datetime.now(datetime.timezone.utc)
     )
 
@@ -38,10 +38,10 @@ def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets():
     assert "`my-test-project`.`_e8166e0cdb`.`anonbb92cd`" in sql
 
 
-def test_create_snapshot_sql_doesnt_timetravel_session_tables():
+def test_create_table_clone_doesnt_timetravel_session_datasets():
     table_ref = bigquery.TableReference.from_string("my-test-project._session.abcdefg")
 
-    sql = bigframes.session._io.bigquery.create_snapshot_sql(
+    sql = bigframes.core.io.create_table_clone(
         table_ref, datetime.datetime.now(datetime.timezone.utc)
     )