chore: import upstream snapshot with attribution

python/ray/data/_internal/datasource/bigquery_datasink.py

@@ -0,0 +1,133 @@
+import logging
+import os
+import tempfile
+import time
+import uuid
+from typing import TYPE_CHECKING, Iterable, Optional
+
+import pyarrow.parquet as pq
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+
+import ray
+from ray.data._internal.datasource import bigquery_datasource
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MAX_RETRY_CNT = 10
+RATE_LIMIT_EXCEEDED_SLEEP_TIME = 11
+
+
+class BigQueryDatasink(Datasink[None]):
+    def __init__(
+        self,
+        project_id: str,
+        dataset: str,
+        max_retry_cnt: int = DEFAULT_MAX_RETRY_CNT,
+        overwrite_table: Optional[bool] = True,
+    ) -> None:
+        _check_import(self, module="google.cloud", package="bigquery")
+        _check_import(self, module="google.cloud", package="bigquery_storage")
+        _check_import(self, module="google.api_core", package="exceptions")
+
+        self.project_id = project_id
+        self.dataset = dataset
+        self.max_retry_cnt = max_retry_cnt
+        self.overwrite_table = overwrite_table
+
+    def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
+        from google.api_core import exceptions
+
+        if self.project_id is None or self.dataset is None:
+            raise ValueError("project_id and dataset are required args")
+
+        # Set up datasets to write
+        client = bigquery_datasource._create_client(project_id=self.project_id)
+        dataset_id = self.dataset.split(".", 1)[0]
+        try:
+            client.get_dataset(dataset_id)
+        except exceptions.NotFound:
+            client.create_dataset(f"{self.project_id}.{dataset_id}", timeout=30)
+            logger.info("Created dataset " + dataset_id)
+
+        # Delete table if overwrite_table is True
+        if self.overwrite_table:
+            logger.info(
+                f"Attempting to delete table {self.dataset}"
+                + " if it already exists since kwarg overwrite_table = True."
+            )
+            client.delete_table(f"{self.project_id}.{self.dataset}", not_found_ok=True)
+        else:
+            logger.info(
+                f"The write will append to table {self.dataset}"
+                + " if it already exists since kwarg overwrite_table = False."
+            )
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        def _write_single_block(block: Block, project_id: str, dataset: str) -> None:
+            from google.api_core import exceptions
+            from google.cloud import bigquery
+
+            block = BlockAccessor.for_block(block).to_arrow()
+
+            client = bigquery_datasource._create_client(project_id=project_id)
+            job_config = bigquery.LoadJobConfig(autodetect=True)
+            job_config.source_format = bigquery.SourceFormat.PARQUET
+            job_config.write_disposition = bigquery.WriteDisposition.WRITE_APPEND
+
+            with tempfile.TemporaryDirectory() as temp_dir:
+                fp = os.path.join(temp_dir, f"block_{uuid.uuid4()}.parquet")
+                pq.write_table(block, fp, compression="SNAPPY")
+
+                retry_cnt = 0
+                while retry_cnt <= self.max_retry_cnt:
+                    with open(fp, "rb") as source_file:
+                        job = client.load_table_from_file(
+                            source_file, dataset, job_config=job_config
+                        )
+                    try:
+                        logger.info(job.result())
+                        break
+                    except (exceptions.Forbidden, exceptions.TooManyRequests) as e:
+                        retry_cnt += 1
+                        if retry_cnt > self.max_retry_cnt:
+                            break
+                        logger.info(
+                            "A block write encountered a rate limit exceeded error"
+                            + f" {retry_cnt} time(s). Sleeping to try again."
+                        )
+                        logging.debug(e)
+                        time.sleep(RATE_LIMIT_EXCEEDED_SLEEP_TIME)
+
+                # Raise exception if retry_cnt exceeds max_retry_cnt
+                if retry_cnt > self.max_retry_cnt:
+                    logger.info(
+                        f"Maximum ({self.max_retry_cnt}) retry count exceeded. Ray"
+                        " will attempt to retry the block write via fault tolerance."
+                    )
+                    raise RuntimeError(
+                        f"Write failed due to {retry_cnt}"
+                        " repeated API rate limit exceeded responses. Consider"
+                        " specifying the max_retry_cnt kwarg with a higher value."
+                    )
+
+        _write_single_block = cached_remote_fn(_write_single_block)
+
+        # Launch a remote task for each block within this write task
+        ray.get(
+            [
+                _write_single_block.remote(block, self.project_id, self.dataset)
+                for block in blocks
+                if BlockAccessor.for_block(block).num_rows() > 0
+            ]
+        )