add config option for default write engine

Author: TrevorBergeron

bigframes/_config/compute_options.py

@@ -15,7 +15,7 @@
 """Options for displaying objects."""
 
 import dataclasses
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Literal, Optional
 
 
 @dataclasses.dataclass
@@ -140,6 +140,17 @@ class ComputeOptions:
         int | None: Number of rows, if set.
     """
 
+    default_write_engine: Literal["bigquery_load", "bigquery_write"] = "bigquery_write"
+    """
+    Sets the default write engine for uploadin local data to bigquery.
+
+    The two options are "bigquery_load" or "bigquery_write". "bigquery_write" is generally
+    preferred as it is faster, but "bigquery_load" may be used if bigquery write api is unavailable.
+
+    Returns:
+        str: "bigquery_load" or "bigquery_write"
+    """
+
     semantic_ops_confirmation_threshold: Optional[int] = 0
     """
     Deprecated.

bigframes/session/bq_caching_executor.py

@@ -598,9 +598,16 @@ def _upload_local_data(self, local_table: local_data.ManagedArrowTable):
         # Might be better as a queue and a worker thread
         with self._upload_lock:
             if local_table not in self.cache._uploaded_local_data:
-                uploaded = self.loader.write_data(
-                    local_table, bigframes.core.guid.generate_guid()
-                )
+                engine = bigframes.options.compute.default_write_engine
+                if engine == "bigquery_load":
+                    uploaded = self.loader.load_data(
+                        local_table, bigframes.core.guid.generate_guid()
+                    )
+                else:
+                    assert engine == "bigquery_write"
+                    uploaded = self.loader.write_data(
+                        local_table, bigframes.core.guid.generate_guid()
+                    )
                 self.cache.cache_remote_replacement(local_table, uploaded)
 
     def _execute_plan_gbq(

tests/system/small/test_large_local_data.py

@@ -23,6 +23,26 @@
 large_dataframe.index = large_dataframe.index.astype("Int64")
 
 
+@pytest.mark.parametrize(
+    ("default_write_engine",),
+    [
+        pytest.param("bigquery_load"),
+        pytest.param("bigquery_write"),
+    ],
+)
+def test_read_pandas_config_default_engine(
+    session: bigframes.Session, default_write_engine
+):
+    pytest.importorskip("pandas", minversion="2.0.0")
+    with bigframes.option_context(
+        "compute.default_write_engine",
+        default_write_engine,
+    ):
+        bf_df = session.read_pandas(large_dataframe)
+
+    assert_frame_equal(large_dataframe, bf_df.to_pandas())
+
+
 def test_read_pandas_defer_noop(session: bigframes.Session):
     pytest.importorskip("pandas", minversion="2.0.0")
     bf_df = session.read_pandas(large_dataframe, write_engine="_deferred")