chore: import upstream snapshot with attribution

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

@@ -0,0 +1,233 @@
+import json
+import logging
+import os
+import time
+from typing import TYPE_CHECKING, List, Optional
+from urllib.parse import urljoin
+
+import numpy as np
+import pyarrow
+import requests
+
+from ray.data._internal.datasource.databricks_credentials import (
+    DatabricksCredentialProvider,
+    build_headers,
+    request_with_401_retry,
+)
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+_STATEMENT_EXEC_POLL_TIME_S = 1
+
+
+@PublicAPI(stability="alpha")
+class DatabricksUCDatasource(Datasource):
+    def __init__(
+        self,
+        warehouse_id: str,
+        catalog: str,
+        schema: str,
+        query: str,
+        credential_provider: DatabricksCredentialProvider,
+    ):
+        self._credential_provider = credential_provider
+
+        # Get host from provider (token is fetched fresh for each request)
+        self.host = self._credential_provider.get_host()
+        self.warehouse_id = warehouse_id
+        self.catalog = catalog
+        self.schema_name = schema
+        self.query = query
+
+        if not self.host.startswith(("http://", "https://")):
+            self.host = f"https://{self.host}"
+
+        url_base = f"{self.host}/api/2.0/sql/statements/"
+
+        payload = json.dumps(
+            {
+                "statement": self.query,
+                "warehouse_id": self.warehouse_id,
+                "wait_timeout": "0s",
+                "disposition": "EXTERNAL_LINKS",
+                "format": "ARROW_STREAM",
+                "catalog": self.catalog,
+                "schema": self.schema_name,
+            }
+        )
+
+        response = request_with_401_retry(
+            requests.post,
+            url_base,
+            self._credential_provider,
+            data=payload,
+        )
+        statement_id = response.json()["statement_id"]
+
+        state = response.json()["status"]["state"]
+
+        logger.info(f"Waiting for query {query!r} execution result.")
+        try:
+            while state in ["PENDING", "RUNNING"]:
+                time.sleep(_STATEMENT_EXEC_POLL_TIME_S)
+                response = request_with_401_retry(
+                    requests.get,
+                    urljoin(url_base, statement_id) + "/",
+                    self._credential_provider,
+                )
+                state = response.json()["status"]["state"]
+        except KeyboardInterrupt:
+            # User cancel the command, so we cancel query execution.
+            requests.post(
+                urljoin(url_base, f"{statement_id}/cancel"),
+                headers=build_headers(self._credential_provider),
+            )
+            try:
+                response.raise_for_status()
+            except Exception as e:
+                logger.warning(
+                    f"Canceling query {query!r} execution failed, reason: {repr(e)}."
+                )
+            raise
+
+        if state != "SUCCEEDED":
+            raise RuntimeError(
+                f"Query {self.query!r} execution failed.\n{response.json()}"
+            )
+
+        manifest = response.json()["manifest"]
+        self.is_truncated = manifest.get("truncated", False)
+
+        if self.is_truncated:
+            logger.warning(
+                f"The resulting size of the dataset of '{query!r}' exceeds "
+                "100GiB and it is truncated."
+            )
+
+        chunks = manifest.get("chunks", [])
+
+        # Make chunks metadata are ordered by index.
+        chunks = sorted(chunks, key=lambda x: x["chunk_index"])
+        num_chunks = len(chunks)
+        self.num_chunks = num_chunks
+        self._estimate_inmemory_data_size = sum(chunk["byte_count"] for chunk in chunks)
+
+        # Capture credential provider (not self) to avoid serializing entire datasource
+        credential_provider_for_tasks = self._credential_provider
+
+        def get_read_task(
+            task_index: int, parallelism: int, per_task_row_limit: Optional[int] = None
+        ):
+            # Handle empty chunk list by yielding an empty PyArrow table
+            if num_chunks == 0:
+                import pyarrow as pa
+
+                metadata = BlockMetadata(
+                    num_rows=0,
+                    size_bytes=0,
+                    input_files=None,
+                    exec_stats=None,
+                )
+
+                def empty_read_fn():
+                    yield pa.Table.from_pydict({})
+
+                return ReadTask(read_fn=empty_read_fn, metadata=metadata)
+
+            # get chunk list to be read in this task and preserve original chunk order
+            chunk_index_list = list(
+                np.array_split(range(num_chunks), parallelism)[task_index]
+            )
+
+            num_rows = sum(
+                chunks[chunk_index]["row_count"] for chunk_index in chunk_index_list
+            )
+            size_bytes = sum(
+                chunks[chunk_index]["byte_count"] for chunk_index in chunk_index_list
+            )
+
+            metadata = BlockMetadata(
+                num_rows=num_rows,
+                size_bytes=size_bytes,
+                input_files=None,
+                exec_stats=None,
+            )
+
+            def _read_fn():
+                for chunk_index in chunk_index_list:
+                    resolve_external_link_url = urljoin(
+                        url_base, f"{statement_id}/result/chunks/{chunk_index}"
+                    )
+
+                    resolve_response = request_with_401_retry(
+                        requests.get,
+                        resolve_external_link_url,
+                        credential_provider_for_tasks,
+                    )
+                    external_url = resolve_response.json()["external_links"][0][
+                        "external_link"
+                    ]
+                    # NOTE: do _NOT_ send the authorization header to external urls
+                    raw_response = requests.get(external_url, auth=None, headers=None)
+                    raw_response.raise_for_status()
+
+                    with pyarrow.ipc.open_stream(raw_response.content) as reader:
+                        arrow_table = reader.read_all()
+
+                    yield arrow_table
+
+            def read_fn():
+                if mock_setup_fn_path := os.environ.get(
+                    "RAY_DATABRICKS_UC_DATASOURCE_READ_FN_MOCK_TEST_SETUP_FN_PATH"
+                ):
+                    import ray.cloudpickle as pickle
+
+                    # This is for testing.
+                    with open(mock_setup_fn_path, "rb") as f:
+                        mock_setup = pickle.load(f)
+                    with mock_setup():
+                        yield from _read_fn()
+                else:
+                    yield from _read_fn()
+
+            return ReadTask(
+                read_fn=read_fn,
+                metadata=metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+
+        self._get_read_task = get_read_task
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return self._estimate_inmemory_data_size
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        # Handle empty dataset case
+        if self.num_chunks == 0:
+            return [self._get_read_task(0, 1, per_task_row_limit)]
+
+        assert parallelism > 0, f"Invalid parallelism {parallelism}"
+
+        if parallelism > self.num_chunks:
+            parallelism = self.num_chunks
+            logger.info(
+                "The parallelism is reduced to chunk number due to "
+                "insufficient chunk parallelism."
+            )
+
+        return [
+            self._get_read_task(index, parallelism, per_task_row_limit)
+            for index in range(parallelism)
+        ]