snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/errors/exceptions/connect.py

@@ -0,0 +1,190 @@
+#
+# 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 json
+from typing import Dict, Optional, TYPE_CHECKING
+
+
+from pyspark.errors.exceptions.base import (
+    AnalysisException as BaseAnalysisException,
+    IllegalArgumentException as BaseIllegalArgumentException,
+    ArithmeticException as BaseArithmeticException,
+    UnsupportedOperationException as BaseUnsupportedOperationException,
+    ArrayIndexOutOfBoundsException as BaseArrayIndexOutOfBoundsException,
+    DateTimeException as BaseDateTimeException,
+    NumberFormatException as BaseNumberFormatException,
+    ParseException as BaseParseException,
+    PySparkException,
+    PythonException as BasePythonException,
+    StreamingQueryException as BaseStreamingQueryException,
+    QueryExecutionException as BaseQueryExecutionException,
+    SparkRuntimeException as BaseSparkRuntimeException,
+    SparkUpgradeException as BaseSparkUpgradeException,
+)
+
+if TYPE_CHECKING:
+    from google.rpc.error_details_pb2 import ErrorInfo
+
+
+class SparkConnectException(PySparkException):
+    """
+    Exception thrown from Spark Connect.
+    """
+
+
+def convert_exception(info: "ErrorInfo", message: str) -> SparkConnectException:
+    classes = []
+    if "classes" in info.metadata:
+        classes = json.loads(info.metadata["classes"])
+
+    if "stackTrace" in info.metadata:
+        stackTrace = info.metadata["stackTrace"]
+        message += f"\n\nJVM stacktrace:\n{stackTrace}"
+
+    if "org.apache.spark.sql.catalyst.parser.ParseException" in classes:
+        return ParseException(message)
+    # Order matters. ParseException inherits AnalysisException.
+    elif "org.apache.spark.sql.AnalysisException" in classes:
+        return AnalysisException(message)
+    elif "org.apache.spark.sql.streaming.StreamingQueryException" in classes:
+        return StreamingQueryException(message)
+    elif "org.apache.spark.sql.execution.QueryExecutionException" in classes:
+        return QueryExecutionException(message)
+    # Order matters. NumberFormatException inherits IllegalArgumentException.
+    elif "java.lang.NumberFormatException" in classes:
+        return NumberFormatException(message)
+    elif "java.lang.IllegalArgumentException" in classes:
+        return IllegalArgumentException(message)
+    elif "java.lang.ArithmeticException" in classes:
+        return ArithmeticException(message)
+    elif "java.lang.UnsupportedOperationException" in classes:
+        return UnsupportedOperationException(message)
+    elif "java.lang.ArrayIndexOutOfBoundsException" in classes:
+        return ArrayIndexOutOfBoundsException(message)
+    elif "java.time.DateTimeException" in classes:
+        return DateTimeException(message)
+    elif "org.apache.spark.SparkRuntimeException" in classes:
+        return SparkRuntimeException(message)
+    elif "org.apache.spark.SparkUpgradeException" in classes:
+        return SparkUpgradeException(message)
+    elif "org.apache.spark.api.python.PythonException" in classes:
+        return PythonException(
+            "\n  An exception was thrown from the Python worker. "
+            "Please see the stack trace below.\n%s" % message
+        )
+    else:
+        return SparkConnectGrpcException(message, reason=info.reason)
+
+
+class SparkConnectGrpcException(SparkConnectException):
+    """
+    Base class to handle the errors from GRPC.
+    """
+
+    def __init__(
+        self,
+        message: Optional[str] = None,
+        error_class: Optional[str] = None,
+        message_parameters: Optional[Dict[str, str]] = None,
+        reason: Optional[str] = None,
+    ) -> None:
+        self.message = message  # type: ignore[assignment]
+        if reason is not None:
+            self.message = f"({reason}) {self.message}"
+
+        super().__init__(
+            message=self.message,
+            error_class=error_class,
+            message_parameters=message_parameters,
+        )
+
+
+class AnalysisException(SparkConnectGrpcException, BaseAnalysisException):
+    """
+    Failed to analyze a SQL query plan, thrown from Spark Connect.
+    """
+
+
+class ParseException(AnalysisException, BaseParseException):
+    """
+    Failed to parse a SQL command, thrown from Spark Connect.
+    """
+
+
+class IllegalArgumentException(SparkConnectGrpcException, BaseIllegalArgumentException):
+    """
+    Passed an illegal or inappropriate argument, thrown from Spark Connect.
+    """
+
+
+class StreamingQueryException(SparkConnectGrpcException, BaseStreamingQueryException):
+    """
+    Exception that stopped a :class:`StreamingQuery` thrown from Spark Connect.
+    """
+
+
+class QueryExecutionException(SparkConnectGrpcException, BaseQueryExecutionException):
+    """
+    Failed to execute a query, thrown from Spark Connect.
+    """
+
+
+class PythonException(SparkConnectGrpcException, BasePythonException):
+    """
+    Exceptions thrown from Spark Connect.
+    """
+
+
+class ArithmeticException(SparkConnectGrpcException, BaseArithmeticException):
+    """
+    Arithmetic exception thrown from Spark Connect.
+    """
+
+
+class UnsupportedOperationException(SparkConnectGrpcException, BaseUnsupportedOperationException):
+    """
+    Unsupported operation exception thrown from Spark Connect.
+    """
+
+
+class ArrayIndexOutOfBoundsException(SparkConnectGrpcException, BaseArrayIndexOutOfBoundsException):
+    """
+    Array index out of bounds exception thrown from Spark Connect.
+    """
+
+
+class DateTimeException(SparkConnectGrpcException, BaseDateTimeException):
+    """
+    Datetime exception thrown from Spark Connect.
+    """
+
+
+class NumberFormatException(IllegalArgumentException, BaseNumberFormatException):
+    """
+    Number format exception thrown from Spark Connect.
+    """
+
+
+class SparkRuntimeException(SparkConnectGrpcException, BaseSparkRuntimeException):
+    """
+    Runtime exception thrown from Spark Connect.
+    """
+
+
+class SparkUpgradeException(SparkConnectGrpcException, BaseSparkUpgradeException):
+    """
+    Exception thrown because of Spark upgrade from Spark Connect.
+    """

snowflake/snowpark_connect/includes/python/pyspark/errors/tests/__init__.py

@@ -0,0 +1,16 @@
+#
+# 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.
+#

snowflake/snowpark_connect/includes/python/pyspark/errors/tests/test_errors.py

@@ -0,0 +1,60 @@
+# -*- encoding: 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 json
+import unittest
+
+from pyspark.errors.error_classes import ERROR_CLASSES_JSON
+from pyspark.errors.utils import ErrorClassesReader
+
+
+class ErrorsTest(unittest.TestCase):
+    def test_error_classes_sorted(self):
+        # Test error classes is sorted alphabetically
+        error_reader = ErrorClassesReader()
+        error_class_names = list(error_reader.error_info_map.keys())
+        for i in range(len(error_class_names) - 1):
+            self.assertTrue(
+                error_class_names[i] < error_class_names[i + 1],
+                f"Error class [{error_class_names[i]}] should place"
+                f"after [{error_class_names[i + 1]}]",
+            )
+
+    def test_error_classes_duplicated(self):
+        # Test error classes is not duplicated
+        def detect_duplication(pairs):
+            error_classes_json = {}
+            for name, message in pairs:
+                self.assertTrue(name not in error_classes_json, f"Duplicate error class: {name}")
+                error_classes_json[name] = message
+            return error_classes_json
+
+        json.loads(ERROR_CLASSES_JSON, object_pairs_hook=detect_duplication)
+
+
+if __name__ == "__main__":
+    import unittest
+    from pyspark.errors.tests.test_errors import *  # noqa: F401
+
+    try:
+        import xmlrunner
+
+        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+    except ImportError:
+        testRunner = None
+    unittest.main(testRunner=testRunner, verbosity=2)

snowflake/snowpark_connect/includes/python/pyspark/errors/utils.py

@@ -0,0 +1,116 @@
+#
+# 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 re
+from typing import Dict
+
+from pyspark.errors.error_classes import ERROR_CLASSES_MAP
+
+
+class ErrorClassesReader:
+    """
+    A reader to load error information from error_classes.py.
+    """
+
+    def __init__(self) -> None:
+        self.error_info_map = ERROR_CLASSES_MAP
+
+    def get_error_message(self, error_class: str, message_parameters: Dict[str, str]) -> str:
+        """
+        Returns the completed error message by applying message parameters to the message template.
+        """
+        message_template = self.get_message_template(error_class)
+        # Verify message parameters.
+        message_parameters_from_template = re.findall("<([a-zA-Z0-9_-]+)>", message_template)
+        assert set(message_parameters_from_template) == set(message_parameters), (
+            f"Undefined error message parameter for error class: {error_class}. "
+            f"Parameters: {message_parameters}"
+        )
+        table = str.maketrans("<>", "{}")
+
+        return message_template.translate(table).format(**message_parameters)
+
+    def get_message_template(self, error_class: str) -> str:
+        """
+        Returns the message template for corresponding error class from error_classes.py.
+
+        For example,
+        when given `error_class` is "EXAMPLE_ERROR_CLASS",
+        and corresponding error class in error_classes.py looks like the below:
+
+        .. code-block:: python
+
+            "EXAMPLE_ERROR_CLASS" : {
+              "message" : [
+                "Problem <A> because of <B>."
+              ]
+            }
+
+        In this case, this function returns:
+        "Problem <A> because of <B>."
+
+        For sub error class, when given `error_class` is "EXAMPLE_ERROR_CLASS.SUB_ERROR_CLASS",
+        and corresponding error class in error_classes.py looks like the below:
+
+        .. code-block:: python
+
+            "EXAMPLE_ERROR_CLASS" : {
+              "message" : [
+                "Problem <A> because of <B>."
+              ],
+              "sub_class" : {
+                "SUB_ERROR_CLASS" : {
+                  "message" : [
+                    "Do <C> to fix the problem."
+                  ]
+                }
+              }
+            }
+
+        In this case, this function returns:
+        "Problem <A> because <B>. Do <C> to fix the problem."
+        """
+        error_classes = error_class.split(".")
+        len_error_classes = len(error_classes)
+        assert len_error_classes in (1, 2)
+
+        # Generate message template for main error class.
+        main_error_class = error_classes[0]
+        if main_error_class in self.error_info_map:
+            main_error_class_info_map = self.error_info_map[main_error_class]
+        else:
+            raise ValueError(f"Cannot find main error class '{main_error_class}'")
+
+        main_message_template = "\n".join(main_error_class_info_map["message"])
+
+        has_sub_class = len_error_classes == 2
+
+        if not has_sub_class:
+            message_template = main_message_template
+        else:
+            # Generate message template for sub error class if exists.
+            sub_error_class = error_classes[1]
+            main_error_class_subclass_info_map = main_error_class_info_map["sub_class"]
+            if sub_error_class in main_error_class_subclass_info_map:
+                sub_error_class_info_map = main_error_class_subclass_info_map[sub_error_class]
+            else:
+                raise ValueError(f"Cannot find sub error class '{sub_error_class}'")
+
+            sub_message_template = "\n".join(sub_error_class_info_map["message"])
+            message_template = main_message_template + " " + sub_message_template
+
+        return message_template

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

@@ -0,0 +1,165 @@
+#
+# 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
+
+
+__all__ = ["SparkFiles"]
+
+from typing import cast, ClassVar, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pyspark import SparkContext
+
+
+class SparkFiles:
+
+    """
+    Resolves paths to files added through :meth:`SparkContext.addFile`.
+
+    SparkFiles contains only classmethods; users should not create SparkFiles
+    instances.
+    """
+
+    _root_directory: ClassVar[Optional[str]] = None
+    _is_running_on_worker: ClassVar[bool] = False
+    _sc: ClassVar[Optional["SparkContext"]] = None
+
+    def __init__(self) -> None:
+        raise NotImplementedError("Do not construct SparkFiles objects")
+
+    @classmethod
+    def get(cls, filename: str) -> str:
+        """
+        Get the absolute path of a file added through
+        :meth:`SparkContext.addFile` or :meth:`SparkContext.addPyFile`.
+
+        .. versionadded:: 0.7.0
+
+        Parameters
+        ----------
+        filename : str
+            file that are added to resources
+
+        Returns
+        -------
+        str
+            the absolute path of the file
+
+        See Also
+        --------
+        :meth:`SparkFiles.getRootDirectory`
+        :meth:`SparkContext.addFile`
+        :meth:`SparkContext.addPyFile`
+        :meth:`SparkContext.listFiles`
+
+        Examples
+        --------
+        >>> import os
+        >>> import tempfile
+        >>> from pyspark import SparkFiles
+
+        >>> with tempfile.TemporaryDirectory() as d:
+        ...     path1 = os.path.join(d, "test.txt")
+        ...     with open(path1, "w") as f:
+        ...         _ = f.write("100")
+        ...
+        ...     sc.addFile(path1)
+        ...     file_list1 = sorted(sc.listFiles)
+        ...
+        ...     def func1(iterator):
+        ...         path = SparkFiles.get("test.txt")
+        ...         assert path.startswith(SparkFiles.getRootDirectory())
+        ...         return [path]
+        ...
+        ...     path_list1 = sc.parallelize([1, 2, 3, 4]).mapPartitions(func1).collect()
+        ...
+        ...     path2 = os.path.join(d, "test.py")
+        ...     with open(path2, "w") as f:
+        ...         _ = f.write("import pyspark")
+        ...
+        ...     # py files
+        ...     sc.addPyFile(path2)
+        ...     file_list2 = sorted(sc.listFiles)
+        ...
+        ...     def func2(iterator):
+        ...         path = SparkFiles.get("test.py")
+        ...         assert path.startswith(SparkFiles.getRootDirectory())
+        ...         return [path]
+        ...
+        ...     path_list2 = sc.parallelize([1, 2, 3, 4]).mapPartitions(func2).collect()
+        >>> file_list1
+        ['file:/.../test.txt']
+        >>> set(path_list1)
+        {'.../test.txt'}
+        >>> file_list2
+        ['file:/.../test.py', 'file:/.../test.txt']
+        >>> set(path_list2)
+        {'.../test.py'}
+        """
+        path = os.path.join(SparkFiles.getRootDirectory(), filename)
+        return os.path.abspath(path)
+
+    @classmethod
+    def getRootDirectory(cls) -> str:
+        """
+        Get the root directory that contains files added through
+        :meth:`SparkContext.addFile` or :meth:`SparkContext.addPyFile`.
+
+        .. versionadded:: 0.7.0
+
+        Returns
+        -------
+        str
+            the root directory that contains files added to resources
+
+        See Also
+        --------
+        :meth:`SparkFiles.get`
+        :meth:`SparkContext.addFile`
+        :meth:`SparkContext.addPyFile`
+
+        Examples
+        --------
+        >>> from pyspark.files import SparkFiles
+        >>> SparkFiles.getRootDirectory()  # doctest: +SKIP
+        '.../spark-a904728e-08d3-400c-a872-cfd82fd6dcd2/userFiles-648cf6d6-bb2c-4f53-82bd-e658aba0c5de'
+        """
+        if cls._is_running_on_worker:
+            return cast(str, cls._root_directory)
+        else:
+            # This will have to change if we support multiple SparkContexts:
+            assert cls._sc is not None
+            assert cls._sc._jvm is not None
+            return cls._sc._jvm.org.apache.spark.SparkFiles.getRootDirectory()
+
+
+def _test() -> None:
+    import doctest
+    import sys
+    from pyspark import SparkContext
+
+    globs = globals().copy()
+    globs["sc"] = SparkContext("local[2]", "files tests")
+    (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS)
+    globs["sc"].stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()

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

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+
+#
+# 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.
+#
+
+# This script attempt to determine the correct setting for SPARK_HOME given
+# that Spark may have been installed on the system with pip.
+
+import os
+import sys
+
+
+def _find_spark_home():
+    """Find the SPARK_HOME."""
+    # If the environment has SPARK_HOME set trust it.
+    if "SPARK_HOME" in os.environ:
+        return os.environ["SPARK_HOME"]
+
+    def is_spark_home(path):
+        """Takes a path and returns true if the provided path could be a reasonable SPARK_HOME"""
+        return os.path.isfile(os.path.join(path, "bin/spark-submit")) and (
+            os.path.isdir(os.path.join(path, "jars"))
+            or os.path.isdir(os.path.join(path, "assembly"))
+        )
+
+    # Spark distribution can be downloaded when PYSPARK_HADOOP_VERSION environment variable is set.
+    # We should look up this directory first, see also SPARK-32017.
+    spark_dist_dir = "spark-distribution"
+    paths = [
+        "../",  # When we're in spark/python.
+    ]
+
+    if "__file__" in globals():
+        paths += [
+            # Two case belows are valid when the current script is called as a library.
+            os.path.join(os.path.dirname(os.path.realpath(__file__)), spark_dist_dir),
+            os.path.dirname(os.path.realpath(__file__)),
+        ]
+
+    # Add the path of the PySpark module if it exists
+    import_error_raised = False
+    from importlib.util import find_spec
+
+    try:
+        module_home = os.path.dirname(find_spec("pyspark").origin)
+        paths.append(os.path.join(module_home, spark_dist_dir))
+        paths.append(module_home)
+        # If we are installed in edit mode also look two dirs up
+        # Downloading different versions are not supported in edit mode.
+        paths.append(os.path.join(module_home, "../../"))
+    except ImportError:
+        # Not pip installed no worries
+        import_error_raised = True
+
+    # Normalize the paths
+    paths = [os.path.abspath(p) for p in paths]
+
+    try:
+        return next(path for path in paths if is_spark_home(path))
+    except StopIteration:
+        print("Could not find valid SPARK_HOME while searching {0}".format(paths), file=sys.stderr)
+        if import_error_raised:
+            print(
+                "\nDid you install PySpark via a package manager such as pip or Conda? If so,\n"
+                "PySpark was not found in your Python environment. It is possible your\n"
+                "Python environment does not properly bind with your package manager.\n"
+                "\nPlease check your default 'python' and if you set PYSPARK_PYTHON and/or\n"
+                "PYSPARK_DRIVER_PYTHON environment variables, and see if you can import\n"
+                "PySpark, for example, 'python -c 'import pyspark'.\n"
+                "\nIf you cannot import, you can install by using the Python executable directly,\n"
+                "for example, 'python -m pip install pyspark [--user]'. Otherwise, you can also\n"
+                "explicitly set the Python executable, that has PySpark installed, to\n"
+                "PYSPARK_PYTHON or PYSPARK_DRIVER_PYTHON environment variables, for example,\n"
+                "'PYSPARK_PYTHON=python3 pyspark'.\n",
+                file=sys.stderr,
+            )
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    print(_find_spark_home())