snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/sql/connect/client/__init__.py

@@ -0,0 +1,22 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from pyspark.sql.connect.utils import check_dependencies
+
+check_dependencies(__name__)
+
+from pyspark.sql.connect.client.core import *  # noqa: F401,F403

snowflake/snowpark_connect/includes/python/pyspark/sql/connect/client/artifact.py

@@ -0,0 +1,412 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from pyspark.sql.connect.utils import check_dependencies
+
+check_dependencies(__name__)
+
+import hashlib
+import importlib
+import io
+import sys
+import os
+import zlib
+from itertools import chain
+from typing import List, Iterable, BinaryIO, Iterator, Optional, Tuple
+import abc
+from pathlib import Path
+from urllib.parse import urlparse
+from urllib.request import url2pathname
+from functools import cached_property
+
+import grpc
+
+import pyspark.sql.connect.proto as proto
+import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib
+
+
+JAR_PREFIX: str = "jars"
+PYFILE_PREFIX: str = "pyfiles"
+ARCHIVE_PREFIX: str = "archives"
+FILE_PREFIX: str = "files"
+FORWARD_TO_FS_PREFIX: str = "forward_to_fs"
+CACHE_PREFIX: str = "cache"
+
+
+class LocalData(metaclass=abc.ABCMeta):
+    """
+    Payload stored on this machine.
+    """
+
+    @cached_property
+    @abc.abstractmethod
+    def stream(self) -> BinaryIO:
+        pass
+
+    @cached_property
+    @abc.abstractmethod
+    def size(self) -> int:
+        pass
+
+
+class LocalFile(LocalData):
+    """
+    Payload stored in a local file.
+    """
+
+    def __init__(self, path: str):
+        self.path = path
+        self._size: int
+        self._stream: int
+
+    @cached_property
+    def size(self) -> int:
+        return os.path.getsize(self.path)
+
+    @cached_property
+    def stream(self) -> BinaryIO:
+        return open(self.path, "rb")
+
+
+class InMemory(LocalData):
+    """
+    Payload stored in memory.
+    """
+
+    def __init__(self, blob: bytes):
+        self.blob = blob
+        self._size: int
+        self._stream: int
+
+    @cached_property
+    def size(self) -> int:
+        return len(self.blob)
+
+    @cached_property
+    def stream(self) -> BinaryIO:
+        return io.BytesIO(self.blob)
+
+
+class Artifact:
+    def __init__(self, path: str, storage: LocalData):
+        assert not Path(path).is_absolute(), f"Bad path: {path}"
+        self.path = path
+        self.storage = storage
+
+    @cached_property
+    def size(self) -> int:
+        if isinstance(self.storage, LocalData):
+            return self.storage.size
+        else:
+            raise RuntimeError(f"Unsupported storage {type(self.storage)}")
+
+
+def new_jar_artifact(file_name: str, storage: LocalData) -> Artifact:
+    return _new_artifact(JAR_PREFIX, ".jar", file_name, storage)
+
+
+def new_pyfile_artifact(file_name: str, storage: LocalData) -> Artifact:
+    assert any(file_name.endswith(s) for s in (".py", ".zip", ".egg", ".jar"))
+    return _new_artifact(PYFILE_PREFIX, "", file_name, storage)
+
+
+def new_archive_artifact(file_name: str, storage: LocalData) -> Artifact:
+    return _new_artifact(ARCHIVE_PREFIX, "", file_name, storage)
+
+
+def new_file_artifact(file_name: str, storage: LocalData) -> Artifact:
+    return _new_artifact(FILE_PREFIX, "", file_name, storage)
+
+
+def new_cache_artifact(id: str, storage: LocalData) -> Artifact:
+    return _new_artifact(CACHE_PREFIX, "", id, storage)
+
+
+def _new_artifact(
+    prefix: str, required_suffix: str, file_name: str, storage: LocalData
+) -> Artifact:
+    assert not Path(file_name).is_absolute()
+    assert file_name.endswith(required_suffix)
+    return Artifact(os.path.join(prefix, file_name), storage)
+
+
+class ArtifactManager:
+    """
+    The Artifact Manager is responsible for handling and transferring artifacts from the local
+    client to the server (local/remote).
+
+    Parameters
+    ----------
+    user_id : str, optional
+        User ID.
+    session_id: str
+        An unique identifier of the session which the artifact manager belongs to.
+    channel: grpc.Channel
+        GRPC Channel instance.
+    """
+
+    # Using the midpoint recommendation of 32KiB for chunk size as specified in
+    # https://github.com/grpc/grpc.github.io/issues/371.
+    CHUNK_SIZE: int = 32 * 1024
+
+    def __init__(
+        self,
+        user_id: Optional[str],
+        session_id: str,
+        channel: grpc.Channel,
+        metadata: Iterable[Tuple[str, str]],
+    ):
+        self._user_context = proto.UserContext()
+        if user_id is not None:
+            self._user_context.user_id = user_id
+        self._stub = grpc_lib.SparkConnectServiceStub(channel)
+        self._session_id = session_id
+        self._metadata = metadata
+
+    def _parse_artifacts(
+        self, path_or_uri: str, pyfile: bool, archive: bool, file: bool
+    ) -> List[Artifact]:
+        # Currently only local files with .jar extension is supported.
+        parsed = urlparse(path_or_uri)
+        # Check if it is a file from the scheme
+        if parsed.scheme == "":
+            # Similar with Utils.resolveURI.
+            fragment = parsed.fragment
+            parsed = urlparse(Path(url2pathname(parsed.path)).absolute().as_uri())
+            parsed = parsed._replace(fragment=fragment)
+
+        if parsed.scheme == "file":
+            local_path = url2pathname(parsed.path)
+            name = Path(local_path).name
+            if pyfile and name.endswith(".py"):
+                artifact = new_pyfile_artifact(name, LocalFile(local_path))
+                importlib.invalidate_caches()
+            elif pyfile and (
+                name.endswith(".zip") or name.endswith(".egg") or name.endswith(".jar")
+            ):
+                sys.path.insert(1, local_path)
+                artifact = new_pyfile_artifact(name, LocalFile(local_path))
+                importlib.invalidate_caches()
+            elif archive and (
+                name.endswith(".zip")
+                or name.endswith(".jar")
+                or name.endswith(".tar.gz")
+                or name.endswith(".tgz")
+                or name.endswith(".tar")
+            ):
+                assert any(name.endswith(s) for s in (".zip", ".jar", ".tar.gz", ".tgz", ".tar"))
+
+                if parsed.fragment != "":
+                    # Minimal fix for the workaround of fragment handling in URI.
+                    # This has a limitation - hash(#) in the file name would not work.
+                    if "#" in local_path:
+                        raise ValueError("'#' in the path is not supported for adding an archive.")
+                    name = f"{name}#{parsed.fragment}"
+
+                artifact = new_archive_artifact(name, LocalFile(local_path))
+            elif file:
+                artifact = new_file_artifact(name, LocalFile(local_path))
+            elif name.endswith(".jar"):
+                artifact = new_jar_artifact(name, LocalFile(local_path))
+            else:
+                raise RuntimeError(f"Unsupported file format: {local_path}")
+            return [artifact]
+        raise RuntimeError(f"Unsupported scheme: {parsed.scheme}")
+
+    def _parse_forward_to_fs_artifacts(self, local_path: str, dest_path: str) -> List[Artifact]:
+        abs_path: Path = Path(local_path).absolute()
+        # TODO: Support directory path.
+        assert abs_path.is_file(), "local path must be a file path."
+        storage = LocalFile(str(abs_path))
+
+        assert Path(dest_path).is_absolute(), "destination FS path must be an absolute path."
+
+        # The `dest_path` is an absolute path, to add the FORWARD_TO_FS_PREFIX,
+        # we cannot use `os.path.join`
+        artifact_path = FORWARD_TO_FS_PREFIX + dest_path
+        return [Artifact(artifact_path, storage)]
+
+    def _create_requests(
+        self, *path: str, pyfile: bool, archive: bool, file: bool
+    ) -> Iterator[proto.AddArtifactsRequest]:
+        """Separated for the testing purpose."""
+        return self._add_artifacts(
+            chain(
+                *(self._parse_artifacts(p, pyfile=pyfile, archive=archive, file=file) for p in path)
+            )
+        )
+
+    def _retrieve_responses(
+        self, requests: Iterator[proto.AddArtifactsRequest]
+    ) -> proto.AddArtifactsResponse:
+        """Separated for the testing purpose."""
+        return self._stub.AddArtifacts(requests, metadata=self._metadata)
+
+    def _request_add_artifacts(self, requests: Iterator[proto.AddArtifactsRequest]) -> None:
+        response: proto.AddArtifactsResponse = self._retrieve_responses(requests)
+        summaries: List[proto.AddArtifactsResponse.ArtifactSummary] = []
+
+        for summary in response.artifacts:
+            summaries.append(summary)
+            # TODO(SPARK-42658): Handle responses containing CRC failures.
+
+    def add_artifacts(self, *path: str, pyfile: bool, archive: bool, file: bool) -> None:
+        """
+        Add a single artifact to the session.
+        Currently only local files with .jar extension is supported.
+        """
+        requests: Iterator[proto.AddArtifactsRequest] = self._create_requests(
+            *path, pyfile=pyfile, archive=archive, file=file
+        )
+        self._request_add_artifacts(requests)
+
+    def _add_forward_to_fs_artifacts(self, local_path: str, dest_path: str) -> None:
+        requests: Iterator[proto.AddArtifactsRequest] = self._add_artifacts(
+            self._parse_forward_to_fs_artifacts(local_path, dest_path)
+        )
+        self._request_add_artifacts(requests)
+
+    def _add_artifacts(self, artifacts: Iterable[Artifact]) -> Iterator[proto.AddArtifactsRequest]:
+        """
+        Add a number of artifacts to the session.
+        """
+
+        current_batch: List[Artifact] = []
+        current_batch_size = 0
+
+        def add_to_batch(dep: Artifact, size: int) -> None:
+            nonlocal current_batch
+            nonlocal current_batch_size
+
+            current_batch.append(dep)
+            current_batch_size += size
+
+        def write_batch() -> Iterator[proto.AddArtifactsRequest]:
+            nonlocal current_batch
+            nonlocal current_batch_size
+
+            yield from self._add_batched_artifacts(current_batch)
+            current_batch = []
+            current_batch_size = 0
+
+        for artifact in artifacts:
+            data = artifact.storage
+            size = data.size
+            if size > ArtifactManager.CHUNK_SIZE:
+                # Payload can either be a batch OR a single chunked artifact.
+                # Write batch if non-empty before chunking current artifact.
+                if len(current_batch) > 0:
+                    yield from write_batch()
+                yield from self._add_chunked_artifact(artifact)
+            else:
+                if current_batch_size + size > ArtifactManager.CHUNK_SIZE:
+                    yield from write_batch()
+                add_to_batch(artifact, size)
+
+        if len(current_batch) > 0:
+            yield from write_batch()
+
+    def _add_batched_artifacts(
+        self, artifacts: Iterable[Artifact]
+    ) -> Iterator[proto.AddArtifactsRequest]:
+        """
+        Add a batch of artifacts to the stream. All the artifacts in this call are packaged into a
+        single :class:`proto.AddArtifactsRequest`.
+        """
+        artifact_chunks = []
+
+        for artifact in artifacts:
+            binary = artifact.storage.stream.read()
+            crc32 = zlib.crc32(binary)
+            data = proto.AddArtifactsRequest.ArtifactChunk(data=binary, crc=crc32)
+            artifact_chunks.append(
+                proto.AddArtifactsRequest.SingleChunkArtifact(name=artifact.path, data=data)
+            )
+
+        # Write the request once
+        yield proto.AddArtifactsRequest(
+            session_id=self._session_id,
+            user_context=self._user_context,
+            batch=proto.AddArtifactsRequest.Batch(artifacts=artifact_chunks),
+        )
+
+    def _add_chunked_artifact(self, artifact: Artifact) -> Iterator[proto.AddArtifactsRequest]:
+        """
+        Add a artifact in chunks to the stream. The artifact's data is spread out over multiple
+        :class:`proto.AddArtifactsRequest requests`.
+        """
+        initial_batch = True
+        # Integer division that rounds up to the nearest whole number.
+        get_num_chunks = int(
+            (artifact.size + (ArtifactManager.CHUNK_SIZE - 1)) / ArtifactManager.CHUNK_SIZE
+        )
+
+        # Consume stream in chunks until there is no data left to read.
+        for chunk in iter(lambda: artifact.storage.stream.read(ArtifactManager.CHUNK_SIZE), b""):
+            if initial_batch:
+                # First RPC contains the `BeginChunkedArtifact` payload (`begin_chunk`).
+                yield proto.AddArtifactsRequest(
+                    session_id=self._session_id,
+                    user_context=self._user_context,
+                    begin_chunk=proto.AddArtifactsRequest.BeginChunkedArtifact(
+                        name=artifact.path,
+                        total_bytes=artifact.size,
+                        num_chunks=get_num_chunks,
+                        initial_chunk=proto.AddArtifactsRequest.ArtifactChunk(
+                            data=chunk, crc=zlib.crc32(chunk)
+                        ),
+                    ),
+                )
+                initial_batch = False
+            else:
+                # Subsequent RPCs contains the `ArtifactChunk` payload (`chunk`).
+                yield proto.AddArtifactsRequest(
+                    session_id=self._session_id,
+                    user_context=self._user_context,
+                    chunk=proto.AddArtifactsRequest.ArtifactChunk(
+                        data=chunk, crc=zlib.crc32(chunk)
+                    ),
+                )
+
+    def is_cached_artifact(self, hash: str) -> bool:
+        """
+        Ask the server either any artifact with `hash` has been cached at the server side or not.
+        """
+        artifactName = CACHE_PREFIX + "/" + hash
+        request = proto.ArtifactStatusesRequest(
+            user_context=self._user_context, session_id=self._session_id, names=[artifactName]
+        )
+        resp: proto.ArtifactStatusesResponse = self._stub.ArtifactStatus(
+            request, metadata=self._metadata
+        )
+        status = resp.statuses.get(artifactName)
+        return status.exists if status is not None else False
+
+    def cache_artifact(self, blob: bytes) -> str:
+        """
+        Cache the give blob at the session.
+        """
+        hash = hashlib.sha256(blob).hexdigest()
+        if not self.is_cached_artifact(hash):
+            requests = self._add_artifacts([new_cache_artifact(hash, InMemory(blob))])
+            response: proto.AddArtifactsResponse = self._retrieve_responses(requests)
+            summaries: List[proto.AddArtifactsResponse.ArtifactSummary] = []
+
+            for summary in response.artifacts:
+                summaries.append(summary)
+                # TODO(SPARK-42658): Handle responses containing CRC failures.
+
+        return hash