snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/install.py

@@ -0,0 +1,203 @@
+#
+# 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.
+#
+import os
+import re
+import tarfile
+import traceback
+import urllib.request
+from shutil import rmtree
+
+# NOTE that we shouldn't import pyspark here because this is used in
+# setup.py, and assume there's no PySpark imported.
+
+DEFAULT_HADOOP = "hadoop3"
+DEFAULT_HIVE = "hive2.3"
+SUPPORTED_HADOOP_VERSIONS = ["hadoop3", "without-hadoop"]
+SUPPORTED_HIVE_VERSIONS = ["hive2.3"]
+UNSUPPORTED_COMBINATIONS = []  # type: ignore
+
+
+def checked_package_name(spark_version, hadoop_version, hive_version):
+    """
+    Check the generated package name, here we need to use the final hadoop version.
+    """
+    return "%s-bin-%s" % (spark_version, hadoop_version)
+
+
+def checked_versions(spark_version, hadoop_version, hive_version):
+    """
+    Check the valid combinations of supported versions in Spark distributions.
+
+    Parameters
+    ----------
+    spark_version : str
+        Spark version. It should be X.X.X such as '3.0.0' or spark-3.0.0.
+    hadoop_version : str
+        Hadoop version. It should be X such as '2' or 'hadoop2'.
+        'without' and 'without-hadoop' are supported as special keywords for Hadoop free
+        distribution.
+    hive_version : str
+        Hive version. It should be X.X such as '2.3' or 'hive2.3'.
+
+    Parameters
+    ----------
+    tuple
+        fully-qualified versions of Spark, Hadoop and Hive in a tuple.
+        For example, spark-3.2.0, hadoop3 and hive2.3.
+    """
+    if re.match("^[0-9]+\\.[0-9]+\\.[0-9]+$", spark_version):
+        spark_version = "spark-%s" % spark_version
+    if not spark_version.startswith("spark-"):
+        raise RuntimeError(
+            "Spark version should start with 'spark-' prefix; however, " "got %s" % spark_version
+        )
+
+    if hadoop_version == "without":
+        hadoop_version = "without-hadoop"
+    elif re.match("^[0-9]+$", hadoop_version):
+        hadoop_version = "hadoop%s" % hadoop_version
+
+    if hadoop_version not in SUPPORTED_HADOOP_VERSIONS:
+        raise RuntimeError(
+            "Spark distribution of %s is not supported. Hadoop version should be "
+            "one of [%s]" % (hadoop_version, ", ".join(SUPPORTED_HADOOP_VERSIONS))
+        )
+
+    if re.match("^[0-9]+\\.[0-9]+$", hive_version):
+        hive_version = "hive%s" % hive_version
+
+    if hive_version not in SUPPORTED_HIVE_VERSIONS:
+        raise RuntimeError(
+            "Spark distribution of %s is not supported. Hive version should be "
+            "one of [%s]" % (hive_version, ", ".join(SUPPORTED_HADOOP_VERSIONS))
+        )
+
+    return spark_version, convert_old_hadoop_version(spark_version, hadoop_version), hive_version
+
+
+def convert_old_hadoop_version(spark_version, hadoop_version):
+    # check if Spark version <= 3.2, if so, convert hadoop3 to hadoop3.2 and hadoop2 to hadoop2.7
+    version_dict = {
+        "hadoop3": "hadoop3.2",
+        "hadoop2": "hadoop2.7",
+        "without": "without",
+        "without-hadoop": "without-hadoop",
+    }
+    spark_version_parts = re.search("^spark-([0-9]+)\\.([0-9]+)\\.[0-9]+$", spark_version)
+    spark_major_version = int(spark_version_parts.group(1))
+    spark_minor_version = int(spark_version_parts.group(2))
+    if spark_major_version < 3 or (spark_major_version == 3 and spark_minor_version <= 2):
+        hadoop_version = version_dict[hadoop_version]
+    return hadoop_version
+
+
+def install_spark(dest, spark_version, hadoop_version, hive_version):
+    """
+    Installs Spark that corresponds to the given Hadoop version in the current
+    library directory.
+
+    Parameters
+    ----------
+    dest : str
+        The location to download and install the Spark.
+    spark_version : str
+        Spark version. It should be spark-X.X.X form.
+    hadoop_version : str
+        Hadoop version. It should be hadoopX.X
+        such as 'hadoop2.7' or 'without-hadoop'.
+    hive_version : str
+        Hive version. It should be hiveX.X such as 'hive2.3'.
+    """
+
+    package_name = checked_package_name(spark_version, hadoop_version, hive_version)
+    package_local_path = os.path.join(dest, "%s.tgz" % package_name)
+    if "PYSPARK_RELEASE_MIRROR" in os.environ:
+        sites = [os.environ["PYSPARK_RELEASE_MIRROR"]]
+    else:
+        sites = get_preferred_mirrors()
+    print("Trying to download Spark %s from [%s]" % (spark_version, ", ".join(sites)))
+
+    pretty_pkg_name = "%s for Hadoop %s" % (
+        spark_version,
+        "Free build" if hadoop_version == "without" else hadoop_version,
+    )
+
+    for site in sites:
+        os.makedirs(dest, exist_ok=True)
+        url = "%s/spark/%s/%s.tgz" % (site, spark_version, package_name)
+
+        tar = None
+        try:
+            print("Downloading %s from:\n- %s" % (pretty_pkg_name, url))
+            download_to_file(urllib.request.urlopen(url), package_local_path)
+
+            print("Installing to %s" % dest)
+            tar = tarfile.open(package_local_path, "r:gz")
+            for member in tar.getmembers():
+                if member.name == package_name:
+                    # Skip the root directory.
+                    continue
+                member.name = os.path.relpath(member.name, package_name + os.path.sep)
+                tar.extract(member, dest)
+            return
+        except Exception:
+            print("Failed to download %s from %s:" % (pretty_pkg_name, url))
+            traceback.print_exc()
+            rmtree(dest, ignore_errors=True)
+        finally:
+            if tar is not None:
+                tar.close()
+            if os.path.exists(package_local_path):
+                os.remove(package_local_path)
+    raise IOError("Unable to download %s." % pretty_pkg_name)
+
+
+def get_preferred_mirrors():
+    mirror_urls = []
+    for _ in range(3):
+        try:
+            response = urllib.request.urlopen(
+                "https://www.apache.org/dyn/closer.lua?preferred=true"
+            )
+            mirror_urls.append(response.read().decode("utf-8"))
+        except Exception:
+            # If we can't get a mirror URL, skip it. No retry.
+            pass
+
+    default_sites = [
+        "https://dlcdn.apache.org/",
+        "https://archive.apache.org/dist",
+        "https://dist.apache.org/repos/dist/release",
+    ]
+    return list(set(mirror_urls)) + [x for x in default_sites if x not in mirror_urls]
+
+
+def download_to_file(response, path, chunk_size=1024 * 1024):
+    total_size = int(response.info().get("Content-Length").strip())
+    bytes_so_far = 0
+
+    with open(path, mode="wb") as dest:
+        while True:
+            chunk = response.read(chunk_size)
+            bytes_so_far += len(chunk)
+            if not chunk:
+                break
+            dest.write(chunk)
+            print(
+                "Downloaded %d of %d bytes (%0.2f%%)"
+                % (bytes_so_far, total_size, round(float(bytes_so_far) / total_size * 100, 2))
+            )

snowflake/snowpark_connect/includes/python/pyspark/instrumentation_utils.py

@@ -0,0 +1,190 @@
+# -*- coding: utf-8 -*-
+#
+# 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.
+#
+
+import functools
+import inspect
+import threading
+import importlib
+import time
+from types import ModuleType
+from typing import Tuple, Union, List, Callable, Any, Type
+
+
+__all__: List[str] = []
+
+_local = threading.local()
+
+
+def _wrap_function(class_name: str, function_name: str, func: Callable, logger: Any) -> Callable:
+
+    signature = inspect.signature(func)
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        if hasattr(_local, "logging") and _local.logging:
+            # no need to log since this should be internal call.
+            return func(*args, **kwargs)
+        _local.logging = True
+        try:
+            start = time.perf_counter()
+            try:
+                res = func(*args, **kwargs)
+                logger.log_success(
+                    class_name, function_name, time.perf_counter() - start, signature
+                )
+                return res
+            except Exception as ex:
+                logger.log_failure(
+                    class_name, function_name, ex, time.perf_counter() - start, signature
+                )
+                raise
+        finally:
+            _local.logging = False
+
+    return wrapper
+
+
+def _wrap_property(class_name: str, property_name: str, prop: Any, logger: Any) -> Any:
+    @property  # type: ignore[misc]
+    def wrapper(self: Any) -> Any:
+        if hasattr(_local, "logging") and _local.logging:
+            # no need to log since this should be internal call.
+            return prop.fget(self)
+        _local.logging = True
+        try:
+            start = time.perf_counter()
+            try:
+                res = prop.fget(self)
+                logger.log_success(class_name, property_name, time.perf_counter() - start)
+                return res
+            except Exception as ex:
+                logger.log_failure(class_name, property_name, ex, time.perf_counter() - start)
+                raise
+        finally:
+            _local.logging = False
+
+    wrapper.__doc__ = prop.__doc__
+
+    if prop.fset is not None:
+        wrapper = wrapper.setter(  # type: ignore[attr-defined]
+            _wrap_function(class_name, prop.fset.__name__, prop.fset, logger)
+        )
+
+    return wrapper
+
+
+def _wrap_missing_function(
+    class_name: str, function_name: str, func: Callable, original: Any, logger: Any
+) -> Any:
+
+    if not hasattr(original, function_name):
+        return func
+
+    signature = inspect.signature(getattr(original, function_name))
+
+    is_deprecated = func.__name__ == "deprecated_function"
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            return func(*args, **kwargs)
+        finally:
+            logger.log_missing(class_name, function_name, is_deprecated, signature)
+
+    return wrapper
+
+
+def _wrap_missing_property(class_name: str, property_name: str, prop: Any, logger: Any) -> Any:
+
+    is_deprecated = prop.fget.__name__ == "deprecated_property"
+
+    @property  # type: ignore[misc]
+    def wrapper(self: Any) -> Any:
+        try:
+            return prop.fget(self)
+        finally:
+            logger.log_missing(class_name, property_name, is_deprecated)
+
+    return wrapper
+
+
+def _attach(
+    logger_module: Union[str, ModuleType],
+    modules: List[ModuleType],
+    classes: List[Type[Any]],
+    missings: List[Tuple[Type[Any], Type[Any]]],
+) -> None:
+    if isinstance(logger_module, str):
+        logger_module = importlib.import_module(logger_module)
+
+    logger = getattr(logger_module, "get_logger")()
+
+    special_functions = set(
+        [
+            "__init__",
+            "__repr__",
+            "__str__",
+            "_repr_html_",
+            "__len__",
+            "__getitem__",
+            "__setitem__",
+            "__getattr__",
+            "__enter__",
+            "__exit__",
+        ]
+    )
+
+    # Modules
+    for target_module in modules:
+        target_name = target_module.__name__.split(".")[-1]
+        for name in getattr(target_module, "__all__"):
+            func = getattr(target_module, name)
+            if not inspect.isfunction(func):
+                continue
+            setattr(target_module, name, _wrap_function(target_name, name, func, logger))
+
+    # Classes
+    for target_class in classes:
+        for name, func in inspect.getmembers(target_class, inspect.isfunction):
+            if name.startswith("_") and name not in special_functions:
+                continue
+            try:
+                isstatic = isinstance(inspect.getattr_static(target_class, name), staticmethod)
+            except AttributeError:
+                isstatic = False
+            wrapped_function = _wrap_function(target_class.__name__, name, func, logger)
+            setattr(
+                target_class, name, staticmethod(wrapped_function) if isstatic else wrapped_function
+            )
+
+        for name, prop in inspect.getmembers(target_class, lambda o: isinstance(o, property)):
+            if name.startswith("_"):
+                continue
+            setattr(target_class, name, _wrap_property(target_class.__name__, name, prop, logger))
+
+    # Missings
+    for original, missing in missings:
+        for name, func in inspect.getmembers(missing, inspect.isfunction):
+            setattr(
+                missing,
+                name,
+                _wrap_missing_function(original.__name__, name, func, original, logger),
+            )
+
+        for name, prop in inspect.getmembers(missing, lambda o: isinstance(o, property)):
+            setattr(missing, name, _wrap_missing_property(original.__name__, name, prop, logger))

snowflake/snowpark_connect/includes/python/pyspark/java_gateway.py

@@ -0,0 +1,248 @@
+#
+# 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.
+#
+
+import atexit
+import os
+import signal
+import shlex
+import shutil
+import socket
+import platform
+import tempfile
+import time
+from subprocess import Popen, PIPE
+
+from py4j.java_gateway import java_import, JavaGateway, JavaObject, GatewayParameters
+from py4j.clientserver import ClientServer, JavaParameters, PythonParameters
+from pyspark.find_spark_home import _find_spark_home
+from pyspark.serializers import read_int, write_with_length, UTF8Deserializer
+from pyspark.errors import PySparkRuntimeError
+
+
+def launch_gateway(conf=None, popen_kwargs=None):
+    """
+    launch jvm gateway
+
+    Parameters
+    ----------
+    conf : :py:class:`pyspark.SparkConf`
+        spark configuration passed to spark-submit
+    popen_kwargs : dict
+        Dictionary of kwargs to pass to Popen when spawning
+        the py4j JVM. This is a developer feature intended for use in
+        customizing how pyspark interacts with the py4j JVM (e.g., capturing
+        stdout/stderr).
+
+    Returns
+    -------
+    ClientServer or JavaGateway
+    """
+    if "PYSPARK_GATEWAY_PORT" in os.environ:
+        gateway_port = int(os.environ["PYSPARK_GATEWAY_PORT"])
+        gateway_secret = os.environ["PYSPARK_GATEWAY_SECRET"]
+        # Process already exists
+        proc = None
+    else:
+        SPARK_HOME = _find_spark_home()
+        # Launch the Py4j gateway using Spark's run command so that we pick up the
+        # proper classpath and settings from spark-env.sh
+        on_windows = platform.system() == "Windows"
+        script = "./bin/spark-submit.cmd" if on_windows else "./bin/spark-submit"
+        command = [os.path.join(SPARK_HOME, script)]
+        if conf:
+            for k, v in conf.getAll():
+                command += ["--conf", "%s=%s" % (k, v)]
+        submit_args = os.environ.get("PYSPARK_SUBMIT_ARGS", "pyspark-shell")
+        if os.environ.get("SPARK_TESTING"):
+            submit_args = " ".join(["--conf spark.ui.enabled=false", submit_args])
+        command = command + shlex.split(submit_args)
+
+        # Create a temporary directory where the gateway server should write the connection
+        # information.
+        conn_info_dir = tempfile.mkdtemp()
+        try:
+            fd, conn_info_file = tempfile.mkstemp(dir=conn_info_dir)
+            os.close(fd)
+            os.unlink(conn_info_file)
+
+            env = dict(os.environ)
+            env["_PYSPARK_DRIVER_CONN_INFO_PATH"] = conn_info_file
+
+            # Launch the Java gateway.
+            popen_kwargs = {} if popen_kwargs is None else popen_kwargs
+            # We open a pipe to stdin so that the Java gateway can die when the pipe is broken
+            popen_kwargs["stdin"] = PIPE
+            # We always set the necessary environment variables.
+            popen_kwargs["env"] = env
+            if not on_windows:
+                # Don't send ctrl-c / SIGINT to the Java gateway:
+                def preexec_func():
+                    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+                popen_kwargs["preexec_fn"] = preexec_func
+                proc = Popen(command, **popen_kwargs)
+            else:
+                # preexec_fn not supported on Windows
+                proc = Popen(command, **popen_kwargs)
+
+            # Wait for the file to appear, or for the process to exit, whichever happens first.
+            while not proc.poll() and not os.path.isfile(conn_info_file):
+                time.sleep(0.1)
+
+            if not os.path.isfile(conn_info_file):
+                raise PySparkRuntimeError(
+                    error_class="JAVA_GATEWAY_EXITED",
+                    message_parameters={},
+                )
+
+            with open(conn_info_file, "rb") as info:
+                gateway_port = read_int(info)
+                gateway_secret = UTF8Deserializer().loads(info)
+        finally:
+            shutil.rmtree(conn_info_dir)
+
+        # In Windows, ensure the Java child processes do not linger after Python has exited.
+        # In UNIX-based systems, the child process can kill itself on broken pipe (i.e. when
+        # the parent process' stdin sends an EOF). In Windows, however, this is not possible
+        # because java.lang.Process reads directly from the parent process' stdin, contending
+        # with any opportunity to read an EOF from the parent. Note that this is only best
+        # effort and will not take effect if the python process is violently terminated.
+        if on_windows:
+            # In Windows, the child process here is "spark-submit.cmd", not the JVM itself
+            # (because the UNIX "exec" command is not available). This means we cannot simply
+            # call proc.kill(), which kills only the "spark-submit.cmd" process but not the
+            # JVMs. Instead, we use "taskkill" with the tree-kill option "/t" to terminate all
+            # child processes in the tree (http://technet.microsoft.com/en-us/library/bb491009.aspx)
+            def killChild():
+                Popen(["cmd", "/c", "taskkill", "/f", "/t", "/pid", str(proc.pid)])
+
+            atexit.register(killChild)
+
+    # Connect to the gateway (or client server to pin the thread between JVM and Python)
+    if os.environ.get("PYSPARK_PIN_THREAD", "true").lower() == "true":
+        gateway = ClientServer(
+            java_parameters=JavaParameters(
+                port=gateway_port, auth_token=gateway_secret, auto_convert=True
+            ),
+            python_parameters=PythonParameters(port=0, eager_load=False),
+        )
+    else:
+        gateway = JavaGateway(
+            gateway_parameters=GatewayParameters(
+                port=gateway_port, auth_token=gateway_secret, auto_convert=True
+            )
+        )
+
+    # Store a reference to the Popen object for use by the caller (e.g., in reading stdout/stderr)
+    gateway.proc = proc
+
+    # Import the classes used by PySpark
+    java_import(gateway.jvm, "org.apache.spark.SparkConf")
+    java_import(gateway.jvm, "org.apache.spark.api.java.*")
+    java_import(gateway.jvm, "org.apache.spark.api.python.*")
+    java_import(gateway.jvm, "org.apache.spark.ml.python.*")
+    java_import(gateway.jvm, "org.apache.spark.mllib.api.python.*")
+    java_import(gateway.jvm, "org.apache.spark.resource.*")
+    # TODO(davies): move into sql
+    java_import(gateway.jvm, "org.apache.spark.sql.*")
+    java_import(gateway.jvm, "org.apache.spark.sql.api.python.*")
+    java_import(gateway.jvm, "org.apache.spark.sql.hive.*")
+    java_import(gateway.jvm, "scala.Tuple2")
+
+    return gateway
+
+
+def _do_server_auth(conn, auth_secret):
+    """
+    Performs the authentication protocol defined by the SocketAuthHelper class on the given
+    file-like object 'conn'.
+    """
+    write_with_length(auth_secret.encode("utf-8"), conn)
+    conn.flush()
+    reply = UTF8Deserializer().loads(conn)
+    if reply != "ok":
+        conn.close()
+        raise PySparkRuntimeError(
+            error_class="UNEXPECTED_RESPONSE_FROM_SERVER",
+            message_parameters={},
+        )
+
+
+def local_connect_and_auth(port, auth_secret):
+    """
+    Connect to local host, authenticate with it, and return a (sockfile,sock) for that connection.
+    Handles IPV4 & IPV6, does some error handling.
+
+    Parameters
+    ----------
+    port : str or int or None
+    auth_secret : str
+
+    Returns
+    -------
+    tuple
+        with (sockfile, sock)
+    """
+    sock = None
+    errors = []
+    # Support for both IPv4 and IPv6.
+    addr = "127.0.0.1"
+    if os.environ.get("SPARK_PREFER_IPV6", "false").lower() == "true":
+        addr = "::1"
+    for res in socket.getaddrinfo(addr, port, socket.AF_UNSPEC, socket.SOCK_STREAM):
+        af, socktype, proto, _, sa = res
+        try:
+            sock = socket.socket(af, socktype, proto)
+            sock.settimeout(int(os.environ.get("SPARK_AUTH_SOCKET_TIMEOUT", 15)))
+            sock.connect(sa)
+            sockfile = sock.makefile("rwb", int(os.environ.get("SPARK_BUFFER_SIZE", 65536)))
+            _do_server_auth(sockfile, auth_secret)
+            return (sockfile, sock)
+        except socket.error as e:
+            emsg = str(e)
+            errors.append("tried to connect to %s, but an error occurred: %s" % (sa, emsg))
+            sock.close()
+            sock = None
+    raise PySparkRuntimeError(
+        error_class="CANNOT_OPEN_SOCKET",
+        message_parameters={
+            "errors": str(errors),
+        },
+    )
+
+
+def ensure_callback_server_started(gw):
+    """
+    Start callback server if not already started. The callback server is needed if the Java
+    driver process needs to callback into the Python driver process to execute Python code.
+    """
+
+    # getattr will fallback to JVM, so we cannot test by hasattr()
+    if "_callback_server" not in gw.__dict__ or gw._callback_server is None:
+        gw.callback_server_parameters.eager_load = True
+        gw.callback_server_parameters.daemonize = True
+        gw.callback_server_parameters.daemonize_connections = True
+        gw.callback_server_parameters.port = 0
+        gw.start_callback_server(gw.callback_server_parameters)
+        cbport = gw._callback_server.server_socket.getsockname()[1]
+        gw._callback_server.port = cbport
+        # gateway with real port
+        gw._python_proxy_port = gw._callback_server.port
+        # get the GatewayServer object in JVM by ID
+        jgws = JavaObject("GATEWAY_SERVER", gw._gateway_client)
+        # update the port of CallbackClient with real port
+        jgws.resetCallbackClient(jgws.getCallbackClient().getAddress(), gw._python_proxy_port)