This is an automated email from the ASF dual-hosted git repository.
zero323 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 6ff760d [SPARK-37154][PYTHON] Inline hints for pyspark.rdd
6ff760d is described below
commit 6ff760d483124b121d79c3a2d5fdc3ee3f27dd00
Author: zero323 <[email protected]>
AuthorDate: Sat Feb 19 11:32:13 2022 +0100
[SPARK-37154][PYTHON] Inline hints for pyspark.rdd
### What changes were proposed in this pull request?
This PR proposes migration of type hints for `pyspark.rdd` from stub file
to inline annotation.
### Why are the changes needed?
As a part of ongoing process of migration of stubs to inline hints.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests + new data tests.
Closes #35252 from zero323/SPARK-37154.
Authored-by: zero323 <[email protected]>
Signed-off-by: zero323 <[email protected]>
---
python/pyspark/_typing.pyi | 14 +-
python/pyspark/conf.py | 12 +
python/pyspark/context.py | 3 +-
python/pyspark/rdd.py | 987 ++++++++++++++++++-------
python/pyspark/rdd.pyi | 481 ------------
python/pyspark/serializers.py | 7 +
python/pyspark/sql/_typing.pyi | 4 +-
python/pyspark/sql/pandas/_typing/__init__.pyi | 5 +-
python/pyspark/tests/typing/test_rdd.yml | 77 +-
9 files changed, 814 insertions(+), 776 deletions(-)
diff --git a/python/pyspark/_typing.pyi b/python/pyspark/_typing.pyi
index 9a36c89..6cc0926 100644
--- a/python/pyspark/_typing.pyi
+++ b/python/pyspark/_typing.pyi
@@ -17,17 +17,27 @@
# under the License.
from typing import Callable, Iterable, Sized, TypeVar, Union
-from typing_extensions import Protocol
+from typing_extensions import Literal, Protocol
+
+from numpy import int32, int64, float32, float64, ndarray
F = TypeVar("F", bound=Callable)
T_co = TypeVar("T_co", covariant=True)
PrimitiveType = Union[bool, float, int, str]
+NonUDFType = Literal[0]
+
class SupportsIAdd(Protocol):
def __iadd__(self, other: SupportsIAdd) -> SupportsIAdd: ...
class SupportsOrdering(Protocol):
- def __le__(self, other: SupportsOrdering) -> bool: ...
+ def __lt__(self, other: SupportsOrdering) -> bool: ...
class SizedIterable(Protocol, Sized, Iterable[T_co]): ...
+
+S = TypeVar("S", bound=SupportsOrdering)
+
+NumberOrArray = TypeVar(
+ "NumberOrArray", float, int, complex, int32, int64, float32, float64,
ndarray
+)
diff --git a/python/pyspark/conf.py b/python/pyspark/conf.py
index 536e1f8..a9e26966 100644
--- a/python/pyspark/conf.py
+++ b/python/pyspark/conf.py
@@ -203,6 +203,18 @@ class SparkConf:
self.set(k, v)
return self
+ @overload
+ def get(self, key: str) -> Optional[str]:
+ ...
+
+ @overload
+ def get(self, key: str, defaultValue: None) -> Optional[str]:
+ ...
+
+ @overload
+ def get(self, key: str, defaultValue: str) -> str:
+ ...
+
def get(self, key: str, defaultValue: Optional[str] = None) ->
Optional[str]:
"""Get the configured value for some key, or return a default
otherwise."""
if defaultValue is None: # Py4J doesn't call the right get() if we
pass None
diff --git a/python/pyspark/context.py b/python/pyspark/context.py
index 3db9630..68f748e 100644
--- a/python/pyspark/context.py
+++ b/python/pyspark/context.py
@@ -35,6 +35,7 @@ from typing import (
List,
NoReturn,
Optional,
+ Sequence,
Tuple,
Type,
TYPE_CHECKING,
@@ -1421,7 +1422,7 @@ class SparkContext:
self,
rdd: "RDD[T]",
partitionFunc: Callable[[Iterable[T]], Iterable[U]],
- partitions: Optional[List[int]] = None,
+ partitions: Optional[Sequence[int]] = None,
allowLocal: bool = False,
) -> List[U]:
"""
diff --git a/python/pyspark/rdd.py b/python/pyspark/rdd.py
index 97b87ea..7cb887f 100644
--- a/python/pyspark/rdd.py
+++ b/python/pyspark/rdd.py
@@ -30,6 +30,26 @@ from collections import defaultdict
from itertools import chain
from functools import reduce
from math import sqrt, log, isinf, isnan, pow, ceil
+from typing import (
+ Any,
+ Callable,
+ Dict,
+ Generic,
+ Hashable,
+ Iterable,
+ Iterator,
+ IO,
+ List,
+ NoReturn,
+ Optional,
+ Sequence,
+ Tuple,
+ Union,
+ TypeVar,
+ cast,
+ overload,
+ TYPE_CHECKING,
+)
from pyspark.java_gateway import local_connect_and_auth
from pyspark.serializers import (
@@ -40,6 +60,7 @@ from pyspark.serializers import (
CloudPickleSerializer,
PairDeserializer,
CPickleSerializer,
+ Serializer,
pack_long,
read_int,
write_int,
@@ -67,6 +88,41 @@ from pyspark.shuffle import (
from pyspark.traceback_utils import SCCallSiteSync
from pyspark.util import fail_on_stopiteration, _parse_memory
+
+if TYPE_CHECKING:
+ import socket
+ import io
+
+ from pyspark._typing import NonUDFType
+ from pyspark._typing import S, NumberOrArray
+ from pyspark.context import SparkContext
+ from pyspark.sql.pandas._typing import (
+ PandasScalarUDFType,
+ PandasGroupedMapUDFType,
+ PandasGroupedAggUDFType,
+ PandasWindowAggUDFType,
+ PandasScalarIterUDFType,
+ PandasMapIterUDFType,
+ PandasCogroupedMapUDFType,
+ ArrowMapIterUDFType,
+ )
+ from pyspark.sql.dataframe import DataFrame
+ from pyspark.sql.types import AtomicType, StructType
+ from pyspark.sql._typing import AtomicValue, RowLike, SQLBatchedUDFType
+
+ from py4j.java_gateway import JavaObject # type: ignore[import]
+ from py4j.java_collections import JavaArray # type: ignore[import]
+
+T = TypeVar("T")
+T_co = TypeVar("T_co", covariant=True)
+U = TypeVar("U")
+K = TypeVar("K", bound=Hashable)
+V = TypeVar("V")
+V1 = TypeVar("V1")
+V2 = TypeVar("V2")
+V3 = TypeVar("V3")
+
+
__all__ = ["RDD"]
@@ -79,21 +135,21 @@ class PythonEvalType:
These values should match values in
org.apache.spark.api.python.PythonEvalType.
"""
- NON_UDF = 0
+ NON_UDF: "NonUDFType" = 0
- SQL_BATCHED_UDF = 100
+ SQL_BATCHED_UDF: "SQLBatchedUDFType" = 100
- SQL_SCALAR_PANDAS_UDF = 200
- SQL_GROUPED_MAP_PANDAS_UDF = 201
- SQL_GROUPED_AGG_PANDAS_UDF = 202
- SQL_WINDOW_AGG_PANDAS_UDF = 203
- SQL_SCALAR_PANDAS_ITER_UDF = 204
- SQL_MAP_PANDAS_ITER_UDF = 205
- SQL_COGROUPED_MAP_PANDAS_UDF = 206
- SQL_MAP_ARROW_ITER_UDF = 207
+ SQL_SCALAR_PANDAS_UDF: "PandasScalarUDFType" = 200
+ SQL_GROUPED_MAP_PANDAS_UDF: "PandasGroupedMapUDFType" = 201
+ SQL_GROUPED_AGG_PANDAS_UDF: "PandasGroupedAggUDFType" = 202
+ SQL_WINDOW_AGG_PANDAS_UDF: "PandasWindowAggUDFType" = 203
+ SQL_SCALAR_PANDAS_ITER_UDF: "PandasScalarIterUDFType" = 204
+ SQL_MAP_PANDAS_ITER_UDF: "PandasMapIterUDFType" = 205
+ SQL_COGROUPED_MAP_PANDAS_UDF: "PandasCogroupedMapUDFType" = 206
+ SQL_MAP_ARROW_ITER_UDF: "ArrowMapIterUDFType" = 207
-def portable_hash(x):
+def portable_hash(x: Hashable) -> int:
"""
This function returns consistent hash code for builtin types, especially
for None and tuple with None.
@@ -137,7 +193,11 @@ class BoundedFloat(float):
100.0
"""
- def __new__(cls, mean, confidence, low, high):
+ confidence: float
+ low: float
+ high: float
+
+ def __new__(cls, mean: float, confidence: float, low: float, high: float)
-> "BoundedFloat":
obj = float.__new__(cls, mean)
obj.confidence = confidence
obj.low = low
@@ -145,7 +205,7 @@ class BoundedFloat(float):
return obj
-def _create_local_socket(sock_info):
+def _create_local_socket(sock_info: "JavaArray") -> "io.BufferedRWPair":
"""
Create a local socket that can be used to load deserialized data from the
JVM
@@ -158,8 +218,10 @@ def _create_local_socket(sock_info):
-------
sockfile file descriptor of the local socket
"""
- port = sock_info[0]
- auth_secret = sock_info[1]
+ sockfile: "io.BufferedRWPair"
+ sock: "socket.socket"
+ port: int = sock_info[0]
+ auth_secret: str = sock_info[1]
sockfile, sock = local_connect_and_auth(port, auth_secret)
# The RDD materialization time is unpredictable, if we set a timeout for
socket reading
# operation, it will very possibly fail. See SPARK-18281.
@@ -167,7 +229,7 @@ def _create_local_socket(sock_info):
return sockfile
-def _load_from_socket(sock_info, serializer):
+def _load_from_socket(sock_info: "JavaArray", serializer: Serializer) ->
Iterator[Any]:
"""
Connect to a local socket described by sock_info and use the given
serializer to yield data
@@ -188,18 +250,21 @@ def _load_from_socket(sock_info, serializer):
return serializer.load_stream(sockfile)
-def _local_iterator_from_socket(sock_info, serializer):
+def _local_iterator_from_socket(sock_info: "JavaArray", serializer:
Serializer) -> Iterator[Any]:
class PyLocalIterable:
"""Create a synchronous local iterable over a socket"""
- def __init__(self, _sock_info, _serializer):
+ def __init__(self, _sock_info: "JavaArray", _serializer: Serializer):
+ port: int
+ auth_secret: str
+ jsocket_auth_server: "JavaObject"
port, auth_secret, self.jsocket_auth_server = _sock_info
self._sockfile = _create_local_socket((port, auth_secret))
self._serializer = _serializer
- self._read_iter = iter([]) # Initialize as empty iterator
+ self._read_iter: Iterator[Any] = iter([]) # Initialize as empty
iterator
self._read_status = 1
- def __iter__(self):
+ def __iter__(self) -> Iterator[Any]:
while self._read_status == 1:
# Request next partition data from Java
write_int(1, self._sockfile)
@@ -218,7 +283,7 @@ def _local_iterator_from_socket(sock_info, serializer):
elif self._read_status == -1:
self.jsocket_auth_server.getResult()
- def __del__(self):
+ def __del__(self) -> None:
# If local iterator is not fully consumed,
if self._read_status == 1:
try:
@@ -236,22 +301,22 @@ def _local_iterator_from_socket(sock_info, serializer):
class Partitioner:
- def __init__(self, numPartitions, partitionFunc):
+ def __init__(self, numPartitions: int, partitionFunc: Callable[[Any],
int]):
self.numPartitions = numPartitions
self.partitionFunc = partitionFunc
- def __eq__(self, other):
+ def __eq__(self, other: Any) -> bool:
return (
isinstance(other, Partitioner)
and self.numPartitions == other.numPartitions
and self.partitionFunc == other.partitionFunc
)
- def __call__(self, k):
+ def __call__(self, k: Any) -> int:
return self.partitionFunc(k) % self.numPartitions
-class RDD:
+class RDD(Generic[T_co]):
"""
A Resilient Distributed Dataset (RDD), the basic abstraction in Spark.
@@ -259,7 +324,12 @@ class RDD:
operated on in parallel.
"""
- def __init__(self, jrdd, ctx,
jrdd_deserializer=AutoBatchedSerializer(CPickleSerializer())):
+ def __init__(
+ self,
+ jrdd: "JavaObject",
+ ctx: "SparkContext",
+ jrdd_deserializer: Serializer =
AutoBatchedSerializer(CPickleSerializer()),
+ ):
self._jrdd = jrdd
self.is_cached = False
self.is_checkpointed = False
@@ -267,21 +337,21 @@ class RDD:
self.ctx = ctx
self._jrdd_deserializer = jrdd_deserializer
self._id = jrdd.id()
- self.partitioner = None
+ self.partitioner: Optional[Partitioner] = None
- def _pickled(self):
+ def _pickled(self: "RDD[T]") -> "RDD[T]":
return self._reserialize(AutoBatchedSerializer(CPickleSerializer()))
- def id(self):
+ def id(self) -> int:
"""
A unique ID for this RDD (within its SparkContext).
"""
return self._id
- def __repr__(self):
+ def __repr__(self) -> str:
return self._jrdd.toString()
- def __getnewargs__(self):
+ def __getnewargs__(self) -> NoReturn:
# This method is called when attempting to pickle an RDD, which is
always an error:
raise RuntimeError(
"It appears that you are attempting to broadcast an RDD or
reference an RDD from an "
@@ -293,13 +363,13 @@ class RDD:
)
@property
- def context(self):
+ def context(self) -> "SparkContext":
"""
The :class:`SparkContext` that this RDD was created on.
"""
return self.ctx
- def cache(self):
+ def cache(self: "RDD[T]") -> "RDD[T]":
"""
Persist this RDD with the default storage level (`MEMORY_ONLY`).
"""
@@ -307,7 +377,7 @@ class RDD:
self.persist(StorageLevel.MEMORY_ONLY)
return self
- def persist(self, storageLevel=StorageLevel.MEMORY_ONLY):
+ def persist(self: "RDD[T]", storageLevel: StorageLevel =
StorageLevel.MEMORY_ONLY) -> "RDD[T]":
"""
Set this RDD's storage level to persist its values across operations
after the first time it is computed. This can only be used to assign
@@ -325,7 +395,7 @@ class RDD:
self._jrdd.persist(javaStorageLevel)
return self
- def unpersist(self, blocking=False):
+ def unpersist(self: "RDD[T]", blocking: bool = False) -> "RDD[T]":
"""
Mark the RDD as non-persistent, and remove all blocks for it from
memory and disk.
@@ -338,7 +408,7 @@ class RDD:
self._jrdd.unpersist(blocking)
return self
- def checkpoint(self):
+ def checkpoint(self) -> None:
"""
Mark this RDD for checkpointing. It will be saved to a file inside the
checkpoint directory set with :meth:`SparkContext.setCheckpointDir` and
@@ -350,13 +420,13 @@ class RDD:
self.is_checkpointed = True
self._jrdd.rdd().checkpoint()
- def isCheckpointed(self):
+ def isCheckpointed(self) -> bool:
"""
Return whether this RDD is checkpointed and materialized, either
reliably or locally.
"""
return self._jrdd.rdd().isCheckpointed()
- def localCheckpoint(self):
+ def localCheckpoint(self) -> None:
"""
Mark this RDD for local checkpointing using Spark's existing caching
layer.
@@ -377,7 +447,7 @@ class RDD:
"""
self._jrdd.rdd().localCheckpoint()
- def isLocallyCheckpointed(self):
+ def isLocallyCheckpointed(self) -> bool:
"""
Return whether this RDD is marked for local checkpointing.
@@ -385,17 +455,17 @@ class RDD:
"""
return self._jrdd.rdd().isLocallyCheckpointed()
- def getCheckpointFile(self):
+ def getCheckpointFile(self) -> Optional[str]:
"""
Gets the name of the file to which this RDD was checkpointed
Not defined if RDD is checkpointed locally.
"""
checkpointFile = self._jrdd.rdd().getCheckpointFile()
- if checkpointFile.isDefined():
- return checkpointFile.get()
- def map(self, f, preservesPartitioning=False):
+ return checkpointFile.get() if checkpointFile.isDefined() else None
+
+ def map(self: "RDD[T]", f: Callable[[T], U], preservesPartitioning: bool =
False) -> "RDD[U]":
"""
Return a new RDD by applying a function to each element of this RDD.
@@ -406,12 +476,14 @@ class RDD:
[('a', 1), ('b', 1), ('c', 1)]
"""
- def func(_, iterator):
+ def func(_: int, iterator: Iterable[T]) -> Iterable[U]:
return map(fail_on_stopiteration(f), iterator)
return self.mapPartitionsWithIndex(func, preservesPartitioning)
- def flatMap(self, f, preservesPartitioning=False):
+ def flatMap(
+ self: "RDD[T]", f: Callable[[T], Iterable[U]], preservesPartitioning:
bool = False
+ ) -> "RDD[U]":
"""
Return a new RDD by first applying a function to all elements of this
RDD, and then flattening the results.
@@ -425,12 +497,14 @@ class RDD:
[(2, 2), (2, 2), (3, 3), (3, 3), (4, 4), (4, 4)]
"""
- def func(s, iterator):
+ def func(_: int, iterator: Iterable[T]) -> Iterable[U]:
return chain.from_iterable(map(fail_on_stopiteration(f), iterator))
return self.mapPartitionsWithIndex(func, preservesPartitioning)
- def mapPartitions(self, f, preservesPartitioning=False):
+ def mapPartitions(
+ self: "RDD[T]", f: Callable[[Iterable[T]], Iterable[U]],
preservesPartitioning: bool = False
+ ) -> "RDD[U]":
"""
Return a new RDD by applying a function to each partition of this RDD.
@@ -442,12 +516,16 @@ class RDD:
[3, 7]
"""
- def func(s, iterator):
+ def func(_: int, iterator: Iterable[T]) -> Iterable[U]:
return f(iterator)
return self.mapPartitionsWithIndex(func, preservesPartitioning)
- def mapPartitionsWithIndex(self, f, preservesPartitioning=False):
+ def mapPartitionsWithIndex(
+ self: "RDD[T]",
+ f: Callable[[int, Iterable[T]], Iterable[U]],
+ preservesPartitioning: bool = False,
+ ) -> "RDD[U]":
"""
Return a new RDD by applying a function to each partition of this RDD,
while tracking the index of the original partition.
@@ -461,7 +539,11 @@ class RDD:
"""
return PipelinedRDD(self, f, preservesPartitioning)
- def mapPartitionsWithSplit(self, f, preservesPartitioning=False):
+ def mapPartitionsWithSplit(
+ self: "RDD[T]",
+ f: Callable[[int, Iterable[T]], Iterable[U]],
+ preservesPartitioning: bool = False,
+ ) -> "RDD[U]":
"""
Return a new RDD by applying a function to each partition of this RDD,
@@ -484,7 +566,7 @@ class RDD:
)
return self.mapPartitionsWithIndex(f, preservesPartitioning)
- def getNumPartitions(self):
+ def getNumPartitions(self) -> int:
"""
Returns the number of partitions in RDD
@@ -496,7 +578,7 @@ class RDD:
"""
return self._jrdd.partitions().size()
- def filter(self, f):
+ def filter(self: "RDD[T]", f: Callable[[T], bool]) -> "RDD[T]":
"""
Return a new RDD containing only the elements that satisfy a predicate.
@@ -507,12 +589,12 @@ class RDD:
[2, 4]
"""
- def func(iterator):
+ def func(iterator: Iterable[T]) -> Iterable[T]:
return filter(fail_on_stopiteration(f), iterator)
return self.mapPartitions(func, True)
- def distinct(self, numPartitions=None):
+ def distinct(self: "RDD[T]", numPartitions: Optional[int] = None) ->
"RDD[T]":
"""
Return a new RDD containing the distinct elements in this RDD.
@@ -527,7 +609,9 @@ class RDD:
.map(lambda x: x[0])
)
- def sample(self, withReplacement, fraction, seed=None):
+ def sample(
+ self: "RDD[T]", withReplacement: bool, fraction: float, seed:
Optional[int] = None
+ ) -> "RDD[T]":
"""
Return a sampled subset of this RDD.
@@ -556,7 +640,9 @@ class RDD:
assert fraction >= 0.0, "Negative fraction value: %s" % fraction
return self.mapPartitionsWithIndex(RDDSampler(withReplacement,
fraction, seed).func, True)
- def randomSplit(self, weights, seed=None):
+ def randomSplit(
+ self: "RDD[T]", weights: Sequence[Union[int, float]], seed:
Optional[int] = None
+ ) -> "List[RDD[T]]":
"""
Randomly splits this RDD with the provided weights.
@@ -593,7 +679,9 @@ class RDD:
]
# this is ported from scala/spark/RDD.scala
- def takeSample(self, withReplacement, num, seed=None):
+ def takeSample(
+ self: "RDD[T]", withReplacement: bool, num: int, seed: Optional[int] =
None
+ ) -> List[T]:
"""
Return a fixed-size sampled subset of this RDD.
@@ -651,7 +739,9 @@ class RDD:
return samples[0:num]
@staticmethod
- def _computeFractionForSampleSize(sampleSizeLowerBound, total,
withReplacement):
+ def _computeFractionForSampleSize(
+ sampleSizeLowerBound: int, total: int, withReplacement: bool
+ ) -> float:
"""
Returns a sampling rate that guarantees a sample of
size >= sampleSizeLowerBound 99.99% of the time.
@@ -683,7 +773,7 @@ class RDD:
gamma = -log(delta) / total
return min(1, fraction + gamma + sqrt(gamma * gamma + 2 * gamma *
fraction))
- def union(self, other):
+ def union(self: "RDD[T]", other: "RDD[U]") -> "RDD[Union[T, U]]":
"""
Return the union of this RDD and another one.
@@ -694,7 +784,9 @@ class RDD:
[1, 1, 2, 3, 1, 1, 2, 3]
"""
if self._jrdd_deserializer == other._jrdd_deserializer:
- rdd = RDD(self._jrdd.union(other._jrdd), self.ctx,
self._jrdd_deserializer)
+ rdd: "RDD[Union[T, U]]" = RDD(
+ self._jrdd.union(other._jrdd), self.ctx,
self._jrdd_deserializer
+ )
else:
# These RDDs contain data in different serialized formats, so we
# must normalize them to the default serializer.
@@ -708,7 +800,7 @@ class RDD:
rdd.partitioner = self.partitioner
return rdd
- def intersection(self, other):
+ def intersection(self: "RDD[T]", other: "RDD[T]") -> "RDD[T]":
"""
Return the intersection of this RDD and another one. The output will
not contain any duplicate elements, even if the input RDDs did.
@@ -731,14 +823,14 @@ class RDD:
.keys()
)
- def _reserialize(self, serializer=None):
+ def _reserialize(self: "RDD[T]", serializer: Optional[Serializer] = None)
-> "RDD[T]":
serializer = serializer or self.ctx.serializer
if self._jrdd_deserializer != serializer:
self = self.map(lambda x: x, preservesPartitioning=True)
self._jrdd_deserializer = serializer
return self
- def __add__(self, other):
+ def __add__(self: "RDD[T]", other: "RDD[U]") -> "RDD[Union[T, U]]":
"""
Return the union of this RDD and another one.
@@ -752,9 +844,43 @@ class RDD:
raise TypeError
return self.union(other)
+ @overload
def repartitionAndSortWithinPartitions(
- self, numPartitions=None, partitionFunc=portable_hash, ascending=True,
keyfunc=lambda x: x
- ):
+ self: "RDD[Tuple[S, V]]",
+ numPartitions: Optional[int] = ...,
+ partitionFunc: Callable[["S"], int] = ...,
+ ascending: bool = ...,
+ ) -> "RDD[Tuple[S, V]]":
+ ...
+
+ @overload
+ def repartitionAndSortWithinPartitions(
+ self: "RDD[Tuple[K, V]]",
+ numPartitions: Optional[int],
+ partitionFunc: Callable[[K], int],
+ ascending: bool,
+ keyfunc: Callable[[K], "S"],
+ ) -> "RDD[Tuple[K, V]]":
+ ...
+
+ @overload
+ def repartitionAndSortWithinPartitions(
+ self: "RDD[Tuple[K, V]]",
+ numPartitions: Optional[int] = ...,
+ partitionFunc: Callable[[K], int] = ...,
+ ascending: bool = ...,
+ *,
+ keyfunc: Callable[[K], "S"],
+ ) -> "RDD[Tuple[K, V]]":
+ ...
+
+ def repartitionAndSortWithinPartitions(
+ self: "RDD[Tuple[Any, Any]]",
+ numPartitions: Optional[int] = None,
+ partitionFunc: Callable[[Any], int] = portable_hash,
+ ascending: bool = True,
+ keyfunc: Callable[[Any], Any] = lambda x: x,
+ ) -> "RDD[Tuple[Any, Any]]":
"""
Repartition the RDD according to the given partitioner and, within
each resulting partition,
sort records by their keys.
@@ -772,13 +898,45 @@ class RDD:
memory = self._memory_limit()
serializer = self._jrdd_deserializer
- def sortPartition(iterator):
+ def sortPartition(iterator: Iterable[Tuple[K, V]]) ->
Iterable[Tuple[K, V]]:
sort = ExternalSorter(memory * 0.9, serializer).sorted
return iter(sort(iterator, key=lambda k_v: keyfunc(k_v[0]),
reverse=(not ascending)))
return self.partitionBy(numPartitions,
partitionFunc).mapPartitions(sortPartition, True)
- def sortByKey(self, ascending=True, numPartitions=None, keyfunc=lambda x:
x):
+ @overload
+ def sortByKey(
+ self: "RDD[Tuple[S, V]]",
+ ascending: bool = ...,
+ numPartitions: Optional[int] = ...,
+ ) -> "RDD[Tuple[K, V]]":
+ ...
+
+ @overload
+ def sortByKey(
+ self: "RDD[Tuple[K, V]]",
+ ascending: bool,
+ numPartitions: int,
+ keyfunc: Callable[[K], "S"],
+ ) -> "RDD[Tuple[K, V]]":
+ ...
+
+ @overload
+ def sortByKey(
+ self: "RDD[Tuple[K, V]]",
+ ascending: bool = ...,
+ numPartitions: Optional[int] = ...,
+ *,
+ keyfunc: Callable[[K], "S"],
+ ) -> "RDD[Tuple[K, V]]":
+ ...
+
+ def sortByKey(
+ self: "RDD[Tuple[K, V]]",
+ ascending: Optional[bool] = True,
+ numPartitions: Optional[int] = None,
+ keyfunc: Callable[[Any], Any] = lambda x: x,
+ ) -> "RDD[Tuple[K, V]]":
"""
Sorts this RDD, which is assumed to consist of (key, value) pairs.
@@ -802,7 +960,7 @@ class RDD:
memory = self._memory_limit()
serializer = self._jrdd_deserializer
- def sortPartition(iterator):
+ def sortPartition(iterator: Iterable[Tuple[K, V]]) ->
Iterable[Tuple[K, V]]:
sort = ExternalSorter(memory * 0.9, serializer).sorted
return iter(sort(iterator, key=lambda kv: keyfunc(kv[0]),
reverse=(not ascending)))
@@ -829,16 +987,21 @@ class RDD:
for i in range(0, numPartitions - 1)
]
- def rangePartitioner(k):
+ def rangePartitioner(k: K) -> int:
p = bisect.bisect_left(bounds, keyfunc(k))
if ascending:
return p
else:
- return numPartitions - 1 - p
+ return numPartitions - 1 - p # type: ignore[operator]
return self.partitionBy(numPartitions,
rangePartitioner).mapPartitions(sortPartition, True)
- def sortBy(self, keyfunc, ascending=True, numPartitions=None):
+ def sortBy(
+ self: "RDD[T]",
+ keyfunc: Callable[[T], "S"],
+ ascending: bool = True,
+ numPartitions: Optional[int] = None,
+ ) -> "RDD[T]":
"""
Sorts this RDD by the given keyfunc
@@ -850,9 +1013,13 @@ class RDD:
>>> sc.parallelize(tmp).sortBy(lambda x: x[1]).collect()
[('a', 1), ('b', 2), ('1', 3), ('d', 4), ('2', 5)]
"""
- return self.keyBy(keyfunc).sortByKey(ascending, numPartitions).values()
+ return (
+ self.keyBy(keyfunc) # type: ignore[type-var]
+ .sortByKey(ascending, numPartitions)
+ .values()
+ )
- def glom(self):
+ def glom(self: "RDD[T]") -> "RDD[List[T]]":
"""
Return an RDD created by coalescing all elements within each partition
into a list.
@@ -864,12 +1031,12 @@ class RDD:
[[1, 2], [3, 4]]
"""
- def func(iterator):
+ def func(iterator: Iterable[T]) -> Iterable[List[T]]:
yield list(iterator)
return self.mapPartitions(func)
- def cartesian(self, other):
+ def cartesian(self: "RDD[T]", other: "RDD[U]") -> "RDD[Tuple[T, U]]":
"""
Return the Cartesian product of this RDD and another one, that is, the
RDD of all pairs of elements ``(a, b)`` where ``a`` is in `self` and
@@ -885,7 +1052,12 @@ class RDD:
deserializer = CartesianDeserializer(self._jrdd_deserializer,
other._jrdd_deserializer)
return RDD(self._jrdd.cartesian(other._jrdd), self.ctx, deserializer)
- def groupBy(self, f, numPartitions=None, partitionFunc=portable_hash):
+ def groupBy(
+ self: "RDD[T]",
+ f: Callable[[T], K],
+ numPartitions: Optional[int] = None,
+ partitionFunc: Callable[[K], int] = portable_hash,
+ ) -> "RDD[Tuple[K, Iterable[T]]]":
"""
Return an RDD of grouped items.
@@ -898,7 +1070,9 @@ class RDD:
"""
return self.map(lambda x: (f(x), x)).groupByKey(numPartitions,
partitionFunc)
- def pipe(self, command, env=None, checkCode=False):
+ def pipe(
+ self, command: str, env: Optional[Dict[str, str]] = None, checkCode:
bool = False
+ ) -> "RDD[str]":
"""
Return an RDD created by piping elements to a forked external process.
@@ -919,10 +1093,10 @@ class RDD:
if env is None:
env = dict()
- def func(iterator):
+ def func(iterator: Iterable[T]) -> Iterable[str]:
pipe = Popen(shlex.split(command), env=env, stdin=PIPE,
stdout=PIPE)
- def pipe_objs(out):
+ def pipe_objs(out: IO[bytes]) -> None:
for obj in iterator:
s = str(obj).rstrip("\n") + "\n"
out.write(s.encode("utf-8"))
@@ -930,7 +1104,7 @@ class RDD:
Thread(target=pipe_objs, args=[pipe.stdin]).start()
- def check_return_code():
+ def check_return_code() -> Iterable[int]:
pipe.wait()
if checkCode and pipe.returncode:
raise RuntimeError(
@@ -942,13 +1116,15 @@ class RDD:
yield i
return (
- x.rstrip(b"\n").decode("utf-8")
- for x in chain(iter(pipe.stdout.readline, b""),
check_return_code())
+ cast(bytes, x).rstrip(b"\n").decode("utf-8")
+ for x in chain(
+ iter(cast(IO[bytes], pipe.stdout).readline, b""),
check_return_code()
+ )
)
return self.mapPartitions(func)
- def foreach(self, f):
+ def foreach(self: "RDD[T]", f: Callable[[T], None]) -> None:
"""
Applies a function to all elements of this RDD.
@@ -959,14 +1135,14 @@ class RDD:
"""
f = fail_on_stopiteration(f)
- def processPartition(iterator):
+ def processPartition(iterator: Iterable[T]) -> Iterable[Any]:
for x in iterator:
f(x)
return iter([])
self.mapPartitions(processPartition).count() # Force evaluation
- def foreachPartition(self, f):
+ def foreachPartition(self: "RDD[T]", f: Callable[[Iterable[T]], None]) ->
None:
"""
Applies a function to each partition of this RDD.
@@ -978,16 +1154,16 @@ class RDD:
>>> sc.parallelize([1, 2, 3, 4, 5]).foreachPartition(f)
"""
- def func(it):
+ def func(it: Iterable[T]) -> Iterable[Any]:
r = f(it)
try:
- return iter(r)
+ return iter(r) # type: ignore[call-overload]
except TypeError:
return iter([])
self.mapPartitions(func).count() # Force evaluation
- def collect(self):
+ def collect(self: "RDD[T]") -> List[T]:
"""
Return a list that contains all of the elements in this RDD.
@@ -997,10 +1173,13 @@ class RDD:
to be small, as all the data is loaded into the driver's memory.
"""
with SCCallSiteSync(self.context):
+ assert self.ctx._jvm is not None
sock_info =
self.ctx._jvm.PythonRDD.collectAndServe(self._jrdd.rdd())
return list(_load_from_socket(sock_info, self._jrdd_deserializer))
- def collectWithJobGroup(self, groupId, description,
interruptOnCancel=False):
+ def collectWithJobGroup(
+ self: "RDD[T]", groupId: str, description: str, interruptOnCancel:
bool = False
+ ) -> "List[T]":
"""
When collect rdd, use this method to specify job group.
@@ -1015,12 +1194,13 @@ class RDD:
)
with SCCallSiteSync(self.context):
+ assert self.ctx._jvm is not None
sock_info = self.ctx._jvm.PythonRDD.collectAndServeWithJobGroup(
self._jrdd.rdd(), groupId, description, interruptOnCancel
)
return list(_load_from_socket(sock_info, self._jrdd_deserializer))
- def reduce(self, f):
+ def reduce(self: "RDD[T]", f: Callable[[T, T], T]) -> T:
"""
Reduces the elements of this RDD using the specified commutative and
associative binary operator. Currently reduces partitions locally.
@@ -1039,7 +1219,7 @@ class RDD:
"""
f = fail_on_stopiteration(f)
- def func(iterator):
+ def func(iterator: Iterable[T]) -> Iterable[T]:
iterator = iter(iterator)
try:
initial = next(iterator)
@@ -1052,7 +1232,7 @@ class RDD:
return reduce(f, vals)
raise ValueError("Can not reduce() empty RDD")
- def treeReduce(self, f, depth=2):
+ def treeReduce(self: "RDD[T]", f: Callable[[T, T], T], depth: int = 2) ->
T:
"""
Reduces the elements of this RDD in a multi-level tree pattern.
@@ -1080,22 +1260,26 @@ class RDD:
if depth < 1:
raise ValueError("Depth cannot be smaller than 1 but got %d." %
depth)
- zeroValue = None, True # Use the second entry to indicate whether
this is a dummy value.
+ # Use the second entry to indicate whether this is a dummy value.
+ zeroValue: Tuple[T, bool] = ( # type: ignore[assignment]
+ None,
+ True,
+ )
- def op(x, y):
+ def op(x: Tuple[T, bool], y: Tuple[T, bool]) -> Tuple[T, bool]:
if x[1]:
return y
elif y[1]:
return x
else:
- return f(x[0], y[0]), False
+ return f(x[0], y[0]), False # type: ignore[arg-type]
reduced = self.map(lambda x: (x, False)).treeAggregate(zeroValue, op,
op, depth)
if reduced[1]:
raise ValueError("Cannot reduce empty RDD.")
return reduced[0]
- def fold(self, zeroValue, op):
+ def fold(self: "RDD[T]", zeroValue: T, op: Callable[[T, T], T]) -> T:
"""
Aggregate the elements of each partition, and then the results for all
the partitions, using a given associative function and a neutral "zero
value."
@@ -1120,7 +1304,7 @@ class RDD:
"""
op = fail_on_stopiteration(op)
- def func(iterator):
+ def func(iterator: Iterable[T]) -> Iterable[T]:
acc = zeroValue
for obj in iterator:
acc = op(acc, obj)
@@ -1132,7 +1316,9 @@ class RDD:
vals = self.mapPartitions(func).collect()
return reduce(op, vals, zeroValue)
- def aggregate(self, zeroValue, seqOp, combOp):
+ def aggregate(
+ self: "RDD[T]", zeroValue: U, seqOp: Callable[[U, T], U], combOp:
Callable[[U, U], U]
+ ) -> U:
"""
Aggregate the elements of each partition, and then the results for all
the partitions, using a given combine functions and a neutral "zero
@@ -1158,7 +1344,7 @@ class RDD:
seqOp = fail_on_stopiteration(seqOp)
combOp = fail_on_stopiteration(combOp)
- def func(iterator):
+ def func(iterator: Iterable[T]) -> Iterable[U]:
acc = zeroValue
for obj in iterator:
acc = seqOp(acc, obj)
@@ -1170,7 +1356,13 @@ class RDD:
vals = self.mapPartitions(func).collect()
return reduce(combOp, vals, zeroValue)
- def treeAggregate(self, zeroValue, seqOp, combOp, depth=2):
+ def treeAggregate(
+ self: "RDD[T]",
+ zeroValue: U,
+ seqOp: Callable[[U, T], U],
+ combOp: Callable[[U, U], U],
+ depth: int = 2,
+ ) -> U:
"""
Aggregates the elements of this RDD in a multi-level tree
pattern.
@@ -1199,7 +1391,7 @@ class RDD:
if self.getNumPartitions() == 0:
return zeroValue
- def aggregatePartition(iterator):
+ def aggregatePartition(iterator: Iterable[T]) -> Iterable[U]:
acc = zeroValue
for obj in iterator:
acc = seqOp(acc, obj)
@@ -1211,10 +1403,10 @@ class RDD:
# If creating an extra level doesn't help reduce the wall-clock time,
we stop the tree
# aggregation.
while numPartitions > scale + numPartitions / scale:
- numPartitions /= scale
+ numPartitions /= scale # type: ignore[assignment]
curNumPartitions = int(numPartitions)
- def mapPartition(i, iterator):
+ def mapPartition(i: int, iterator: Iterable[U]) ->
Iterable[Tuple[int, U]]:
for obj in iterator:
yield (i % curNumPartitions, obj)
@@ -1226,7 +1418,15 @@ class RDD:
return partiallyAggregated.reduce(combOp)
- def max(self, key=None):
+ @overload
+ def max(self: "RDD[S]") -> "S":
+ ...
+
+ @overload
+ def max(self: "RDD[T]", key: Callable[[T], "S"]) -> T:
+ ...
+
+ def max(self: "RDD[T]", key: Optional[Callable[[T], "S"]] = None) -> T:
"""
Find the maximum item in this RDD.
@@ -1244,10 +1444,18 @@ class RDD:
5.0
"""
if key is None:
- return self.reduce(max)
- return self.reduce(lambda a, b: max(a, b, key=key))
+ return self.reduce(max) # type: ignore[arg-type]
+ return self.reduce(lambda a, b: max(a, b, key=key)) # type:
ignore[arg-type]
- def min(self, key=None):
+ @overload
+ def min(self: "RDD[S]") -> "S":
+ ...
+
+ @overload
+ def min(self: "RDD[T]", key: Callable[[T], "S"]) -> T:
+ ...
+
+ def min(self: "RDD[T]", key: Optional[Callable[[T], "S"]] = None) -> T:
"""
Find the minimum item in this RDD.
@@ -1265,10 +1473,10 @@ class RDD:
10.0
"""
if key is None:
- return self.reduce(min)
- return self.reduce(lambda a, b: min(a, b, key=key))
+ return self.reduce(min) # type: ignore[arg-type]
+ return self.reduce(lambda a, b: min(a, b, key=key)) # type:
ignore[arg-type]
- def sum(self):
+ def sum(self: "RDD[NumberOrArray]") -> "NumberOrArray":
"""
Add up the elements in this RDD.
@@ -1277,9 +1485,11 @@ class RDD:
>>> sc.parallelize([1.0, 2.0, 3.0]).sum()
6.0
"""
- return self.mapPartitions(lambda x: [sum(x)]).fold(0, operator.add)
+ return self.mapPartitions(lambda x: [sum(x)]).fold( # type:
ignore[return-value]
+ 0, operator.add
+ )
- def count(self):
+ def count(self) -> int:
"""
Return the number of elements in this RDD.
@@ -1290,18 +1500,22 @@ class RDD:
"""
return self.mapPartitions(lambda i: [sum(1 for _ in i)]).sum()
- def stats(self):
+ def stats(self: "RDD[NumberOrArray]") -> StatCounter:
"""
Return a :class:`StatCounter` object that captures the mean, variance
and count of the RDD's elements in one operation.
"""
- def redFunc(left_counter, right_counter):
+ def redFunc(left_counter: StatCounter, right_counter: StatCounter) ->
StatCounter:
return left_counter.mergeStats(right_counter)
- return self.mapPartitions(lambda i: [StatCounter(i)]).reduce(redFunc)
+ return self.mapPartitions(lambda i: [StatCounter(i)]).reduce( # type:
ignore[arg-type]
+ redFunc
+ )
- def histogram(self, buckets):
+ def histogram(
+ self: "RDD[S]", buckets: Union[int, List["S"], Tuple["S", ...]]
+ ) -> Tuple[Sequence["S"], List[int]]:
"""
Compute a histogram using the provided buckets. The buckets
are all open to the right except for the last which is closed.
@@ -1345,7 +1559,7 @@ class RDD:
raise ValueError("number of buckets must be >= 1")
# filter out non-comparable elements
- def comparable(x):
+ def comparable(x: Any) -> bool:
if x is None:
return False
if type(x) is float and isnan(x):
@@ -1355,7 +1569,7 @@ class RDD:
filtered = self.filter(comparable)
# faster than stats()
- def minmax(a, b):
+ def minmax(a: Tuple["S", "S"], b: Tuple["S", "S"]) -> Tuple["S",
"S"]:
return min(a[0], b[0]), max(a[1], b[1])
try:
@@ -1369,7 +1583,7 @@ class RDD:
return [minv, maxv], [filtered.count()]
try:
- inc = (maxv - minv) / buckets
+ inc = (maxv - minv) / buckets # type: ignore[operator]
except TypeError:
raise TypeError("Can not generate buckets with non-number in
RDD")
@@ -1378,8 +1592,8 @@ class RDD:
# keep them as integer if possible
inc = int(inc)
- if inc * buckets != maxv - minv:
- inc = (maxv - minv) * 1.0 / buckets
+ if inc * buckets != maxv - minv: # type: ignore[operator]
+ inc = (maxv - minv) * 1.0 / buckets # type: ignore[operator]
buckets = [i * inc + minv for i in range(buckets)]
buckets.append(maxv) # fix accumulated error
@@ -1403,35 +1617,47 @@ class RDD:
even = False
inc = None
try:
- steps = [buckets[i + 1] - buckets[i] for i in
range(len(buckets) - 1)]
+ steps = [
+ buckets[i + 1] - buckets[i] # type: ignore[operator]
+ for i in range(len(buckets) - 1)
+ ]
except TypeError:
pass # objects in buckets do not support '-'
else:
if max(steps) - min(steps) < 1e-10: # handle precision errors
even = True
- inc = (maxv - minv) / (len(buckets) - 1)
+ inc = (maxv - minv) / (len(buckets) - 1) # type:
ignore[operator]
else:
raise TypeError("buckets should be a list or tuple or number(int
or long)")
- def histogram(iterator):
- counters = [0] * len(buckets)
+ def histogram(iterator: Iterable["S"]) -> Iterable[List[int]]:
+ counters = [0] * len(buckets) # type: ignore[arg-type]
for i in iterator:
- if i is None or (type(i) is float and isnan(i)) or i > maxv or
i < minv:
+ if (
+ i is None
+ or (isinstance(i, float) and isnan(i)) # type:
ignore[arg-type]
+ or i > maxv
+ or i < minv
+ ):
continue
- t = int((i - minv) / inc) if even else
bisect.bisect_right(buckets, i) - 1
+ t = (
+ int((i - minv) / inc) # type: ignore[operator]
+ if even
+ else bisect.bisect_right(buckets, i) - 1 # type:
ignore[arg-type]
+ )
counters[t] += 1
# add last two together
last = counters.pop()
counters[-1] += last
return [counters]
- def mergeCounters(a, b):
+ def mergeCounters(a: List[int], b: List[int]) -> List[int]:
return [i + j for i, j in zip(a, b)]
return buckets, self.mapPartitions(histogram).reduce(mergeCounters)
- def mean(self):
+ def mean(self: "RDD[NumberOrArray]") -> "NumberOrArray":
"""
Compute the mean of this RDD's elements.
@@ -1440,9 +1666,9 @@ class RDD:
>>> sc.parallelize([1, 2, 3]).mean()
2.0
"""
- return self.stats().mean()
+ return self.stats().mean() # type: ignore[return-value]
- def variance(self):
+ def variance(self: "RDD[NumberOrArray]") -> "NumberOrArray":
"""
Compute the variance of this RDD's elements.
@@ -1451,9 +1677,9 @@ class RDD:
>>> sc.parallelize([1, 2, 3]).variance()
0.666...
"""
- return self.stats().variance()
+ return self.stats().variance() # type: ignore[return-value]
- def stdev(self):
+ def stdev(self: "RDD[NumberOrArray]") -> "NumberOrArray":
"""
Compute the standard deviation of this RDD's elements.
@@ -1462,9 +1688,9 @@ class RDD:
>>> sc.parallelize([1, 2, 3]).stdev()
0.816...
"""
- return self.stats().stdev()
+ return self.stats().stdev() # type: ignore[return-value]
- def sampleStdev(self):
+ def sampleStdev(self: "RDD[NumberOrArray]") -> "NumberOrArray":
"""
Compute the sample standard deviation of this RDD's elements (which
corrects for bias in estimating the standard deviation by dividing by
@@ -1475,9 +1701,9 @@ class RDD:
>>> sc.parallelize([1, 2, 3]).sampleStdev()
1.0
"""
- return self.stats().sampleStdev()
+ return self.stats().sampleStdev() # type: ignore[return-value]
- def sampleVariance(self):
+ def sampleVariance(self: "RDD[NumberOrArray]") -> "NumberOrArray":
"""
Compute the sample variance of this RDD's elements (which corrects
for bias in estimating the variance by dividing by N-1 instead of N).
@@ -1487,9 +1713,9 @@ class RDD:
>>> sc.parallelize([1, 2, 3]).sampleVariance()
1.0
"""
- return self.stats().sampleVariance()
+ return self.stats().sampleVariance() # type: ignore[return-value]
- def countByValue(self):
+ def countByValue(self: "RDD[K]") -> Dict[K, int]:
"""
Return the count of each unique value in this RDD as a dictionary of
(value, count) pairs.
@@ -1500,20 +1726,28 @@ class RDD:
[(1, 2), (2, 3)]
"""
- def countPartition(iterator):
- counts = defaultdict(int)
+ def countPartition(iterator: Iterable[K]) -> Iterable[Dict[K, int]]:
+ counts: Dict[K, int] = defaultdict(int)
for obj in iterator:
counts[obj] += 1
yield counts
- def mergeMaps(m1, m2):
+ def mergeMaps(m1: Dict[K, int], m2: Dict[K, int]) -> Dict[K, int]:
for k, v in m2.items():
m1[k] += v
return m1
return self.mapPartitions(countPartition).reduce(mergeMaps)
- def top(self, num, key=None):
+ @overload
+ def top(self: "RDD[S]", num: int) -> List["S"]:
+ ...
+
+ @overload
+ def top(self: "RDD[T]", num: int, key: Callable[[T], "S"]) -> List[T]:
+ ...
+
+ def top(self: "RDD[T]", num: int, key: Optional[Callable[[T], "S"]] =
None) -> List[T]:
"""
Get the top N elements from an RDD.
@@ -1534,15 +1768,23 @@ class RDD:
[4, 3, 2]
"""
- def topIterator(iterator):
+ def topIterator(iterator: Iterable[T]) -> Iterable[List[T]]:
yield heapq.nlargest(num, iterator, key=key)
- def merge(a, b):
+ def merge(a: List[T], b: List[T]) -> List[T]:
return heapq.nlargest(num, a + b, key=key)
return self.mapPartitions(topIterator).reduce(merge)
- def takeOrdered(self, num, key=None):
+ @overload
+ def takeOrdered(self: "RDD[S]", num: int) -> List["S"]:
+ ...
+
+ @overload
+ def takeOrdered(self: "RDD[T]", num: int, key: Callable[[T], "S"]) ->
List[T]:
+ ...
+
+ def takeOrdered(self: "RDD[T]", num: int, key: Optional[Callable[[T],
"S"]] = None) -> List[T]:
"""
Get the N elements from an RDD ordered in ascending order or as
specified by the optional key function.
@@ -1560,12 +1802,12 @@ class RDD:
[10, 9, 7, 6, 5, 4]
"""
- def merge(a, b):
+ def merge(a: List[T], b: List[T]) -> List[T]:
return heapq.nsmallest(num, a + b, key)
return self.mapPartitions(lambda it: [heapq.nsmallest(num, it,
key)]).reduce(merge)
- def take(self, num):
+ def take(self: "RDD[T]", num: int) -> List[T]:
"""
Take the first num elements of the RDD.
@@ -1589,7 +1831,7 @@ class RDD:
>>> sc.parallelize(range(100), 100).filter(lambda x: x > 90).take(3)
[91, 92, 93]
"""
- items = []
+ items: List[T] = []
totalParts = self.getNumPartitions()
partsScanned = 0
@@ -1612,7 +1854,7 @@ class RDD:
left = num - len(items)
- def takeUpToNumLeft(iterator):
+ def takeUpToNumLeft(iterator: Iterable[T]) -> Iterable[T]:
iterator = iter(iterator)
taken = 0
while taken < left:
@@ -1630,7 +1872,7 @@ class RDD:
return items[:num]
- def first(self):
+ def first(self: "RDD[T]") -> T:
"""
Return the first element in this RDD.
@@ -1648,7 +1890,7 @@ class RDD:
return rs[0]
raise ValueError("RDD is empty")
- def isEmpty(self):
+ def isEmpty(self) -> bool:
"""
Returns true if and only if the RDD contains no elements at all.
@@ -1665,7 +1907,12 @@ class RDD:
"""
return self.getNumPartitions() == 0 or len(self.take(1)) == 0
- def saveAsNewAPIHadoopDataset(self, conf, keyConverter=None,
valueConverter=None):
+ def saveAsNewAPIHadoopDataset(
+ self: "RDD[Tuple[K, V]]",
+ conf: Dict[str, str],
+ keyConverter: Optional[str] = None,
+ valueConverter: Optional[str] = None,
+ ) -> None:
"""
Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to
any Hadoop file
system, using the new Hadoop OutputFormat API (mapreduce package).
Keys/values are
@@ -1683,20 +1930,22 @@ class RDD:
"""
jconf = self.ctx._dictToJavaMap(conf)
pickledRDD = self._pickled()
+ assert self.ctx._jvm is not None
+
self.ctx._jvm.PythonRDD.saveAsHadoopDataset(
pickledRDD._jrdd, True, jconf, keyConverter, valueConverter, True
)
def saveAsNewAPIHadoopFile(
- self,
- path,
- outputFormatClass,
- keyClass=None,
- valueClass=None,
- keyConverter=None,
- valueConverter=None,
- conf=None,
- ):
+ self: "RDD[Tuple[K, V]]",
+ path: str,
+ outputFormatClass: str,
+ keyClass: Optional[str] = None,
+ valueClass: Optional[str] = None,
+ keyConverter: Optional[str] = None,
+ valueConverter: Optional[str] = None,
+ conf: Optional[Dict[str, str]] = None,
+ ) -> None:
"""
Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to
any Hadoop file
system, using the new Hadoop OutputFormat API (mapreduce package). Key
and value types
@@ -1725,6 +1974,8 @@ class RDD:
"""
jconf = self.ctx._dictToJavaMap(conf)
pickledRDD = self._pickled()
+ assert self.ctx._jvm is not None
+
self.ctx._jvm.PythonRDD.saveAsNewAPIHadoopFile(
pickledRDD._jrdd,
True,
@@ -1737,7 +1988,12 @@ class RDD:
jconf,
)
- def saveAsHadoopDataset(self, conf, keyConverter=None,
valueConverter=None):
+ def saveAsHadoopDataset(
+ self: "RDD[Tuple[K, V]]",
+ conf: Dict[str, str],
+ keyConverter: Optional[str] = None,
+ valueConverter: Optional[str] = None,
+ ) -> None:
"""
Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to
any Hadoop file
system, using the old Hadoop OutputFormat API (mapred package).
Keys/values are
@@ -1755,21 +2011,23 @@ class RDD:
"""
jconf = self.ctx._dictToJavaMap(conf)
pickledRDD = self._pickled()
+ assert self.ctx._jvm is not None
+
self.ctx._jvm.PythonRDD.saveAsHadoopDataset(
pickledRDD._jrdd, True, jconf, keyConverter, valueConverter, False
)
def saveAsHadoopFile(
- self,
- path,
- outputFormatClass,
- keyClass=None,
- valueClass=None,
- keyConverter=None,
- valueConverter=None,
- conf=None,
- compressionCodecClass=None,
- ):
+ self: "RDD[Tuple[K, V]]",
+ path: str,
+ outputFormatClass: str,
+ keyClass: Optional[str] = None,
+ valueClass: Optional[str] = None,
+ keyConverter: Optional[str] = None,
+ valueConverter: Optional[str] = None,
+ conf: Optional[Dict[str, str]] = None,
+ compressionCodecClass: Optional[str] = None,
+ ) -> None:
"""
Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to
any Hadoop file
system, using the old Hadoop OutputFormat API (mapred package). Key
and value types
@@ -1803,6 +2061,8 @@ class RDD:
"""
jconf = self.ctx._dictToJavaMap(conf)
pickledRDD = self._pickled()
+ assert self.ctx._jvm is not None
+
self.ctx._jvm.PythonRDD.saveAsHadoopFile(
pickledRDD._jrdd,
True,
@@ -1816,7 +2076,9 @@ class RDD:
compressionCodecClass,
)
- def saveAsSequenceFile(self, path, compressionCodecClass=None):
+ def saveAsSequenceFile(
+ self: "RDD[Tuple[K, V]]", path: str, compressionCodecClass:
Optional[str] = None
+ ) -> None:
"""
Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to
any Hadoop file
system, using the "org.apache.hadoop.io.Writable" types that we
convert from the
@@ -1834,11 +2096,13 @@ class RDD:
i.e. "org.apache.hadoop.io.compress.GzipCodec" (None by default)
"""
pickledRDD = self._pickled()
+ assert self.ctx._jvm is not None
+
self.ctx._jvm.PythonRDD.saveAsSequenceFile(
pickledRDD._jrdd, True, path, compressionCodecClass
)
- def saveAsPickleFile(self, path, batchSize=10):
+ def saveAsPickleFile(self, path: str, batchSize: int = 10) -> None:
"""
Save this RDD as a SequenceFile of serialized objects. The serializer
used is :class:`pyspark.serializers.CPickleSerializer`, default batch
size
@@ -1853,13 +2117,14 @@ class RDD:
>>> sorted(sc.pickleFile(tmpFile.name, 5).map(str).collect())
['1', '2', 'rdd', 'spark']
"""
+ ser: Serializer
if batchSize == 0:
ser = AutoBatchedSerializer(CPickleSerializer())
else:
ser = BatchedSerializer(CPickleSerializer(), batchSize)
self._reserialize(ser)._jrdd.saveAsObjectFile(path)
- def saveAsTextFile(self, path, compressionCodecClass=None):
+ def saveAsTextFile(self, path: str, compressionCodecClass: Optional[str] =
None) -> None:
"""
Save this RDD as a text file, using string representations of elements.
@@ -1904,16 +2169,20 @@ class RDD:
'bar\\nfoo\\n'
"""
- def func(split, iterator):
+ def func(split: int, iterator: Iterable[Any]) -> Iterable[bytes]:
for x in iterator:
- if not isinstance(x, (str, bytes)):
- x = str(x)
- if isinstance(x, str):
- x = x.encode("utf-8")
- yield x
+ if isinstance(x, bytes):
+ yield x
+ elif isinstance(x, str):
+ yield x.encode("utf-8")
+ else:
+ yield str(x).encode("utf-8")
keyed = self.mapPartitionsWithIndex(func)
- keyed._bypass_serializer = True
+ keyed._bypass_serializer = True # type: ignore[attr-defined]
+
+ assert self.ctx._jvm is not None
+
if compressionCodecClass:
compressionCodec =
self.ctx._jvm.java.lang.Class.forName(compressionCodecClass)
keyed._jrdd.map(self.ctx._jvm.BytesToString()).saveAsTextFile(path,
compressionCodec)
@@ -1922,7 +2191,7 @@ class RDD:
# Pair functions
- def collectAsMap(self):
+ def collectAsMap(self: "RDD[Tuple[K, V]]") -> Dict[K, V]:
"""
Return the key-value pairs in this RDD to the master as a dictionary.
@@ -1941,7 +2210,7 @@ class RDD:
"""
return dict(self.collect())
- def keys(self):
+ def keys(self: "RDD[Tuple[K, V]]") -> "RDD[K]":
"""
Return an RDD with the keys of each tuple.
@@ -1953,7 +2222,7 @@ class RDD:
"""
return self.map(lambda x: x[0])
- def values(self):
+ def values(self: "RDD[Tuple[K, V]]") -> "RDD[V]":
"""
Return an RDD with the values of each tuple.
@@ -1965,7 +2234,12 @@ class RDD:
"""
return self.map(lambda x: x[1])
- def reduceByKey(self, func, numPartitions=None,
partitionFunc=portable_hash):
+ def reduceByKey(
+ self: "RDD[Tuple[K, V]]",
+ func: Callable[[V, V], V],
+ numPartitions: Optional[int] = None,
+ partitionFunc: Callable[[K], int] = portable_hash,
+ ) -> "RDD[Tuple[K, V]]":
"""
Merge the values for each key using an associative and commutative
reduce function.
@@ -1985,7 +2259,7 @@ class RDD:
"""
return self.combineByKey(lambda x: x, func, func, numPartitions,
partitionFunc)
- def reduceByKeyLocally(self, func):
+ def reduceByKeyLocally(self: "RDD[Tuple[K, V]]", func: Callable[[V, V],
V]) -> Dict[K, V]:
"""
Merge the values for each key using an associative and commutative
reduce function, but
return the results immediately to the master as a dictionary.
@@ -2002,20 +2276,20 @@ class RDD:
"""
func = fail_on_stopiteration(func)
- def reducePartition(iterator):
- m = {}
+ def reducePartition(iterator: Iterable[Tuple[K, V]]) ->
Iterable[Dict[K, V]]:
+ m: Dict[K, V] = {}
for k, v in iterator:
m[k] = func(m[k], v) if k in m else v
yield m
- def mergeMaps(m1, m2):
+ def mergeMaps(m1: Dict[K, V], m2: Dict[K, V]) -> Dict[K, V]:
for k, v in m2.items():
m1[k] = func(m1[k], v) if k in m1 else v
return m1
return self.mapPartitions(reducePartition).reduce(mergeMaps)
- def countByKey(self):
+ def countByKey(self: "RDD[Tuple[K, V]]") -> Dict[K, int]:
"""
Count the number of elements for each key, and return the result to the
master as a dictionary.
@@ -2028,7 +2302,11 @@ class RDD:
"""
return self.map(lambda x: x[0]).countByValue()
- def join(self, other, numPartitions=None):
+ def join(
+ self: "RDD[Tuple[K, V]]",
+ other: "RDD[Tuple[K, U]]",
+ numPartitions: Optional[int] = None,
+ ) -> "RDD[Tuple[K, Tuple[V, U]]]":
"""
Return an RDD containing all pairs of elements with matching keys in
`self` and `other`.
@@ -2047,7 +2325,11 @@ class RDD:
"""
return python_join(self, other, numPartitions)
- def leftOuterJoin(self, other, numPartitions=None):
+ def leftOuterJoin(
+ self: "RDD[Tuple[K, V]]",
+ other: "RDD[Tuple[K, U]]",
+ numPartitions: Optional[int] = None,
+ ) -> "RDD[Tuple[K, Tuple[V, Optional[U]]]]":
"""
Perform a left outer join of `self` and `other`.
@@ -2066,7 +2348,11 @@ class RDD:
"""
return python_left_outer_join(self, other, numPartitions)
- def rightOuterJoin(self, other, numPartitions=None):
+ def rightOuterJoin(
+ self: "RDD[Tuple[K, V]]",
+ other: "RDD[Tuple[K, U]]",
+ numPartitions: Optional[int] = None,
+ ) -> "RDD[Tuple[K, Tuple[Optional[V], U]]]":
"""
Perform a right outer join of `self` and `other`.
@@ -2085,7 +2371,11 @@ class RDD:
"""
return python_right_outer_join(self, other, numPartitions)
- def fullOuterJoin(self, other, numPartitions=None):
+ def fullOuterJoin(
+ self: "RDD[Tuple[K, V]]",
+ other: "RDD[Tuple[K, U]]",
+ numPartitions: Optional[int] = None,
+ ) -> "RDD[Tuple[K, Tuple[Optional[V], Optional[U]]]]":
"""
Perform a right outer join of `self` and `other`.
@@ -2111,7 +2401,11 @@ class RDD:
# TODO: add option to control map-side combining
# portable_hash is used as default, because builtin hash of None is
different
# cross machines.
- def partitionBy(self, numPartitions, partitionFunc=portable_hash):
+ def partitionBy(
+ self: "RDD[Tuple[K, V]]",
+ numPartitions: Optional[int],
+ partitionFunc: Callable[[K], int] = portable_hash,
+ ) -> "RDD[Tuple[K, V]]":
"""
Return a copy of the RDD partitioned using the specified partitioner.
@@ -2138,13 +2432,13 @@ class RDD:
limit = self._memory_limit() / 2
- def add_shuffle_key(split, iterator):
+ def add_shuffle_key(split: int, iterator: Iterable[Tuple[K, V]]) ->
Iterable[bytes]:
buckets = defaultdict(list)
- c, batch = 0, min(10 * numPartitions, 1000)
+ c, batch = 0, min(10 * numPartitions, 1000) # type:
ignore[operator]
for k, v in iterator:
- buckets[partitionFunc(k) % numPartitions].append((k, v))
+ buckets[partitionFunc(k) % numPartitions].append((k, v)) #
type: ignore[operator]
c += 1
# check used memory and avg size of chunk of objects
@@ -2160,7 +2454,7 @@ class RDD:
avg = int(size / n) >> 20
# let 1M < avg < 10M
if avg < 1:
- batch = min(sys.maxsize, batch * 1.5)
+ batch = min(sys.maxsize, batch * 1.5) # type:
ignore[assignment]
elif avg > 10:
batch = max(int(batch / 1.5), 1)
c = 0
@@ -2170,24 +2464,26 @@ class RDD:
yield outputSerializer.dumps(items)
keyed = self.mapPartitionsWithIndex(add_shuffle_key,
preservesPartitioning=True)
- keyed._bypass_serializer = True
+ keyed._bypass_serializer = True # type: ignore[attr-defined]
+ assert self.ctx._jvm is not None
+
with SCCallSiteSync(self.context):
pairRDD =
self.ctx._jvm.PairwiseRDD(keyed._jrdd.rdd()).asJavaPairRDD()
jpartitioner = self.ctx._jvm.PythonPartitioner(numPartitions,
id(partitionFunc))
jrdd =
self.ctx._jvm.PythonRDD.valueOfPair(pairRDD.partitionBy(jpartitioner))
- rdd = RDD(jrdd, self.ctx, BatchedSerializer(outputSerializer))
+ rdd: "RDD[Tuple[K, V]]" = RDD(jrdd, self.ctx,
BatchedSerializer(outputSerializer))
rdd.partitioner = partitioner
return rdd
# TODO: add control over map-side aggregation
def combineByKey(
- self,
- createCombiner,
- mergeValue,
- mergeCombiners,
- numPartitions=None,
- partitionFunc=portable_hash,
- ):
+ self: "RDD[Tuple[K, V]]",
+ createCombiner: Callable[[V], U],
+ mergeValue: Callable[[U, V], U],
+ mergeCombiners: Callable[[U, U], U],
+ numPartitions: Optional[int] = None,
+ partitionFunc: Callable[[K], int] = portable_hash,
+ ) -> "RDD[Tuple[K, U]]":
"""
Generic function to combine the elements for each key using a custom
set of aggregation functions.
@@ -2238,7 +2534,7 @@ class RDD:
memory = self._memory_limit()
agg = Aggregator(createCombiner, mergeValue, mergeCombiners)
- def combineLocally(iterator):
+ def combineLocally(iterator: Iterable[Tuple[K, V]]) ->
Iterable[Tuple[K, U]]:
merger = ExternalMerger(agg, memory * 0.9, serializer)
merger.mergeValues(iterator)
return merger.items()
@@ -2246,7 +2542,7 @@ class RDD:
locally_combined = self.mapPartitions(combineLocally,
preservesPartitioning=True)
shuffled = locally_combined.partitionBy(numPartitions, partitionFunc)
- def _mergeCombiners(iterator):
+ def _mergeCombiners(iterator: Iterable[Tuple[K, U]]) ->
Iterable[Tuple[K, U]]:
merger = ExternalMerger(agg, memory, serializer)
merger.mergeCombiners(iterator)
return merger.items()
@@ -2254,8 +2550,13 @@ class RDD:
return shuffled.mapPartitions(_mergeCombiners,
preservesPartitioning=True)
def aggregateByKey(
- self, zeroValue, seqFunc, combFunc, numPartitions=None,
partitionFunc=portable_hash
- ):
+ self: "RDD[Tuple[K, V]]",
+ zeroValue: U,
+ seqFunc: Callable[[U, V], U],
+ combFunc: Callable[[U, U], U],
+ numPartitions: Optional[int] = None,
+ partitionFunc: Callable[[K], int] = portable_hash,
+ ) -> "RDD[Tuple[K, U]]":
"""
Aggregate the values of each key, using given combine functions and a
neutral
"zero value". This function can return a different result type, U,
than the type
@@ -2266,14 +2567,20 @@ class RDD:
allowed to modify and return their first argument instead of creating
a new U.
"""
- def createZero():
+ def createZero() -> U:
return copy.deepcopy(zeroValue)
return self.combineByKey(
lambda v: seqFunc(createZero(), v), seqFunc, combFunc,
numPartitions, partitionFunc
)
- def foldByKey(self, zeroValue, func, numPartitions=None,
partitionFunc=portable_hash):
+ def foldByKey(
+ self: "RDD[Tuple[K, V]]",
+ zeroValue: V,
+ func: Callable[[V, V], V],
+ numPartitions: Optional[int] = None,
+ partitionFunc: Callable[[K], int] = portable_hash,
+ ) -> "RDD[Tuple[K, V]]":
"""
Merge the values for each key using an associative function "func"
and a neutral "zeroValue" which may be added to the result an
@@ -2288,18 +2595,22 @@ class RDD:
[('a', 2), ('b', 1)]
"""
- def createZero():
+ def createZero() -> V:
return copy.deepcopy(zeroValue)
return self.combineByKey(
lambda v: func(createZero(), v), func, func, numPartitions,
partitionFunc
)
- def _memory_limit(self):
+ def _memory_limit(self) -> int:
return _parse_memory(self.ctx._conf.get("spark.python.worker.memory",
"512m"))
# TODO: support variant with custom partitioner
- def groupByKey(self, numPartitions=None, partitionFunc=portable_hash):
+ def groupByKey(
+ self: "RDD[Tuple[K, V]]",
+ numPartitions: Optional[int] = None,
+ partitionFunc: Callable[[K], int] = portable_hash,
+ ) -> "RDD[Tuple[K, Iterable[V]]]":
"""
Group the values for each key in the RDD into a single sequence.
Hash-partitions the resulting RDD with numPartitions partitions.
@@ -2319,14 +2630,14 @@ class RDD:
[('a', [1, 1]), ('b', [1])]
"""
- def createCombiner(x):
+ def createCombiner(x: V) -> List[V]:
return [x]
- def mergeValue(xs, x):
+ def mergeValue(xs: List[V], x: V) -> List[V]:
xs.append(x)
return xs
- def mergeCombiners(a, b):
+ def mergeCombiners(a: List[V], b: List[V]) -> List[V]:
a.extend(b)
return a
@@ -2334,7 +2645,7 @@ class RDD:
serializer = self._jrdd_deserializer
agg = Aggregator(createCombiner, mergeValue, mergeCombiners)
- def combine(iterator):
+ def combine(iterator: Iterable[Tuple[K, V]]) -> Iterable[Tuple[K,
List[V]]]:
merger = ExternalMerger(agg, memory * 0.9, serializer)
merger.mergeValues(iterator)
return merger.items()
@@ -2342,14 +2653,16 @@ class RDD:
locally_combined = self.mapPartitions(combine,
preservesPartitioning=True)
shuffled = locally_combined.partitionBy(numPartitions, partitionFunc)
- def groupByKey(it):
+ def groupByKey(it: Iterable[Tuple[K, List[V]]]) -> Iterable[Tuple[K,
List[V]]]:
merger = ExternalGroupBy(agg, memory, serializer)
merger.mergeCombiners(it)
return merger.items()
return shuffled.mapPartitions(groupByKey,
True).mapValues(ResultIterable)
- def flatMapValues(self, f):
+ def flatMapValues(
+ self: "RDD[Tuple[K, V]]", f: Callable[[V], Iterable[U]]
+ ) -> "RDD[Tuple[K, U]]":
"""
Pass each value in the key-value pair RDD through a flatMap function
without changing the keys; this also retains the original RDD's
@@ -2363,12 +2676,12 @@ class RDD:
[('a', 'x'), ('a', 'y'), ('a', 'z'), ('b', 'p'), ('b', 'r')]
"""
- def flat_map_fn(kv):
+ def flat_map_fn(kv: Tuple[K, V]) -> Iterable[Tuple[K, U]]:
return ((kv[0], x) for x in f(kv[1]))
return self.flatMap(flat_map_fn, preservesPartitioning=True)
- def mapValues(self, f):
+ def mapValues(self: "RDD[Tuple[K, V]]", f: Callable[[V], U]) ->
"RDD[Tuple[K, U]]":
"""
Pass each value in the key-value pair RDD through a map function
without changing the keys; this also retains the original RDD's
@@ -2382,12 +2695,45 @@ class RDD:
[('a', 3), ('b', 1)]
"""
- def map_values_fn(kv):
+ def map_values_fn(kv: Tuple[K, V]) -> Tuple[K, U]:
return kv[0], f(kv[1])
return self.map(map_values_fn, preservesPartitioning=True)
- def groupWith(self, other, *others):
+ @overload
+ def groupWith(
+ self: "RDD[Tuple[K, V]]", other: "RDD[Tuple[K, V1]]"
+ ) -> "RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[V1]]]]":
+ ...
+
+ @overload
+ def groupWith(
+ self: "RDD[Tuple[K, V]]", other: "RDD[Tuple[K, V1]]", __o1:
"RDD[Tuple[K, V2]]"
+ ) -> "RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[V1],
ResultIterable[V2]]]]":
+ ...
+
+ @overload
+ def groupWith(
+ self: "RDD[Tuple[K, V]]",
+ other: "RDD[Tuple[K, V1]]",
+ _o1: "RDD[Tuple[K, V2]]",
+ _o2: "RDD[Tuple[K, V3]]",
+ ) -> """RDD[
+ Tuple[
+ K,
+ Tuple[
+ ResultIterable[V],
+ ResultIterable[V1],
+ ResultIterable[V2],
+ ResultIterable[V3],
+ ],
+ ]
+ ]""":
+ ...
+
+ def groupWith( # type: ignore[misc]
+ self: "RDD[Tuple[Any, Any]]", other: "RDD[Tuple[Any, Any]]", *others:
"RDD[Tuple[Any, Any]]"
+ ) -> "RDD[Tuple[Any, Tuple[ResultIterable[Any], ...]]]":
"""
Alias for cogroup but with support for multiple RDDs.
@@ -2404,7 +2750,11 @@ class RDD:
return python_cogroup((self, other) + others, numPartitions=None)
# TODO: add variant with custom partitioner
- def cogroup(self, other, numPartitions=None):
+ def cogroup(
+ self: "RDD[Tuple[K, V]]",
+ other: "RDD[Tuple[K, U]]",
+ numPartitions: Optional[int] = None,
+ ) -> "RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[U]]]]":
"""
For each key k in `self` or `other`, return a resulting RDD that
contains a tuple with the list of values for that key in `self` as
@@ -2419,7 +2769,12 @@ class RDD:
"""
return python_cogroup((self, other), numPartitions)
- def sampleByKey(self, withReplacement, fractions, seed=None):
+ def sampleByKey(
+ self: "RDD[Tuple[K, V]]",
+ withReplacement: bool,
+ fractions: Dict[K, Union[float, int]],
+ seed: Optional[int] = None,
+ ) -> "RDD[Tuple[K, V]]":
"""
Return a subset of this RDD sampled by key (via stratified sampling).
Create a sample of this RDD using variable sampling rates for
@@ -2443,7 +2798,11 @@ class RDD:
RDDStratifiedSampler(withReplacement, fractions, seed).func, True
)
- def subtractByKey(self, other, numPartitions=None):
+ def subtractByKey(
+ self: "RDD[Tuple[K, V]]",
+ other: "RDD[Tuple[K, Any]]",
+ numPartitions: Optional[int] = None,
+ ) -> "RDD[Tuple[K, V]]":
"""
Return each (key, value) pair in `self` that has no pair with matching
key in `other`.
@@ -2456,13 +2815,17 @@ class RDD:
[('b', 4), ('b', 5)]
"""
- def filter_func(pair):
+ def filter_func(pair: Tuple[K, Tuple[V, Any]]) -> bool:
key, (val1, val2) = pair
- return val1 and not val2
+ return val1 and not val2 # type: ignore[return-value]
- return self.cogroup(other,
numPartitions).filter(filter_func).flatMapValues(lambda x: x[0])
+ return (
+ self.cogroup(other, numPartitions)
+ .filter(filter_func) # type: ignore[arg-type]
+ .flatMapValues(lambda x: x[0])
+ )
- def subtract(self, other, numPartitions=None):
+ def subtract(self: "RDD[T]", other: "RDD[T]", numPartitions: Optional[int]
= None) -> "RDD[T]":
"""
Return each value in `self` that is not contained in `other`.
@@ -2477,7 +2840,7 @@ class RDD:
rdd = other.map(lambda x: (x, True))
return self.map(lambda x: (x, True)).subtractByKey(rdd,
numPartitions).keys()
- def keyBy(self, f):
+ def keyBy(self: "RDD[T]", f: Callable[[T], K]) -> "RDD[Tuple[K, T]]":
"""
Creates tuples of the elements in this RDD by applying `f`.
@@ -2490,7 +2853,7 @@ class RDD:
"""
return self.map(lambda x: (f(x), x))
- def repartition(self, numPartitions):
+ def repartition(self: "RDD[T]", numPartitions: int) -> "RDD[T]":
"""
Return a new RDD that has exactly numPartitions partitions.
@@ -2511,7 +2874,7 @@ class RDD:
"""
return self.coalesce(numPartitions, shuffle=True)
- def coalesce(self, numPartitions, shuffle=False):
+ def coalesce(self: "RDD[T]", numPartitions: int, shuffle: bool = False) ->
"RDD[T]":
"""
Return a new RDD that is reduced into `numPartitions` partitions.
@@ -2535,7 +2898,7 @@ class RDD:
jrdd = self._jrdd.coalesce(numPartitions, shuffle)
return RDD(jrdd, self.ctx, jrdd_deserializer)
- def zip(self, other):
+ def zip(self: "RDD[T]", other: "RDD[U]") -> "RDD[Tuple[T, U]]":
"""
Zips this RDD with another one, returning key-value pairs with the
first element in each RDD second element in each RDD, etc. Assumes
@@ -2551,12 +2914,12 @@ class RDD:
[(0, 1000), (1, 1001), (2, 1002), (3, 1003), (4, 1004)]
"""
- def get_batch_size(ser):
+ def get_batch_size(ser: Serializer) -> int:
if isinstance(ser, BatchedSerializer):
return ser.batchSize
return 1 # not batched
- def batch_as(rdd, batchSize):
+ def batch_as(rdd: "RDD[V]", batchSize: int) -> "RDD[V]":
return rdd._reserialize(BatchedSerializer(CPickleSerializer(),
batchSize))
my_batch = get_batch_size(self._jrdd_deserializer)
@@ -2579,7 +2942,7 @@ class RDD:
deserializer = PairDeserializer(self._jrdd_deserializer,
other._jrdd_deserializer)
return RDD(pairRDD, self.ctx, deserializer)
- def zipWithIndex(self):
+ def zipWithIndex(self: "RDD[T]") -> "RDD[Tuple[T, int]]":
"""
Zips this RDD with its element indices.
@@ -2602,13 +2965,13 @@ class RDD:
for i in range(len(nums) - 1):
starts.append(starts[-1] + nums[i])
- def func(k, it):
+ def func(k: int, it: Iterable[T]) -> Iterable[Tuple[T, int]]:
for i, v in enumerate(it, starts[k]):
yield v, i
return self.mapPartitionsWithIndex(func)
- def zipWithUniqueId(self):
+ def zipWithUniqueId(self: "RDD[T]") -> "RDD[Tuple[T, int]]":
"""
Zips this RDD with generated unique Long ids.
@@ -2624,21 +2987,20 @@ class RDD:
"""
n = self.getNumPartitions()
- def func(k, it):
+ def func(k: int, it: Iterable[T]) -> Iterable[Tuple[T, int]]:
for i, v in enumerate(it):
yield v, i * n + k
return self.mapPartitionsWithIndex(func)
- def name(self):
+ def name(self) -> Optional[str]:
"""
Return the name of this RDD.
"""
n = self._jrdd.name()
- if n:
- return n
+ return n if n else None
- def setName(self, name):
+ def setName(self: "RDD[T]", name: str) -> "RDD[T]":
"""
Assign a name to this RDD.
@@ -2651,15 +3013,15 @@ class RDD:
self._jrdd.setName(name)
return self
- def toDebugString(self):
+ def toDebugString(self) -> Optional[bytes]:
"""
A description of this RDD and its recursive dependencies for debugging.
"""
debug_string = self._jrdd.toDebugString()
- if debug_string:
- return debug_string.encode("utf-8")
- def getStorageLevel(self):
+ return debug_string.encode("utf-8") if debug_string else None
+
+ def getStorageLevel(self) -> StorageLevel:
"""
Get the RDD's current storage level.
@@ -2681,7 +3043,7 @@ class RDD:
)
return storage_level
- def _defaultReducePartitions(self):
+ def _defaultReducePartitions(self) -> int:
"""
Returns the default number of partitions to use during reduce tasks
(e.g., groupBy).
If spark.default.parallelism is set, then we'll use the value from
SparkContext
@@ -2696,7 +3058,7 @@ class RDD:
else:
return self.getNumPartitions()
- def lookup(self, key):
+ def lookup(self: "RDD[Tuple[K, V]]", key: K) -> List[V]:
"""
Return the list of values in the RDD for key `key`. This operation
is done efficiently if the RDD has a known partitioner by only
@@ -2724,16 +3086,18 @@ class RDD:
return values.collect()
- def _to_java_object_rdd(self):
+ def _to_java_object_rdd(self) -> "JavaObject":
"""Return a JavaRDD of Object by unpickling
It will convert each Python object into Java object by Pickle,
whenever the
RDD is serialized in batch or not.
"""
rdd = self._pickled()
+ assert self.ctx._jvm is not None
+
return self.ctx._jvm.SerDeUtil.pythonToJava(rdd._jrdd, True)
- def countApprox(self, timeout, confidence=0.95):
+ def countApprox(self, timeout: int, confidence: float = 0.95) -> int:
"""
Approximate version of count() that returns a potentially incomplete
result within a timeout, even if not all tasks have finished.
@@ -2747,7 +3111,9 @@ class RDD:
drdd = self.mapPartitions(lambda it: [float(sum(1 for i in it))])
return int(drdd.sumApprox(timeout, confidence))
- def sumApprox(self, timeout, confidence=0.95):
+ def sumApprox(
+ self: "RDD[Union[float, int]]", timeout: int, confidence: float = 0.95
+ ) -> BoundedFloat:
"""
Approximate operation to return the sum within a timeout
or meet the confidence.
@@ -2760,11 +3126,14 @@ class RDD:
True
"""
jrdd = self.mapPartitions(lambda it:
[float(sum(it))])._to_java_object_rdd()
+ assert self.ctx._jvm is not None
jdrdd = self.ctx._jvm.JavaDoubleRDD.fromRDD(jrdd.rdd())
r = jdrdd.sumApprox(timeout, confidence).getFinalValue()
return BoundedFloat(r.mean(), r.confidence(), r.low(), r.high())
- def meanApprox(self, timeout, confidence=0.95):
+ def meanApprox(
+ self: "RDD[Union[float, int]]", timeout: int, confidence: float = 0.95
+ ) -> BoundedFloat:
"""
Approximate operation to return the mean within a timeout
or meet the confidence.
@@ -2777,11 +3146,12 @@ class RDD:
True
"""
jrdd = self.map(float)._to_java_object_rdd()
+ assert self.ctx._jvm is not None
jdrdd = self.ctx._jvm.JavaDoubleRDD.fromRDD(jrdd.rdd())
r = jdrdd.meanApprox(timeout, confidence).getFinalValue()
return BoundedFloat(r.mean(), r.confidence(), r.low(), r.high())
- def countApproxDistinct(self, relativeSD=0.05):
+ def countApproxDistinct(self: "RDD[T]", relativeSD: float = 0.05) -> int:
"""
Return approximate number of distinct elements in the RDD.
@@ -2814,7 +3184,7 @@ class RDD:
hashRDD = self.map(lambda x: portable_hash(x) & 0xFFFFFFFF)
return hashRDD._to_java_object_rdd().countApproxDistinct(relativeSD)
- def toLocalIterator(self, prefetchPartitions=False):
+ def toLocalIterator(self: "RDD[T]", prefetchPartitions: bool = False) ->
Iterator[T]:
"""
Return an iterator that contains all of the elements in this RDD.
The iterator will consume as much memory as the largest partition in
this RDD.
@@ -2832,13 +3202,15 @@ class RDD:
>>> [x for x in rdd.toLocalIterator()]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
"""
+ assert self.ctx._jvm is not None
+
with SCCallSiteSync(self.context):
sock_info = self.ctx._jvm.PythonRDD.toLocalIteratorAndServe(
self._jrdd.rdd(), prefetchPartitions
)
return _local_iterator_from_socket(sock_info, self._jrdd_deserializer)
- def barrier(self):
+ def barrier(self: "RDD[T]") -> "RDDBarrier[T]":
"""
Marks the current stage as a barrier stage, where Spark must launch
all tasks together.
In case of a task failure, instead of only restarting the failed task,
Spark will abort the
@@ -2868,13 +3240,13 @@ class RDD:
"""
return RDDBarrier(self)
- def _is_barrier(self):
+ def _is_barrier(self) -> bool:
"""
Whether this RDD is in a barrier stage.
"""
return self._jrdd.rdd().isBarrier()
- def withResources(self, profile):
+ def withResources(self: "RDD[T]", profile: ResourceProfile) -> "RDD[T]":
"""
Specify a :class:`pyspark.resource.ResourceProfile` to use when
calculating this RDD.
This is only supported on certain cluster managers and currently
requires dynamic
@@ -2891,6 +3263,8 @@ class RDD:
if profile._java_resource_profile is not None:
jrp = profile._java_resource_profile
else:
+ assert self.ctx._jvm is not None
+
builder =
self.ctx._jvm.org.apache.spark.resource.ResourceProfileBuilder()
ereqs = ExecutorResourceRequests(self.ctx._jvm,
profile._executor_resource_requests)
treqs = TaskResourceRequests(self.ctx._jvm,
profile._task_resource_requests)
@@ -2901,7 +3275,7 @@ class RDD:
self._jrdd.withResources(jrp)
return self
- def getResourceProfile(self):
+ def getResourceProfile(self) -> Optional[ResourceProfile]:
"""
Get the :class:`pyspark.resource.ResourceProfile` specified with this
RDD or None
if it wasn't specified.
@@ -2923,11 +3297,38 @@ class RDD:
else:
return None
+ @overload
+ def toDF(
+ self: "RDD[RowLike]",
+ schema: Optional[Union[List[str], Tuple[str, ...]]] = None,
+ sampleRatio: Optional[float] = None,
+ ) -> "DataFrame":
+ ...
+
+ @overload
+ def toDF(
+ self: "RDD[RowLike]", schema: Optional[Union["StructType", str]] = None
+ ) -> "DataFrame":
+ ...
+
+ @overload
+ def toDF(
+ self: "RDD[AtomicValue]",
+ schema: Union["AtomicType", str],
+ ) -> "DataFrame":
+ ...
+
+ def toDF(
+ self: "RDD[Any]", schema: Optional[Any] = None, sampleRatio:
Optional[float] = None
+ ) -> "DataFrame":
+ raise RuntimeError("""RDD.toDF was called before SparkSession was
initialized.""")
-def _prepare_for_python_RDD(sc, command):
+
+def _prepare_for_python_RDD(sc: "SparkContext", command: Any) -> Tuple[bytes,
Any, Any, Any]:
# the serialized command will be compressed by broadcast
ser = CloudPickleSerializer()
pickled_command = ser.dumps(command)
+ assert sc._jvm is not None
if len(pickled_command) >
sc._jvm.PythonUtils.getBroadcastThreshold(sc._jsc): # Default 1M
# The broadcast will have same life cycle as created PythonRDD
broadcast = sc.broadcast(pickled_command)
@@ -2937,11 +3338,14 @@ def _prepare_for_python_RDD(sc, command):
return pickled_command, broadcast_vars, sc.environment, sc._python_includes
-def _wrap_function(sc, func, deserializer, serializer, profiler=None):
+def _wrap_function(
+ sc: "SparkContext", func: Callable, deserializer: Any, serializer: Any,
profiler: Any = None
+) -> "JavaObject":
assert deserializer, "deserializer should not be empty"
assert serializer, "serializer should not be empty"
command = (func, profiler, deserializer, serializer)
pickled_command, broadcast_vars, env, includes =
_prepare_for_python_RDD(sc, command)
+ assert sc._jvm is not None
return sc._jvm.PythonFunction(
bytearray(pickled_command),
env,
@@ -2953,7 +3357,7 @@ def _wrap_function(sc, func, deserializer, serializer,
profiler=None):
)
-class RDDBarrier:
+class RDDBarrier(Generic[T]):
"""
Wraps an RDD in a barrier stage, which forces Spark to launch tasks of
this stage together.
@@ -2966,10 +3370,12 @@ class RDDBarrier:
This API is experimental
"""
- def __init__(self, rdd):
+ def __init__(self, rdd: RDD[T]):
self.rdd = rdd
- def mapPartitions(self, f, preservesPartitioning=False):
+ def mapPartitions(
+ self, f: Callable[[Iterable[T]], Iterable[U]], preservesPartitioning:
bool = False
+ ) -> RDD[U]:
"""
Returns a new RDD by applying a function to each partition of the
wrapped RDD,
where tasks are launched together in a barrier stage.
@@ -2983,12 +3389,16 @@ class RDDBarrier:
This API is experimental
"""
- def func(s, iterator):
+ def func(s: int, iterator: Iterable[T]) -> Iterable[U]:
return f(iterator)
return PipelinedRDD(self.rdd, func, preservesPartitioning,
isFromBarrier=True)
- def mapPartitionsWithIndex(self, f, preservesPartitioning=False):
+ def mapPartitionsWithIndex(
+ self,
+ f: Callable[[int, Iterable[T]], Iterable[U]],
+ preservesPartitioning: bool = False,
+ ) -> RDD[U]:
"""
Returns a new RDD by applying a function to each partition of the
wrapped RDD, while
tracking the index of the original partition. And all tasks are
launched together
@@ -3005,7 +3415,7 @@ class RDDBarrier:
return PipelinedRDD(self.rdd, f, preservesPartitioning,
isFromBarrier=True)
-class PipelinedRDD(RDD):
+class PipelinedRDD(RDD[U], Generic[T, U]):
"""
Examples
@@ -3027,7 +3437,13 @@ class PipelinedRDD(RDD):
20
"""
- def __init__(self, prev, func, preservesPartitioning=False,
isFromBarrier=False):
+ def __init__(
+ self,
+ prev: RDD[T],
+ func: Callable[[int, Iterable[T]], Iterable[U]],
+ preservesPartitioning: bool = False,
+ isFromBarrier: bool = False,
+ ):
if not isinstance(prev, PipelinedRDD) or not prev._is_pipelinable():
# This transformation is the first in its stage:
self.func = func
@@ -3035,9 +3451,9 @@ class PipelinedRDD(RDD):
self._prev_jrdd = prev._jrdd
self._prev_jrdd_deserializer = prev._jrdd_deserializer
else:
- prev_func = prev.func
+ prev_func: Callable[[int, Iterable[V]], Iterable[T]] = prev.func
- def pipeline_func(split, iterator):
+ def pipeline_func(split: int, iterator: Iterable[V]) ->
Iterable[U]:
return func(split, prev_func(split, iterator))
self.func = pipeline_func
@@ -3049,18 +3465,18 @@ class PipelinedRDD(RDD):
self.is_checkpointed = False
self.ctx = prev.ctx
self.prev = prev
- self._jrdd_val = None
+ self._jrdd_val: Optional["JavaObject"] = None
self._id = None
self._jrdd_deserializer = self.ctx.serializer
self._bypass_serializer = False
self.partitioner = prev.partitioner if self.preservesPartitioning else
None
self.is_barrier = isFromBarrier or prev._is_barrier()
- def getNumPartitions(self):
+ def getNumPartitions(self) -> int:
return self._prev_jrdd.partitions().size()
@property
- def _jrdd(self):
+ def _jrdd(self) -> "JavaObject":
if self._jrdd_val:
return self._jrdd_val
if self._bypass_serializer:
@@ -3074,29 +3490,32 @@ class PipelinedRDD(RDD):
wrapped_func = _wrap_function(
self.ctx, self.func, self._prev_jrdd_deserializer,
self._jrdd_deserializer, profiler
)
+
+ assert self.ctx._jvm is not None
python_rdd = self.ctx._jvm.PythonRDD(
self._prev_jrdd.rdd(), wrapped_func, self.preservesPartitioning,
self.is_barrier
)
self._jrdd_val = python_rdd.asJavaRDD()
if profiler:
+ assert self._jrdd_val is not None
self._id = self._jrdd_val.id()
self.ctx.profiler_collector.add_profiler(self._id, profiler)
return self._jrdd_val
- def id(self):
+ def id(self) -> int:
if self._id is None:
self._id = self._jrdd.id()
return self._id
- def _is_pipelinable(self):
+ def _is_pipelinable(self) -> bool:
return not (self.is_cached or self.is_checkpointed or
self.has_resource_profile)
- def _is_barrier(self):
+ def _is_barrier(self) -> bool:
return self.is_barrier
-def _test():
+def _test() -> None:
import doctest
from pyspark.context import SparkContext
diff --git a/python/pyspark/rdd.pyi b/python/pyspark/rdd.pyi
deleted file mode 100644
index c4eddbf..0000000
--- a/python/pyspark/rdd.pyi
+++ /dev/null
@@ -1,481 +0,0 @@
-#
-# 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
-from typing import (
- Any,
- Callable,
- Dict,
- Generic,
- Hashable,
- Iterable,
- Iterator,
- List,
- Optional,
- Tuple,
- Union,
- TypeVar,
-)
-from typing_extensions import Literal
-
-from numpy import int32, int64, float32, float64, ndarray
-
-from pyspark._typing import SupportsOrdering
-from pyspark.sql.pandas._typing import (
- PandasScalarUDFType,
- PandasScalarIterUDFType,
- PandasGroupedMapUDFType,
- PandasCogroupedMapUDFType,
- PandasGroupedAggUDFType,
- PandasMapIterUDFType,
- ArrowMapIterUDFType,
-)
-import pyspark.context
-from pyspark.resultiterable import ResultIterable
-from pyspark.serializers import Serializer
-from pyspark.storagelevel import StorageLevel
-from pyspark.resource.requests import ( # noqa: F401
- ExecutorResourceRequests,
- TaskResourceRequests,
-)
-from pyspark.resource.profile import ResourceProfile
-from pyspark.statcounter import StatCounter
-from pyspark.sql.dataframe import DataFrame
-from pyspark.sql.types import AtomicType, StructType
-from pyspark.sql._typing import AtomicValue, RowLike
-from py4j.java_gateway import JavaObject # type: ignore[import]
-
-T = TypeVar("T")
-T_co = TypeVar("T_co", covariant=True)
-U = TypeVar("U")
-K = TypeVar("K", bound=Hashable)
-V = TypeVar("V")
-V1 = TypeVar("V1")
-V2 = TypeVar("V2")
-V3 = TypeVar("V3")
-O = TypeVar("O", bound=SupportsOrdering)
-NumberOrArray = TypeVar(
- "NumberOrArray", float, int, complex, int32, int64, float32, float64,
ndarray
-)
-
-def portable_hash(x: Hashable) -> int: ...
-
-class PythonEvalType:
- NON_UDF: Literal[0]
- SQL_BATCHED_UDF: Literal[100]
- SQL_SCALAR_PANDAS_UDF: PandasScalarUDFType
- SQL_GROUPED_MAP_PANDAS_UDF: PandasGroupedMapUDFType
- SQL_GROUPED_AGG_PANDAS_UDF: PandasGroupedAggUDFType
- SQL_WINDOW_AGG_PANDAS_UDF: Literal[203]
- SQL_SCALAR_PANDAS_ITER_UDF: PandasScalarIterUDFType
- SQL_MAP_PANDAS_ITER_UDF: PandasMapIterUDFType
- SQL_COGROUPED_MAP_PANDAS_UDF: PandasCogroupedMapUDFType
- SQL_MAP_ARROW_ITER_UDF: ArrowMapIterUDFType
-
-class BoundedFloat(float):
- def __new__(cls, mean: float, confidence: float, low: float, high: float)
-> BoundedFloat: ...
-
-class Partitioner:
- numPartitions: int
- partitionFunc: Callable[[Any], int]
- def __init__(self, numPartitions: int, partitionFunc: Callable[[Any],
int]) -> None: ...
- def __eq__(self, other: Any) -> bool: ...
- def __call__(self, k: Any) -> int: ...
-
-class RDD(Generic[T_co]):
- is_cached: bool
- is_checkpointed: bool
- ctx: pyspark.context.SparkContext
- partitioner: Optional[Partitioner]
- def __init__(
- self,
- jrdd: JavaObject,
- ctx: pyspark.context.SparkContext,
- jrdd_deserializer: Serializer = ...,
- ) -> None: ...
- def id(self) -> int: ...
- def __getnewargs__(self) -> Any: ...
- @property
- def context(self) -> pyspark.context.SparkContext: ...
- def cache(self) -> RDD[T_co]: ...
- def persist(self, storageLevel: StorageLevel = ...) -> RDD[T_co]: ...
- def unpersist(self, blocking: bool = ...) -> RDD[T_co]: ...
- def checkpoint(self) -> None: ...
- def isCheckpointed(self) -> bool: ...
- def localCheckpoint(self) -> None: ...
- def isLocallyCheckpointed(self) -> bool: ...
- def getCheckpointFile(self) -> Optional[str]: ...
- def map(self, f: Callable[[T_co], U], preservesPartitioning: bool = ...)
-> RDD[U]: ...
- def flatMap(
- self, f: Callable[[T_co], Iterable[U]], preservesPartitioning: bool =
...
- ) -> RDD[U]: ...
- def mapPartitions(
- self, f: Callable[[Iterable[T_co]], Iterable[U]],
preservesPartitioning: bool = ...
- ) -> RDD[U]: ...
- def mapPartitionsWithIndex(
- self,
- f: Callable[[int, Iterable[T_co]], Iterable[U]],
- preservesPartitioning: bool = ...,
- ) -> RDD[U]: ...
- def mapPartitionsWithSplit(
- self,
- f: Callable[[int, Iterable[T_co]], Iterable[U]],
- preservesPartitioning: bool = ...,
- ) -> RDD[U]: ...
- def getNumPartitions(self) -> int: ...
- def filter(self, f: Callable[[T_co], bool]) -> RDD[T_co]: ...
- def distinct(self, numPartitions: Optional[int] = ...) -> RDD[T_co]: ...
- def sample(
- self, withReplacement: bool, fraction: float, seed: Optional[int] = ...
- ) -> RDD[T_co]: ...
- def randomSplit(
- self, weights: List[Union[int, float]], seed: Optional[int] = ...
- ) -> List[RDD[T_co]]: ...
- def takeSample(
- self, withReplacement: bool, num: int, seed: Optional[int] = ...
- ) -> List[T_co]: ...
- def union(self, other: RDD[U]) -> RDD[Union[T_co, U]]: ...
- def intersection(self, other: RDD[T_co]) -> RDD[T_co]: ...
- def __add__(self, other: RDD[T_co]) -> RDD[T_co]: ...
- @overload
- def repartitionAndSortWithinPartitions(
- self: RDD[Tuple[O, V]],
- numPartitions: Optional[int] = ...,
- partitionFunc: Callable[[O], int] = ...,
- ascending: bool = ...,
- ) -> RDD[Tuple[O, V]]: ...
- @overload
- def repartitionAndSortWithinPartitions(
- self: RDD[Tuple[K, V]],
- numPartitions: Optional[int],
- partitionFunc: Callable[[K], int],
- ascending: bool,
- keyfunc: Callable[[K], O],
- ) -> RDD[Tuple[K, V]]: ...
- @overload
- def repartitionAndSortWithinPartitions(
- self: RDD[Tuple[K, V]],
- numPartitions: Optional[int] = ...,
- partitionFunc: Callable[[K], int] = ...,
- ascending: bool = ...,
- *,
- keyfunc: Callable[[K], O],
- ) -> RDD[Tuple[K, V]]: ...
- @overload
- def sortByKey(
- self: RDD[Tuple[O, V]],
- ascending: bool = ...,
- numPartitions: Optional[int] = ...,
- ) -> RDD[Tuple[K, V]]: ...
- @overload
- def sortByKey(
- self: RDD[Tuple[K, V]],
- ascending: bool,
- numPartitions: int,
- keyfunc: Callable[[K], O],
- ) -> RDD[Tuple[K, V]]: ...
- @overload
- def sortByKey(
- self: RDD[Tuple[K, V]],
- ascending: bool = ...,
- numPartitions: Optional[int] = ...,
- *,
- keyfunc: Callable[[K], O],
- ) -> RDD[Tuple[K, V]]: ...
- def sortBy(
- self,
- keyfunc: Callable[[T_co], O],
- ascending: bool = ...,
- numPartitions: Optional[int] = ...,
- ) -> RDD[T_co]: ...
- def glom(self) -> RDD[List[T_co]]: ...
- def cartesian(self, other: RDD[U]) -> RDD[Tuple[T_co, U]]: ...
- def groupBy(
- self,
- f: Callable[[T_co], K],
- numPartitions: Optional[int] = ...,
- partitionFunc: Callable[[K], int] = ...,
- ) -> RDD[Tuple[K, Iterable[T_co]]]: ...
- def pipe(
- self, command: str, env: Optional[Dict[str, str]] = ..., checkCode:
bool = ...
- ) -> RDD[str]: ...
- def foreach(self, f: Callable[[T_co], None]) -> None: ...
- def foreachPartition(self, f: Callable[[Iterable[T_co]], None]) -> None:
...
- def collect(self) -> List[T_co]: ...
- def collectWithJobGroup(
- self, groupId: str, description: str, interruptOnCancel: bool = ...
- ) -> List[T_co]: ...
- def reduce(self, f: Callable[[T_co, T_co], T_co]) -> T_co: ...
- def treeReduce(self, f: Callable[[T_co, T_co], T_co], depth: int = ...) ->
T_co: ...
- def fold(self, zeroValue: T, op: Callable[[T_co, T_co], T_co]) -> T_co: ...
- def aggregate(
- self, zeroValue: U, seqOp: Callable[[U, T_co], U], combOp:
Callable[[U, U], U]
- ) -> U: ...
- def treeAggregate(
- self,
- zeroValue: U,
- seqOp: Callable[[U, T_co], U],
- combOp: Callable[[U, U], U],
- depth: int = ...,
- ) -> U: ...
- @overload
- def max(self: RDD[O]) -> O: ...
- @overload
- def max(self, key: Callable[[T_co], O]) -> T_co: ...
- @overload
- def min(self: RDD[O]) -> O: ...
- @overload
- def min(self, key: Callable[[T_co], O]) -> T_co: ...
- def sum(self: RDD[NumberOrArray]) -> NumberOrArray: ...
- def count(self) -> int: ...
- def stats(self: RDD[NumberOrArray]) -> StatCounter: ...
- def histogram(
- self, buckets: Union[int, List[T_co], Tuple[T_co, ...]]
- ) -> Tuple[List[T_co], List[int]]: ...
- def mean(self: RDD[NumberOrArray]) -> NumberOrArray: ...
- def variance(self: RDD[NumberOrArray]) -> NumberOrArray: ...
- def stdev(self: RDD[NumberOrArray]) -> NumberOrArray: ...
- def sampleStdev(self: RDD[NumberOrArray]) -> NumberOrArray: ...
- def sampleVariance(self: RDD[NumberOrArray]) -> NumberOrArray: ...
- def countByValue(self: RDD[K]) -> Dict[K, int]: ...
- @overload
- def top(self: RDD[O], num: int) -> List[O]: ...
- @overload
- def top(self, num: int, key: Callable[[T_co], O]) -> List[T_co]: ...
- @overload
- def takeOrdered(self: RDD[O], num: int) -> List[O]: ...
- @overload
- def takeOrdered(self, num: int, key: Callable[[T_co], O]) -> List[T_co]:
...
- def take(self, num: int) -> List[T_co]: ...
- def first(self) -> T_co: ...
- def isEmpty(self) -> bool: ...
- def saveAsNewAPIHadoopDataset(
- self: RDD[Tuple[K, V]],
- conf: Dict[str, str],
- keyConverter: Optional[str] = ...,
- valueConverter: Optional[str] = ...,
- ) -> None: ...
- def saveAsNewAPIHadoopFile(
- self: RDD[Tuple[K, V]],
- path: str,
- outputFormatClass: str,
- keyClass: Optional[str] = ...,
- valueClass: Optional[str] = ...,
- keyConverter: Optional[str] = ...,
- valueConverter: Optional[str] = ...,
- conf: Optional[Dict[str, str]] = ...,
- ) -> None: ...
- def saveAsHadoopDataset(
- self: RDD[Tuple[K, V]],
- conf: Dict[str, str],
- keyConverter: Optional[str] = ...,
- valueConverter: Optional[str] = ...,
- ) -> None: ...
- def saveAsHadoopFile(
- self: RDD[Tuple[K, V]],
- path: str,
- outputFormatClass: str,
- keyClass: Optional[str] = ...,
- valueClass: Optional[str] = ...,
- keyConverter: Optional[str] = ...,
- valueConverter: Optional[str] = ...,
- conf: Optional[str] = ...,
- compressionCodecClass: Optional[str] = ...,
- ) -> None: ...
- def saveAsSequenceFile(
- self: RDD[Tuple[K, V]], path: str, compressionCodecClass:
Optional[str] = ...
- ) -> None: ...
- def saveAsPickleFile(self, path: str, batchSize: int = ...) -> None: ...
- def saveAsTextFile(self, path: str, compressionCodecClass: Optional[str] =
...) -> None: ...
- def collectAsMap(self: RDD[Tuple[K, V]]) -> Dict[K, V]: ...
- def keys(self: RDD[Tuple[K, V]]) -> RDD[K]: ...
- def values(self: RDD[Tuple[K, V]]) -> RDD[V]: ...
- def reduceByKey(
- self: RDD[Tuple[K, V]],
- func: Callable[[V, V], V],
- numPartitions: Optional[int] = ...,
- partitionFunc: Callable[[K], int] = ...,
- ) -> RDD[Tuple[K, V]]: ...
- def reduceByKeyLocally(self: RDD[Tuple[K, V]], func: Callable[[V, V], V])
-> Dict[K, V]: ...
- def countByKey(self: RDD[Tuple[K, V]]) -> Dict[K, int]: ...
- def join(
- self: RDD[Tuple[K, V]],
- other: RDD[Tuple[K, U]],
- numPartitions: Optional[int] = ...,
- ) -> RDD[Tuple[K, Tuple[V, U]]]: ...
- def leftOuterJoin(
- self: RDD[Tuple[K, V]],
- other: RDD[Tuple[K, U]],
- numPartitions: Optional[int] = ...,
- ) -> RDD[Tuple[K, Tuple[V, Optional[U]]]]: ...
- def rightOuterJoin(
- self: RDD[Tuple[K, V]],
- other: RDD[Tuple[K, U]],
- numPartitions: Optional[int] = ...,
- ) -> RDD[Tuple[K, Tuple[Optional[V], U]]]: ...
- def fullOuterJoin(
- self: RDD[Tuple[K, V]],
- other: RDD[Tuple[K, U]],
- numPartitions: Optional[int] = ...,
- ) -> RDD[Tuple[K, Tuple[Optional[V], Optional[U]]]]: ...
- def partitionBy(
- self: RDD[Tuple[K, V]],
- numPartitions: int,
- partitionFunc: Callable[[K], int] = ...,
- ) -> RDD[Tuple[K, V]]: ...
- def combineByKey(
- self: RDD[Tuple[K, V]],
- createCombiner: Callable[[V], U],
- mergeValue: Callable[[U, V], U],
- mergeCombiners: Callable[[U, U], U],
- numPartitions: Optional[int] = ...,
- partitionFunc: Callable[[K], int] = ...,
- ) -> RDD[Tuple[K, U]]: ...
- def aggregateByKey(
- self: RDD[Tuple[K, V]],
- zeroValue: U,
- seqFunc: Callable[[U, V], U],
- combFunc: Callable[[U, U], U],
- numPartitions: Optional[int] = ...,
- partitionFunc: Callable[[K], int] = ...,
- ) -> RDD[Tuple[K, U]]: ...
- def foldByKey(
- self: RDD[Tuple[K, V]],
- zeroValue: V,
- func: Callable[[V, V], V],
- numPartitions: Optional[int] = ...,
- partitionFunc: Callable[[K], int] = ...,
- ) -> RDD[Tuple[K, V]]: ...
- def groupByKey(
- self: RDD[Tuple[K, V]],
- numPartitions: Optional[int] = ...,
- partitionFunc: Callable[[K], int] = ...,
- ) -> RDD[Tuple[K, Iterable[V]]]: ...
- def flatMapValues(
- self: RDD[Tuple[K, V]], f: Callable[[V], Iterable[U]]
- ) -> RDD[Tuple[K, U]]: ...
- def mapValues(self: RDD[Tuple[K, V]], f: Callable[[V], U]) -> RDD[Tuple[K,
U]]: ...
- @overload
- def groupWith(
- self: RDD[Tuple[K, V]], __o: RDD[Tuple[K, V1]]
- ) -> RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[V1]]]]: ...
- @overload
- def groupWith(
- self: RDD[Tuple[K, V]], __o1: RDD[Tuple[K, V1]], __o2: RDD[Tuple[K,
V2]]
- ) -> RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[V1],
ResultIterable[V2]]]]: ...
- @overload
- def groupWith(
- self: RDD[Tuple[K, V]],
- other1: RDD[Tuple[K, V1]],
- other2: RDD[Tuple[K, V2]],
- other3: RDD[Tuple[K, V3]],
- ) -> RDD[
- Tuple[
- K,
- Tuple[
- ResultIterable[V],
- ResultIterable[V1],
- ResultIterable[V2],
- ResultIterable[V3],
- ],
- ]
- ]: ...
- def cogroup(
- self: RDD[Tuple[K, V]],
- other: RDD[Tuple[K, U]],
- numPartitions: Optional[int] = ...,
- ) -> RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[U]]]]: ...
- def sampleByKey(
- self: RDD[Tuple[K, V]],
- withReplacement: bool,
- fractions: Dict[K, Union[float, int]],
- seed: Optional[int] = ...,
- ) -> RDD[Tuple[K, V]]: ...
- def subtractByKey(
- self: RDD[Tuple[K, V]],
- other: RDD[Tuple[K, U]],
- numPartitions: Optional[int] = ...,
- ) -> RDD[Tuple[K, V]]: ...
- def subtract(self, other: RDD[T_co], numPartitions: Optional[int] = ...)
-> RDD[T_co]: ...
- def keyBy(self, f: Callable[[T_co], K]) -> RDD[Tuple[K, T_co]]: ...
- def repartition(self, numPartitions: int) -> RDD[T_co]: ...
- def coalesce(self, numPartitions: int, shuffle: bool = ...) -> RDD[T_co]:
...
- def zip(self, other: RDD[U]) -> RDD[Tuple[T_co, U]]: ...
- def zipWithIndex(self) -> RDD[Tuple[T_co, int]]: ...
- def zipWithUniqueId(self) -> RDD[Tuple[T_co, int]]: ...
- def name(self) -> str: ...
- def setName(self, name: str) -> RDD[T_co]: ...
- def toDebugString(self) -> bytes: ...
- def getStorageLevel(self) -> StorageLevel: ...
- def lookup(self: RDD[Tuple[K, V]], key: K) -> List[V]: ...
- def countApprox(self, timeout: int, confidence: float = ...) -> int: ...
- def sumApprox(
- self: RDD[Union[float, int]], timeout: int, confidence: float = ...
- ) -> BoundedFloat: ...
- def meanApprox(
- self: RDD[Union[float, int]], timeout: int, confidence: float = ...
- ) -> BoundedFloat: ...
- def countApproxDistinct(self, relativeSD: float = ...) -> int: ...
- def toLocalIterator(self, prefetchPartitions: bool = ...) ->
Iterator[T_co]: ...
- def barrier(self) -> RDDBarrier[T_co]: ...
- def withResources(self, profile: ResourceProfile) -> RDD[T_co]: ...
- def getResourceProfile(self) -> Optional[ResourceProfile]: ...
- @overload
- def toDF(
- self: RDD[RowLike],
- schema: Optional[Union[List[str], Tuple[str, ...]]] = ...,
- sampleRatio: Optional[float] = ...,
- ) -> DataFrame: ...
- @overload
- def toDF(self: RDD[RowLike], schema: Optional[Union[StructType, str]] =
...) -> DataFrame: ...
- @overload
- def toDF(
- self: RDD[AtomicValue],
- schema: Union[AtomicType, str],
- ) -> DataFrame: ...
-
-class RDDBarrier(Generic[T]):
- rdd: RDD[T]
- def __init__(self, rdd: RDD[T]) -> None: ...
- def mapPartitions(
- self, f: Callable[[Iterable[T]], Iterable[U]], preservesPartitioning:
bool = ...
- ) -> RDD[U]: ...
- def mapPartitionsWithIndex(
- self,
- f: Callable[[int, Iterable[T]], Iterable[U]],
- preservesPartitioning: bool = ...,
- ) -> RDD[U]: ...
-
-class PipelinedRDD(RDD[U], Generic[T, U]):
- func: Callable[[T], U]
- preservesPartitioning: bool
- is_cached: bool
- is_checkpointed: bool
- ctx: pyspark.context.SparkContext
- prev: RDD[T]
- partitioner: Optional[Partitioner]
- is_barrier: bool
- def __init__(
- self,
- prev: RDD[T],
- func: Callable[[Iterable[T]], Iterable[U]],
- preservesPartitioning: bool = ...,
- isFromBarrier: bool = ...,
- ) -> None: ...
- def getNumPartitions(self) -> int: ...
- def id(self) -> int: ...
diff --git a/python/pyspark/serializers.py b/python/pyspark/serializers.py
index a0941afd..628ef18 100644
--- a/python/pyspark/serializers.py
+++ b/python/pyspark/serializers.py
@@ -100,6 +100,13 @@ class Serializer:
"""
raise NotImplementedError
+ def dumps(self, obj):
+ """
+ Serialize an object into a byte array.
+ When batching is used, this will be called with an array of objects.
+ """
+ raise NotImplementedError
+
def _load_stream_without_unbatching(self, stream):
"""
Return an iterator of deserialized batches (iterable) of objects from
the input stream.
diff --git a/python/pyspark/sql/_typing.pyi b/python/pyspark/sql/_typing.pyi
index 2adae6c..209bb70 100644
--- a/python/pyspark/sql/_typing.pyi
+++ b/python/pyspark/sql/_typing.pyi
@@ -25,7 +25,7 @@ from typing import (
TypeVar,
Union,
)
-from typing_extensions import Protocol
+from typing_extensions import Literal, Protocol
import datetime
import decimal
@@ -56,6 +56,8 @@ AtomicValue = TypeVar(
RowLike = TypeVar("RowLike", List[Any], Tuple[Any, ...], pyspark.sql.types.Row)
+SQLBatchedUDFType = Literal[100]
+
class SupportsOpen(Protocol):
def open(self, partition_id: int, epoch_id: int) -> bool: ...
diff --git a/python/pyspark/sql/pandas/_typing/__init__.pyi
b/python/pyspark/sql/pandas/_typing/__init__.pyi
index d3796f4..906703d 100644
--- a/python/pyspark/sql/pandas/_typing/__init__.pyi
+++ b/python/pyspark/sql/pandas/_typing/__init__.pyi
@@ -42,11 +42,12 @@ DataFrameOrSeriesLike_ = TypeVar("DataFrameOrSeriesLike_",
bound=DataFrameOrSeri
# UDF annotations
PandasScalarUDFType = Literal[200]
-PandasScalarIterUDFType = Literal[204]
PandasGroupedMapUDFType = Literal[201]
-PandasCogroupedMapUDFType = Literal[206]
PandasGroupedAggUDFType = Literal[202]
+PandasWindowAggUDFType = Literal[203]
+PandasScalarIterUDFType = Literal[204]
PandasMapIterUDFType = Literal[205]
+PandasCogroupedMapUDFType = Literal[206]
ArrowMapIterUDFType = Literal[207]
class PandasVariadicScalarToScalarFunction(Protocol):
diff --git a/python/pyspark/tests/typing/test_rdd.yml
b/python/pyspark/tests/typing/test_rdd.yml
index 749ad53..4896582 100644
--- a/python/pyspark/tests/typing/test_rdd.yml
+++ b/python/pyspark/tests/typing/test_rdd.yml
@@ -18,11 +18,11 @@
- case: toDF
main: |
from pyspark.sql.types import (
- IntegerType,
- Row,
- StructType,
- StringType,
- StructField,
+ IntegerType,
+ Row,
+ StructType,
+ StringType,
+ StructField,
)
from collections import namedtuple
from pyspark.sql import SparkSession
@@ -60,3 +60,70 @@
rdd_named_tuple.toDF(sampleRatio=0.4)
rdd_named_tuple.toDF(["a", "b"], sampleRatio=0.4)
rdd_named_tuple.toDF(struct)
+
+
+- case: rddMethods
+ main: |
+ from operator import add
+ from typing import Iterable, Set, Tuple
+ from pyspark.sql import SparkSession
+
+ spark = SparkSession.builder.getOrCreate()
+ sc = spark.sparkContext
+
+ def f1(x: int) -> str:
+ return str(x)
+
+ reveal_type(sc.range(10).map(f1))
+
+ def f2(x: int) -> Iterable[int]:
+ return range(x)
+
+ reveal_type(sc.range(10).flatMap(f2))
+
+ reveal_type(sc.parallelize([("a", 1), ("b", 0)]).filter(lambda x: x[1] !=
0))
+
+ reveal_type(sc.parallelize([("a", 1), ("b", 0)]).max())
+
+ reveal_type(sc.range(10).reduce(add))
+
+ def seq_func(xs: Set[str], x: int) -> Set[str]:
+ xs.add(str(x % 11))
+ return xs
+
+ def comb_func(xs: Set[str], ys: Set[str]) -> Set[str]:
+ xs.update(ys)
+ return xs
+
+ zero: Set[str] = set()
+
+ reveal_type(sc.parallelize([("a", 1)]).aggregateByKey(zero, seq_func,
comb_func))
+
+ out: |
+ main:11: note: Revealed type is "pyspark.rdd.RDD[builtins.str*]"
+ main:16: note: Revealed type is "pyspark.rdd.RDD[builtins.int*]"
+ main:18: note: Revealed type is "pyspark.rdd.RDD[Tuple[builtins.str,
builtins.int]]"
+ main:20: note: Revealed type is "Tuple[builtins.str, builtins.int]"
+ main:22: note: Revealed type is "builtins.int"
+ main:34: note: Revealed type is "pyspark.rdd.RDD[Tuple[builtins.str,
builtins.set[builtins.str]]]"
+
+- case: rddMethodsErrors
+ main: |
+ from pyspark.sql import SparkSession
+
+ spark = SparkSession.builder.getOrCreate()
+ sc = spark.sparkContext
+
+ def f1(x: str) -> str:
+ return x
+
+ sc.range(10).map(f1)
+
+ def f2(x: int) -> str:
+ return str(x)
+
+ sc.range(10).reduce(f2)
+
+ out: |
+ main:9: error: Argument 1 to "map" of "RDD" has incompatible type
"Callable[[str], str]"; expected "Callable[[int], str]" [arg-type]
+ main:14: error: Argument 1 to "reduce" of "RDD" has incompatible type
"Callable[[int], str]"; expected "Callable[[int, int], int]" [arg-type]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
