fix: avoid table modification for to_gbq

Author: chelsea-lin

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -17,6 +17,7 @@
 import dataclasses
 import typing
 
+from google.cloud import bigquery
 import pyarrow as pa
 import sqlglot as sg
 import sqlglot.dialects.bigquery
@@ -104,6 +105,24 @@ def from_pyarrow(
         )
         return cls(expr=sg.select(sge.Star()).from_(expr), uid_gen=uid_gen)
 
+    @classmethod
+    def from_query_string(
+        cls,
+        query_string: str,
+    ) -> SQLGlotIR:
+        """Builds SQLGlot expression from a query string"""
+        uid_gen: guid.SequentialUIDGenerator = guid.SequentialUIDGenerator()
+        cte_name = sge.to_identifier(
+            next(uid_gen.get_uid_stream("bfcte_")), quoted=cls.quoted
+        )
+        cte = sge.CTE(
+            this=query_string,
+            alias=cte_name,
+        )
+        select_expr = sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
+        select_expr.set("with", sge.With(expressions=[cte]))
+        return cls(expr=select_expr, uid_gen=uid_gen)
+
     def select(
         self,
         selected_cols: tuple[tuple[str, sge.Expression], ...],
@@ -133,6 +152,36 @@ def project(
         select_expr = self.expr.select(*projected_cols_expr, append=True)
         return SQLGlotIR(expr=select_expr)
 
+    def insert(
+        self,
+        destination: bigquery.TableReference,
+    ) -> str:
+        return sge.insert(self.expr.subquery(), _table(destination)).sql(
+            dialect=self.dialect, pretty=self.pretty
+        )
+
+    def replace(
+        self,
+        destination: bigquery.TableReference,
+    ) -> str:
+        # Workaround for SQLGlot breaking change:
+        # https://github.com/tobymao/sqlglot/pull/4495
+        whens_expr = [
+            sge.When(matched=False, source=True, then=sge.Delete()),
+            sge.When(matched=False, then=sge.Insert(this=sge.Var(this="ROW"))),
+        ]
+        whens_str = "\n".join(
+            when_expr.sql(dialect=self.dialect, pretty=self.pretty)
+            for when_expr in whens_expr
+        )
+
+        merge_str = sge.Merge(
+            this=_table(destination),
+            using=self.expr.subquery(),
+            on=_literal(False, dtypes.BOOL_DTYPE),
+        ).sql(dialect=self.dialect, pretty=self.pretty)
+        return f"{merge_str}\n{whens_str}"
+
     def _encapsulate_as_cte(
         self,
     ) -> sge.Select:
@@ -190,3 +239,11 @@ def _literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
 
 def _cast(arg: typing.Any, to: str) -> sge.Cast:
     return sge.Cast(this=arg, to=to)
+
+
+def _table(table: bigquery.TableReference) -> sge.Table:
+    return sge.Table(
+        this=sg.to_identifier(table.table_id, quoted=True),
+        db=sg.to_identifier(table.dataset_id, quoted=True),
+        catalog=sg.to_identifier(table.project, quoted=True),
+    )

bigframes/session/bq_caching_executor.py

@@ -29,6 +29,7 @@
 
 import bigframes.core
 from bigframes.core import compile, rewrite
+import bigframes.core.compile.sqlglot.sqlglot_ir as sqlglot_ir
 import bigframes.core.guid
 import bigframes.core.nodes as nodes
 import bigframes.core.ordering as order
@@ -206,17 +207,31 @@ def export_gbq(
         if bigframes.options.compute.enable_multi_query_execution:
             self._simplify_with_caching(array_value)
 
-        dispositions = {
-            "fail": bigquery.WriteDisposition.WRITE_EMPTY,
-            "replace": bigquery.WriteDisposition.WRITE_TRUNCATE,
-            "append": bigquery.WriteDisposition.WRITE_APPEND,
-        }
+        table_exists = True
+        try:
+            table = self.bqclient.get_table(destination)
+            if if_exists == "fail":
+                raise ValueError(f"Table already exists: {destination.__str__()}")
+        except google.api_core.exceptions.NotFound:
+            table_exists = False
+
+        # b/409086472: Uses DML for table appends and replacements to avoid
+        # BigQuery `RATE_LIMIT_EXCEEDED` errors, as per quota limits:
+        # https://cloud.google.com/bigquery/quotas#standard_tables
         sql = self.to_sql(array_value, ordered=False)
-        job_config = bigquery.QueryJobConfig(
-            write_disposition=dispositions[if_exists],
-            destination=destination,
-            clustering_fields=cluster_cols if cluster_cols else None,
-        )
+        if not table_exists:
+            job_config = bigquery.QueryJobConfig(
+                destination=destination,
+                clustering_fields=cluster_cols if cluster_cols else None,
+            )
+        else:
+            job_config = bigquery.QueryJobConfig()
+            ir = sqlglot_ir.SQLGlotIR.from_query_string(sql)
+            if if_exists == "append":
+                sql = ir.insert(destination)
+            else:  # for "replace"
+                sql = ir.replace(destination)
+
         # TODO(swast): plumb through the api_name of the user-facing api that
         # caused this query.
         _, query_job = self._run_execute_query(

tests/system/small/test_dataframe_io.py

@@ -515,7 +515,9 @@ def test_to_gbq_if_exists(
     expected_index,
 ):
     """Test the `to_gbq` API with the `if_exists` parameter."""
-    destination_table = f"{dataset_id}.test_to_gbq_if_exists_{if_exists}"
+    destination_table = (
+        f"{dataset_id}.test_to_gbq_if_exists_{if_exists}_{expected_index}"
+    )
 
     scalars_df_default_index.to_gbq(destination_table)
     scalars_df_default_index.to_gbq(destination_table, if_exists=if_exists)