snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/resource/requests.py

@@ -0,0 +1,539 @@
+#
+# 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 typing import overload, Optional, Dict
+
+from py4j.java_gateway import JavaObject, JVMView
+
+from pyspark.util import _parse_memory
+
+
+class ExecutorResourceRequest:
+    """
+    An Executor resource request. This is used in conjunction with the ResourceProfile to
+    programmatically specify the resources needed for an RDD that will be applied at the
+    stage level.
+
+    This is used to specify what the resource requirements are for an Executor and how
+    Spark can find out specific details about those resources. Not all the parameters are
+    required for every resource type. Resources like GPUs are supported and have same limitations
+    as using the global spark configs spark.executor.resource.gpu.*. The amount, discoveryScript,
+    and vendor parameters for resources are all the same parameters a user would specify through the
+    configs: spark.executor.resource.{resourceName}.{amount, discoveryScript, vendor}.
+
+    For instance, a user wants to allocate an Executor with GPU resources on YARN. The user has
+    to specify the resource name (gpu), the amount or number of GPUs per Executor,
+    the discovery script would be specified so that when the Executor starts up it can
+    discovery what GPU addresses are available for it to use because YARN doesn't tell
+    Spark that, then vendor would not be used because its specific for Kubernetes.
+
+    See the configuration and cluster specific docs for more details.
+
+    Use :class:`pyspark.ExecutorResourceRequests` class as a convenience API.
+
+    .. versionadded:: 3.1.0
+
+    Parameters
+    ----------
+    resourceName : str
+        Name of the resource
+    amount : str
+        Amount requesting
+    discoveryScript : str, optional
+        Optional script used to discover the resources. This is required on some
+        cluster managers that don't tell Spark the addresses of the resources
+        allocated. The script runs on Executors startup to discover the addresses
+        of the resources available.
+    vendor : str, optional
+        Vendor, required for some cluster managers
+
+    See Also
+    --------
+    :class:`pyspark.resource.ResourceProfile`
+
+    Notes
+    -----
+    This API is evolving.
+    """
+
+    def __init__(
+        self,
+        resourceName: str,
+        amount: int,
+        discoveryScript: str = "",
+        vendor: str = "",
+    ):
+        self._name = resourceName
+        self._amount = amount
+        self._discovery_script = discoveryScript
+        self._vendor = vendor
+
+    @property
+    def resourceName(self) -> str:
+        """
+        Returns
+        -------
+        str
+            Name of the resource
+        """
+        return self._name
+
+    @property
+    def amount(self) -> int:
+        """
+        Returns
+        -------
+        str
+            Amount requesting
+        """
+        return self._amount
+
+    @property
+    def discoveryScript(self) -> str:
+        """
+        Returns
+        -------
+        str
+            Amount requesting
+        """
+        return self._discovery_script
+
+    @property
+    def vendor(self) -> str:
+        """
+        Returns
+        -------
+        str
+            Vendor, required for some cluster managers
+        """
+        return self._vendor
+
+
+class ExecutorResourceRequests:
+
+    """
+    A set of Executor resource requests. This is used in conjunction with the
+    :class:`pyspark.resource.ResourceProfileBuilder` to programmatically specify the
+    resources needed for an RDD that will be applied at the stage level.
+
+    .. versionadded:: 3.1.0
+
+    See Also
+    --------
+    :class:`pyspark.resource.ResourceProfile`
+
+    Notes
+    -----
+    This API is evolving.
+    """
+
+    _CORES = "cores"
+    _MEMORY = "memory"
+    _OVERHEAD_MEM = "memoryOverhead"
+    _PYSPARK_MEM = "pyspark.memory"
+    _OFFHEAP_MEM = "offHeap"
+
+    @overload
+    def __init__(self, _jvm: JVMView):
+        ...
+
+    @overload
+    def __init__(
+        self,
+        _jvm: None = ...,
+        _requests: Optional[Dict[str, ExecutorResourceRequest]] = ...,
+    ):
+        ...
+
+    def __init__(
+        self,
+        _jvm: Optional[JVMView] = None,
+        _requests: Optional[Dict[str, ExecutorResourceRequest]] = None,
+    ):
+        from pyspark import SparkContext
+
+        _jvm = _jvm or SparkContext._jvm
+        if _jvm is not None:
+            self._java_executor_resource_requests = (
+                _jvm.org.apache.spark.resource.ExecutorResourceRequests()
+            )
+            if _requests is not None:
+                for k, v in _requests.items():
+                    if k == self._MEMORY:
+                        self._java_executor_resource_requests.memory(str(v.amount))
+                    elif k == self._OVERHEAD_MEM:
+                        self._java_executor_resource_requests.memoryOverhead(str(v.amount))
+                    elif k == self._PYSPARK_MEM:
+                        self._java_executor_resource_requests.pysparkMemory(str(v.amount))
+                    elif k == self._CORES:
+                        self._java_executor_resource_requests.cores(v.amount)
+                    else:
+                        self._java_executor_resource_requests.resource(
+                            v.resourceName, v.amount, v.discoveryScript, v.vendor
+                        )
+        else:
+            self._java_executor_resource_requests = None
+            self._executor_resources: Dict[str, ExecutorResourceRequest] = {}
+
+    def memory(self, amount: str) -> "ExecutorResourceRequests":
+        """
+        Specify heap memory. The value specified will be converted to MiB.
+        This is a convenient API to add :class:`ExecutorResourceRequest` for "memory" resource.
+
+        Parameters
+        ----------
+        amount : str
+            Amount of memory. In the same format as JVM memory strings (e.g. 512m, 2g).
+            Default unit is MiB if not specified.
+
+        Returns
+        -------
+        :class:`ExecutorResourceRequests`
+        """
+        if self._java_executor_resource_requests is not None:
+            self._java_executor_resource_requests.memory(amount)
+        else:
+            self._executor_resources[self._MEMORY] = ExecutorResourceRequest(
+                self._MEMORY, _parse_memory(amount)
+            )
+        return self
+
+    def memoryOverhead(self, amount: str) -> "ExecutorResourceRequests":
+        """
+        Specify overhead memory. The value specified will be converted to MiB.
+        This is a convenient API to add :class:`ExecutorResourceRequest` for "memoryOverhead"
+        resource.
+
+        Parameters
+        ----------
+        amount : str
+            Amount of memory. In the same format as JVM memory strings (e.g. 512m, 2g).
+            Default unit is MiB if not specified.
+
+        Returns
+        -------
+        :class:`ExecutorResourceRequests`
+        """
+        if self._java_executor_resource_requests is not None:
+            self._java_executor_resource_requests.memoryOverhead(amount)
+        else:
+            self._executor_resources[self._OVERHEAD_MEM] = ExecutorResourceRequest(
+                self._OVERHEAD_MEM, _parse_memory(amount)
+            )
+        return self
+
+    def pysparkMemory(self, amount: str) -> "ExecutorResourceRequests":
+        """
+        Specify pyspark memory. The value specified will be converted to MiB.
+        This is a convenient API to add :class:`ExecutorResourceRequest` for "pyspark.memory"
+        resource.
+
+        Parameters
+        ----------
+        amount : str
+            Amount of memory. In the same format as JVM memory strings (e.g. 512m, 2g).
+            Default unit is MiB if not specified.
+
+        Returns
+        -------
+        :class:`ExecutorResourceRequests`
+        """
+        if self._java_executor_resource_requests is not None:
+            self._java_executor_resource_requests.pysparkMemory(amount)
+        else:
+            self._executor_resources[self._PYSPARK_MEM] = ExecutorResourceRequest(
+                self._PYSPARK_MEM, _parse_memory(amount)
+            )
+        return self
+
+    def offheapMemory(self, amount: str) -> "ExecutorResourceRequests":
+        """
+        Specify off heap memory. The value specified will be converted to MiB.
+        This value only take effect when MEMORY_OFFHEAP_ENABLED is true.
+        This is a convenient API to add :class:`ExecutorResourceRequest` for "offHeap"
+        resource.
+
+        Parameters
+        ----------
+        amount : str
+            Amount of memory. In the same format as JVM memory strings (e.g. 512m, 2g).
+            Default unit is MiB if not specified.
+
+        Returns
+        -------
+        :class:`ExecutorResourceRequests`
+        """
+        if self._java_executor_resource_requests is not None:
+            self._java_executor_resource_requests.offHeapMemory(amount)
+        else:
+            self._executor_resources[self._OFFHEAP_MEM] = ExecutorResourceRequest(
+                self._OFFHEAP_MEM, _parse_memory(amount)
+            )
+        return self
+
+    def cores(self, amount: int) -> "ExecutorResourceRequests":
+        """
+        Specify number of cores per Executor.
+        This is a convenient API to add :class:`ExecutorResourceRequest` for "cores" resource.
+
+        Parameters
+        ----------
+        amount : int
+            Number of cores to allocate per Executor.
+
+        Returns
+        -------
+        :class:`ExecutorResourceRequests`
+        """
+        if self._java_executor_resource_requests is not None:
+            self._java_executor_resource_requests.cores(amount)
+        else:
+            self._executor_resources[self._CORES] = ExecutorResourceRequest(self._CORES, amount)
+        return self
+
+    def resource(
+        self,
+        resourceName: str,
+        amount: int,
+        discoveryScript: str = "",
+        vendor: str = "",
+    ) -> "ExecutorResourceRequests":
+        """
+        Amount of a particular custom resource(GPU, FPGA, etc) to use. The resource names supported
+        correspond to the regular Spark configs with the prefix removed. For instance, resources
+        like GPUs are gpu (spark configs `spark.executor.resource.gpu.*`). If you pass in a resource
+        that the cluster manager doesn't support the result is undefined, it may error or may just
+        be ignored.
+        This is a convenient API to add :class:`ExecutorResourceRequest` for custom resources.
+
+        Parameters
+        ----------
+        resourceName : str
+            Name of the resource.
+        amount : str
+            amount of that resource per executor to use.
+        discoveryScript : str, optional
+            Optional script used to discover the resources. This is required on
+            some cluster managers that don't tell Spark the addresses of
+            the resources allocated. The script runs on Executors startup to
+            of the resources available.
+        vendor : str
+            Optional vendor, required for some cluster managers
+
+        Returns
+        -------
+        :class:`ExecutorResourceRequests`
+        """
+        if self._java_executor_resource_requests is not None:
+            self._java_executor_resource_requests.resource(
+                resourceName, amount, discoveryScript, vendor
+            )
+        else:
+            self._executor_resources[resourceName] = ExecutorResourceRequest(
+                resourceName, amount, discoveryScript, vendor
+            )
+        return self
+
+    @property
+    def requests(self) -> Dict[str, ExecutorResourceRequest]:
+        """
+        Returns
+        -------
+        dict
+            Returns all the resource requests for the executor.
+        """
+        if self._java_executor_resource_requests is not None:
+            result = {}
+            execRes = self._java_executor_resource_requests.requestsJMap()
+            for k, v in execRes.items():
+                result[k] = ExecutorResourceRequest(
+                    v.resourceName(), v.amount(), v.discoveryScript(), v.vendor()
+                )
+            return result
+        else:
+            return self._executor_resources
+
+
+class TaskResourceRequest:
+    """
+    A task resource request. This is used in conjunction with the
+    :class:`pyspark.resource.ResourceProfile` to programmatically specify the resources
+    needed for an RDD that will be applied at the stage level. The amount is specified
+    as a float to allow for saying you want more than 1 task per resource. Valid values
+    are less than or equal to 0.5 or whole numbers.
+    Use :class:`pyspark.resource.TaskResourceRequests` class as a convenience API.
+
+    Parameters
+    ----------
+    resourceName : str
+        Name of the resource
+    amount : float
+        Amount requesting as a float to support fractional resource requests.
+        Valid values are less than or equal to 0.5 or whole numbers. This essentially
+        lets you configure X number of tasks to run on a single resource,
+        ie amount equals 0.5 translates into 2 tasks per resource address.
+
+    .. versionadded:: 3.1.0
+
+    See Also
+    --------
+    :class:`pyspark.resource.ResourceProfile`
+
+    Notes
+    -----
+    This API is evolving.
+    """
+
+    def __init__(self, resourceName: str, amount: float):
+        self._name = resourceName
+        self._amount = float(amount)
+
+    @property
+    def resourceName(self) -> str:
+        """
+        Returns
+        -------
+        str
+            Name of the resource.
+        """
+        return self._name
+
+    @property
+    def amount(self) -> float:
+        """
+        Returns
+        -------
+        str
+            Amount requesting as a float to support fractional resource requests.
+        """
+        return self._amount
+
+
+class TaskResourceRequests:
+
+    """
+    A set of task resource requests. This is used in conjunction with the
+    :class:`pyspark.resource.ResourceProfileBuilder` to programmatically specify the resources
+    needed for an RDD that will be applied at the stage level.
+
+    .. versionadded:: 3.1.0
+
+    See Also
+    --------
+    :class:`pyspark.resource.ResourceProfile`
+
+    Notes
+    -----
+    This API is evolving.
+    """
+
+    _CPUS = "cpus"
+
+    @overload
+    def __init__(self, _jvm: JVMView):
+        ...
+
+    @overload
+    def __init__(
+        self,
+        _jvm: None = ...,
+        _requests: Optional[Dict[str, TaskResourceRequest]] = ...,
+    ):
+        ...
+
+    def __init__(
+        self,
+        _jvm: Optional[JVMView] = None,
+        _requests: Optional[Dict[str, TaskResourceRequest]] = None,
+    ):
+        from pyspark import SparkContext
+
+        _jvm = _jvm or SparkContext._jvm
+        if _jvm is not None:
+            self._java_task_resource_requests: Optional[
+                JavaObject
+            ] = _jvm.org.apache.spark.resource.TaskResourceRequests()
+            if _requests is not None:
+                for k, v in _requests.items():
+                    if k == self._CPUS:
+                        self._java_task_resource_requests.cpus(int(v.amount))
+                    else:
+                        self._java_task_resource_requests.resource(v.resourceName, v.amount)
+        else:
+            self._java_task_resource_requests = None
+            self._task_resources: Dict[str, TaskResourceRequest] = {}
+
+    def cpus(self, amount: int) -> "TaskResourceRequests":
+        """
+        Specify number of cpus per Task.
+        This is a convenient API to add :class:`TaskResourceRequest` for cpus.
+
+        Parameters
+        ----------
+        amount : int
+            Number of cpus to allocate per Task.
+
+        Returns
+        -------
+        :class:`TaskResourceRequests`
+        """
+        if self._java_task_resource_requests is not None:
+            self._java_task_resource_requests.cpus(amount)
+        else:
+            self._task_resources[self._CPUS] = TaskResourceRequest(self._CPUS, amount)
+        return self
+
+    def resource(self, resourceName: str, amount: float) -> "TaskResourceRequests":
+        """
+        Amount of a particular custom resource(GPU, FPGA, etc) to use.
+        This is a convenient API to add :class:`TaskResourceRequest` for custom resources.
+
+        Parameters
+        ----------
+        resourceName : str
+            Name of the resource.
+        amount : float
+            Amount requesting as a float to support fractional resource requests.
+            Valid values are less than or equal to 0.5 or whole numbers. This essentially
+            lets you configure X number of tasks to run on a single resource,
+            ie amount equals 0.5 translates into 2 tasks per resource address.
+
+        Returns
+        -------
+        :class:`TaskResourceRequests`
+        """
+        if self._java_task_resource_requests is not None:
+            self._java_task_resource_requests.resource(resourceName, float(amount))
+        else:
+            self._task_resources[resourceName] = TaskResourceRequest(resourceName, amount)
+        return self
+
+    @property
+    def requests(self) -> Dict[str, TaskResourceRequest]:
+        """
+        Returns
+        -------
+        dict
+            Returns all the resource requests for the task.
+        """
+        if self._java_task_resource_requests is not None:
+            result = {}
+            taskRes = self._java_task_resource_requests.requestsJMap()
+            for k, v in taskRes.items():
+                result[k] = TaskResourceRequest(v.resourceName(), v.amount())
+            return result
+        else:
+            return self._task_resources

snowflake/snowpark_connect/includes/python/pyspark/resource/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/resource/tests/test_resources.py

@@ -0,0 +1,83 @@
+#
+# 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 unittest
+
+from pyspark.resource import ExecutorResourceRequests, ResourceProfileBuilder, TaskResourceRequests
+
+
+class ResourceProfileTests(unittest.TestCase):
+    def test_profile_before_sc(self):
+        rpb = ResourceProfileBuilder()
+        ereqs = ExecutorResourceRequests().cores(2).memory("6g").memoryOverhead("1g")
+        ereqs.pysparkMemory("2g").offheapMemory("3g").resource("gpu", 2, "testGpus", "nvidia.com")
+        treqs = TaskResourceRequests().cpus(2).resource("gpu", 2)
+
+        def assert_request_contents(exec_reqs, task_reqs):
+            self.assertEqual(len(exec_reqs), 6)
+            self.assertEqual(exec_reqs["cores"].amount, 2)
+            self.assertEqual(exec_reqs["memory"].amount, 6144)
+            self.assertEqual(exec_reqs["memoryOverhead"].amount, 1024)
+            self.assertEqual(exec_reqs["pyspark.memory"].amount, 2048)
+            self.assertEqual(exec_reqs["offHeap"].amount, 3072)
+            self.assertEqual(exec_reqs["gpu"].amount, 2)
+            self.assertEqual(exec_reqs["gpu"].discoveryScript, "testGpus")
+            self.assertEqual(exec_reqs["gpu"].resourceName, "gpu")
+            self.assertEqual(exec_reqs["gpu"].vendor, "nvidia.com")
+            self.assertEqual(len(task_reqs), 2)
+            self.assertEqual(task_reqs["cpus"].amount, 2.0)
+            self.assertEqual(task_reqs["gpu"].amount, 2.0)
+
+        assert_request_contents(ereqs.requests, treqs.requests)
+        rp = rpb.require(ereqs).require(treqs).build
+        assert_request_contents(rp.executorResources, rp.taskResources)
+        from pyspark import SparkContext, SparkConf
+
+        sc = SparkContext(conf=SparkConf())
+        rdd = sc.parallelize(range(10)).withResources(rp)
+        return_rp = rdd.getResourceProfile()
+        assert_request_contents(return_rp.executorResources, return_rp.taskResources)
+        # intermix objects created before SparkContext init and after
+        rpb2 = ResourceProfileBuilder()
+        # use reqs created before SparkContext with Builder after
+        rpb2.require(ereqs)
+        rpb2.require(treqs)
+        rp2 = rpb2.build
+        self.assertTrue(rp2.id > 0)
+        rdd2 = sc.parallelize(range(10)).withResources(rp2)
+        return_rp2 = rdd2.getResourceProfile()
+        assert_request_contents(return_rp2.executorResources, return_rp2.taskResources)
+        ereqs2 = ExecutorResourceRequests().cores(2).memory("6g").memoryOverhead("1g")
+        ereqs.pysparkMemory("2g").resource("gpu", 2, "testGpus", "nvidia.com")
+        treqs2 = TaskResourceRequests().cpus(2).resource("gpu", 2)
+        # use reqs created after SparkContext with Builder before
+        rpb.require(ereqs2)
+        rpb.require(treqs2)
+        rp3 = rpb.build
+        assert_request_contents(rp3.executorResources, rp3.taskResources)
+        sc.stop()
+
+
+if __name__ == "__main__":
+    from pyspark.resource.tests.test_resources 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/resultiterable.py

@@ -0,0 +1,45 @@
+#
+# 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 typing import TypeVar, TYPE_CHECKING, Iterator, Iterable
+
+if TYPE_CHECKING:
+    from pyspark._typing import SizedIterable
+
+
+__all__ = ["ResultIterable"]
+
+T = TypeVar("T")
+
+
+class ResultIterable(Iterable[T]):
+
+    """
+    A special result iterable. This is used because the standard
+    iterator can not be pickled
+    """
+
+    def __init__(self, data: "SizedIterable[T]"):
+        self.data: "SizedIterable[T]" = data
+        self.index: int = 0
+        self.maxindex: int = len(data)
+
+    def __iter__(self) -> Iterator[T]:
+        return iter(self.data)
+
+    def __len__(self) -> int:
+        return len(self.data)