feat: add allow_large_results option to many I/O methods. Set to False to reduce latency (#1428)

* feat: add allow_large_results option

* add to_arrow

* add the ones that only uses to_pandas()

* add to_csv/json/parquet

* mypy fix

* gcs logic update and execute logic update.

* add to_pandas_batches and to_pandas large test.

* add to_pandas_batches and to_pandas large test.

* add unit tests

* add to_pandas and to_arrow override test

* add to_pandas and to_arrow override test

* add copyright

* modify index to_pandas to match behavior of series and df, add tests.

* update warning message

* update warning message

* update warning message test

* update parameters

* test fix

Author: Genesis929

bigframes/_config/bigquery_options.py

@@ -87,6 +87,7 @@ def __init__(
         kms_key_name: Optional[str] = None,
         skip_bq_connection_check: bool = False,
         *,
+        allow_large_results: bool = True,
         ordering_mode: Literal["strict", "partial"] = "strict",
         client_endpoints_override: Optional[dict] = None,
     ):
@@ -98,6 +99,7 @@ def __init__(
         self._application_name = application_name
         self._kms_key_name = kms_key_name
         self._skip_bq_connection_check = skip_bq_connection_check
+        self._allow_large_results = allow_large_results
         self._session_started = False
         # Determines the ordering strictness for the session.
         self._ordering_mode = _validate_ordering_mode(ordering_mode)
@@ -232,6 +234,26 @@ def skip_bq_connection_check(self, value: bool):
             )
         self._skip_bq_connection_check = value
 
+    @property
+    def allow_large_results(self) -> bool:
+        """
+        Sets the flag to allow or disallow query results larger than 10 GB.
+
+        The default setting for this flag is True, which allows queries to return results
+        exceeding 10 GB by creating an explicit destination table. If set to False, it
+        restricts the result size to 10 GB, and BigQuery will raise an error if this limit
+        is exceeded.
+
+        Returns:
+            bool: True if large results are allowed with an explicit destination table,
+            False if results are limited to 10 GB and errors are raised when exceeded.
+        """
+        return self._allow_large_results
+
+    @allow_large_results.setter
+    def allow_large_results(self, value: bool):
+        self._allow_large_results = value
+
     @property
     def use_regional_endpoints(self) -> bool:
         """Flag to connect to regional API endpoints.

bigframes/core/blocks.py

@@ -112,6 +112,7 @@ class MaterializationOptions:
     downsampling: sampling_options.SamplingOptions = dataclasses.field(
         default_factory=sampling_options.SamplingOptions
     )
+    allow_large_results: Optional[bool] = None
     ordered: bool = True
 
 
@@ -479,9 +480,12 @@ def to_arrow(
         self,
         *,
         ordered: bool = True,
+        allow_large_results: Optional[bool] = None,
     ) -> Tuple[pa.Table, bigquery.QueryJob]:
         """Run query and download results as a pyarrow Table."""
-        execute_result = self.session._executor.execute(self.expr, ordered=ordered)
+        execute_result = self.session._executor.execute(
+            self.expr, ordered=ordered, use_explicit_destination=allow_large_results
+        )
         pa_table = execute_result.to_arrow_table()
 
         pa_index_labels = []
@@ -503,6 +507,7 @@ def to_pandas(
         random_state: Optional[int] = None,
         *,
         ordered: bool = True,
+        allow_large_results: Optional[bool] = None,
     ) -> Tuple[pd.DataFrame, Optional[bigquery.QueryJob]]:
         """Run query and download results as a pandas DataFrame.
 
@@ -545,7 +550,9 @@ def to_pandas(
 
         df, query_job = self._materialize_local(
             materialize_options=MaterializationOptions(
-                downsampling=sampling, ordered=ordered
+                downsampling=sampling,
+                allow_large_results=allow_large_results,
+                ordered=ordered,
             )
         )
         df.set_axis(self.column_labels, axis=1, copy=False)
@@ -563,7 +570,10 @@ def try_peek(
             return None
 
     def to_pandas_batches(
-        self, page_size: Optional[int] = None, max_results: Optional[int] = None
+        self,
+        page_size: Optional[int] = None,
+        max_results: Optional[int] = None,
+        allow_large_results: Optional[bool] = None,
     ):
         """Download results one message at a time.
 
@@ -572,7 +582,7 @@ def to_pandas_batches(
         execute_result = self.session._executor.execute(
             self.expr,
             ordered=True,
-            use_explicit_destination=True,
+            use_explicit_destination=allow_large_results,
             page_size=page_size,
             max_results=max_results,
         )
@@ -601,7 +611,10 @@ def _materialize_local(
         """Run query and download results as a pandas DataFrame. Return the total number of results as well."""
         # TODO(swast): Allow for dry run and timeout.
         execute_result = self.session._executor.execute(
-            self.expr, ordered=materialize_options.ordered, get_size_bytes=True
+            self.expr,
+            ordered=materialize_options.ordered,
+            use_explicit_destination=materialize_options.allow_large_results,
+            get_size_bytes=True,
         )
         assert execute_result.total_bytes is not None
         table_mb = execute_result.total_bytes / _BYTES_TO_MEGABYTES
@@ -1698,7 +1711,7 @@ def transpose(
         original_row_index = (
             original_row_index
             if original_row_index is not None
-            else self.index.to_pandas(ordered=True)
+            else self.index.to_pandas(ordered=True)[0]
         )
         original_row_count = len(original_row_index)
         if original_row_count > bigframes.constants.MAX_COLUMNS:
@@ -2657,14 +2670,22 @@ def column_ids(self) -> Sequence[str]:
     def is_null(self) -> bool:
         return len(self._block._index_columns) == 0
 
-    def to_pandas(self, *, ordered: Optional[bool] = None) -> pd.Index:
+    def to_pandas(
+        self,
+        *,
+        ordered: Optional[bool] = None,
+        allow_large_results: Optional[bool] = None,
+    ) -> Tuple[pd.Index, Optional[bigquery.QueryJob]]:
         """Executes deferred operations and downloads the results."""
         if len(self.column_ids) == 0:
             raise bigframes.exceptions.NullIndexError(
                 "Cannot materialize index, as this object does not have an index. Set index column(s) using set_index."
             )
         ordered = ordered if ordered is not None else True
-        return self._block.select_columns([]).to_pandas(ordered=ordered)[0].index
+        df, query_job = self._block.select_columns([]).to_pandas(
+            ordered=ordered, allow_large_results=allow_large_results
+        )
+        return df.index, query_job
 
     def resolve_level(self, level: LevelsType) -> typing.Sequence[str]:
         if utils.is_list_like(level):

bigframes/core/indexes/base.py

@@ -490,17 +490,28 @@ def __getitem__(self, key: int) -> typing.Any:
         else:
             raise NotImplementedError(f"Index key not supported {key}")
 
-    def to_pandas(self) -> pandas.Index:
+    def to_pandas(self, *, allow_large_results: Optional[bool] = None) -> pandas.Index:
         """Gets the Index as a pandas Index.
 
+        Args:
+            allow_large_results (bool, default None):
+                If not None, overrides the global setting to allow or disallow large query results
+                over the default size limit of 10 GB.
+
         Returns:
             pandas.Index:
                 A pandas Index with all of the labels from this Index.
         """
-        return self._block.index.to_pandas(ordered=True)
+        df, query_job = self._block.index.to_pandas(
+            ordered=True, allow_large_results=allow_large_results
+        )
+        self._query_job = query_job
+        return df
 
-    def to_numpy(self, dtype=None, **kwargs) -> np.ndarray:
-        return self.to_pandas().to_numpy(dtype, **kwargs)
+    def to_numpy(self, dtype=None, *, allow_large_results=None, **kwargs) -> np.ndarray:
+        return self.to_pandas(allow_large_results=allow_large_results).to_numpy(
+            dtype, **kwargs
+        )
 
     __array__ = to_numpy