snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/streaming/dstream.py

@@ -0,0 +1,933 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import operator
+import time
+from itertools import chain
+from datetime import datetime
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Hashable,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    TYPE_CHECKING,
+    cast,
+    overload,
+)
+
+from py4j.protocol import Py4JJavaError
+
+from pyspark.storagelevel import StorageLevel
+from pyspark.streaming.util import rddToFileName, TransformFunction
+from pyspark.rdd import portable_hash, RDD
+from pyspark.resultiterable import ResultIterable
+from py4j.java_gateway import JavaObject
+
+if TYPE_CHECKING:
+    from pyspark.serializers import Serializer
+    from pyspark.streaming.context import StreamingContext
+
+__all__ = ["DStream"]
+
+S = TypeVar("S")
+T = TypeVar("T")
+T_co = TypeVar("T_co", covariant=True)
+U = TypeVar("U")
+K = TypeVar("K", bound=Hashable)
+V = TypeVar("V")
+
+
+class DStream(Generic[T_co]):
+    """
+    A Discretized Stream (DStream), the basic abstraction in Spark Streaming,
+    is a continuous sequence of RDDs (of the same type) representing a
+    continuous stream of data (see :class:`RDD` in the Spark core documentation
+    for more details on RDDs).
+
+    DStreams can either be created from live data (such as, data from TCP
+    sockets, etc.) using a :class:`StreamingContext` or it can be
+    generated by transforming existing DStreams using operations such as
+    `map`, `window` and `reduceByKeyAndWindow`. While a Spark Streaming
+    program is running, each DStream periodically generates a RDD, either
+    from live data or by transforming the RDD generated by a parent DStream.
+
+    DStreams internally is characterized by a few basic properties:
+     - A list of other DStreams that the DStream depends on
+     - A time interval at which the DStream generates an RDD
+     - A function that is used to generate an RDD after each time interval
+    """
+
+    def __init__(
+        self,
+        jdstream: JavaObject,
+        ssc: "StreamingContext",
+        jrdd_deserializer: "Serializer",
+    ):
+        self._jdstream = jdstream
+        self._ssc = ssc
+        self._sc = ssc._sc
+        self._jrdd_deserializer = jrdd_deserializer
+        self.is_cached = False
+        self.is_checkpointed = False
+
+    def context(self) -> "StreamingContext":
+        """
+        Return the StreamingContext associated with this DStream
+        """
+        return self._ssc
+
+    def count(self) -> "DStream[int]":
+        """
+        Return a new DStream in which each RDD has a single element
+        generated by counting each RDD of this DStream.
+        """
+        return self.mapPartitions(lambda i: [sum(1 for _ in i)]).reduce(operator.add)
+
+    def filter(self: "DStream[T]", f: Callable[[T], bool]) -> "DStream[T]":
+        """
+        Return a new DStream containing only the elements that satisfy predicate.
+        """
+
+        def func(iterator: Iterable[T]) -> Iterable[T]:
+            return filter(f, iterator)
+
+        return self.mapPartitions(func, True)
+
+    def flatMap(
+        self: "DStream[T]",
+        f: Callable[[T], Iterable[U]],
+        preservesPartitioning: bool = False,
+    ) -> "DStream[U]":
+        """
+        Return a new DStream by applying a function to all elements of
+        this DStream, and then flattening the results
+        """
+
+        def func(s: int, iterator: Iterable[T]) -> Iterable[U]:
+            return chain.from_iterable(map(f, iterator))
+
+        return self.mapPartitionsWithIndex(func, preservesPartitioning)
+
+    def map(
+        self: "DStream[T]", f: Callable[[T], U], preservesPartitioning: bool = False
+    ) -> "DStream[U]":
+        """
+        Return a new DStream by applying a function to each element of DStream.
+        """
+
+        def func(iterator: Iterable[T]) -> Iterable[U]:
+            return map(f, iterator)
+
+        return self.mapPartitions(func, preservesPartitioning)
+
+    def mapPartitions(
+        self: "DStream[T]",
+        f: Callable[[Iterable[T]], Iterable[U]],
+        preservesPartitioning: bool = False,
+    ) -> "DStream[U]":
+        """
+        Return a new DStream in which each RDD is generated by applying
+        mapPartitions() to each RDDs of this DStream.
+        """
+
+        def func(s: int, iterator: Iterable[T]) -> Iterable[U]:
+            return f(iterator)
+
+        return self.mapPartitionsWithIndex(func, preservesPartitioning)
+
+    def mapPartitionsWithIndex(
+        self: "DStream[T]",
+        f: Callable[[int, Iterable[T]], Iterable[U]],
+        preservesPartitioning: bool = False,
+    ) -> "DStream[U]":
+        """
+        Return a new DStream in which each RDD is generated by applying
+        mapPartitionsWithIndex() to each RDDs of this DStream.
+        """
+        return self.transform(lambda rdd: rdd.mapPartitionsWithIndex(f, preservesPartitioning))
+
+    def reduce(self: "DStream[T]", func: Callable[[T, T], T]) -> "DStream[T]":
+        """
+        Return a new DStream in which each RDD has a single element
+        generated by reducing each RDD of this DStream.
+        """
+        return self.map(lambda x: (None, x)).reduceByKey(func, 1).map(lambda x: x[1])
+
+    def reduceByKey(
+        self: "DStream[Tuple[K, V]]",
+        func: Callable[[V, V], V],
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[K, V]]":
+        """
+        Return a new DStream by applying reduceByKey to each RDD.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+        return self.combineByKey(lambda x: x, func, func, numPartitions)
+
+    def combineByKey(
+        self: "DStream[Tuple[K, V]]",
+        createCombiner: Callable[[V], U],
+        mergeValue: Callable[[U, V], U],
+        mergeCombiners: Callable[[U, U], U],
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[K, U]]":
+        """
+        Return a new DStream by applying combineByKey to each RDD.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+
+        def func(rdd: RDD[Tuple[K, V]]) -> RDD[Tuple[K, U]]:
+            return rdd.combineByKey(createCombiner, mergeValue, mergeCombiners, numPartitions)
+
+        return self.transform(func)
+
+    def partitionBy(
+        self: "DStream[Tuple[K, V]]",
+        numPartitions: int,
+        partitionFunc: Callable[[K], int] = portable_hash,
+    ) -> "DStream[Tuple[K, V]]":
+        """
+        Return a copy of the DStream in which each RDD are partitioned
+        using the specified partitioner.
+        """
+        return self.transform(lambda rdd: rdd.partitionBy(numPartitions, partitionFunc))
+
+    @overload
+    def foreachRDD(self: "DStream[T]", func: Callable[[RDD[T]], None]) -> None:
+        ...
+
+    @overload
+    def foreachRDD(self: "DStream[T]", func: Callable[[datetime, RDD[T]], None]) -> None:
+        ...
+
+    def foreachRDD(
+        self: "DStream[T]",
+        func: Union[Callable[[RDD[T]], None], Callable[[datetime, RDD[T]], None]],
+    ) -> None:
+        """
+        Apply a function to each RDD in this DStream.
+        """
+        if func.__code__.co_argcount == 1:
+            old_func = func
+
+            def func(_: datetime, rdd: "RDD[T]") -> None:
+                return old_func(rdd)  # type: ignore[call-arg, arg-type]
+
+        jfunc = TransformFunction(self._sc, func, self._jrdd_deserializer)
+        assert self._ssc._jvm is not None
+        api = self._ssc._jvm.PythonDStream
+        api.callForeachRDD(self._jdstream, jfunc)
+
+    def pprint(self, num: int = 10) -> None:
+        """
+        Print the first num elements of each RDD generated in this DStream.
+
+        Parameters
+        ----------
+        num : int, optional
+            the number of elements from the first will be printed.
+        """
+
+        def takeAndPrint(time: datetime, rdd: RDD[T]) -> None:
+            taken = rdd.take(num + 1)
+            print("-------------------------------------------")
+            print("Time: %s" % time)
+            print("-------------------------------------------")
+            for record in taken[:num]:
+                print(record)
+            if len(taken) > num:
+                print("...")
+            print("")
+
+        self.foreachRDD(takeAndPrint)
+
+    def mapValues(self: "DStream[Tuple[K, V]]", f: Callable[[V], U]) -> "DStream[Tuple[K, U]]":
+        """
+        Return a new DStream by applying a map function to the value of
+        each key-value pairs in this DStream without changing the key.
+        """
+
+        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 flatMapValues(
+        self: "DStream[Tuple[K, V]]", f: Callable[[V], Iterable[U]]
+    ) -> "DStream[Tuple[K, U]]":
+        """
+        Return a new DStream by applying a flatmap function to the value
+        of each key-value pairs in this DStream without changing the key.
+        """
+
+        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 glom(self: "DStream[T]") -> "DStream[List[T]]":
+        """
+        Return a new DStream in which RDD is generated by applying glom()
+        to RDD of this DStream.
+        """
+
+        def func(iterator: Iterable[T]) -> Iterable[List[T]]:
+            yield list(iterator)
+
+        return self.mapPartitions(func)
+
+    def cache(self: "DStream[T]") -> "DStream[T]":
+        """
+        Persist the RDDs of this DStream with the default storage level
+        (`MEMORY_ONLY`).
+        """
+        self.is_cached = True
+        self.persist(StorageLevel.MEMORY_ONLY)
+        return self
+
+    def persist(self: "DStream[T]", storageLevel: StorageLevel) -> "DStream[T]":
+        """
+        Persist the RDDs of this DStream with the given storage level
+        """
+        self.is_cached = True
+        javaStorageLevel = self._sc._getJavaStorageLevel(storageLevel)
+        self._jdstream.persist(javaStorageLevel)
+        return self
+
+    def checkpoint(self: "DStream[T]", interval: int) -> "DStream[T]":
+        """
+        Enable periodic checkpointing of RDDs of this DStream
+
+        Parameters
+        ----------
+        interval : int
+            time in seconds, after each period of that, generated
+            RDD will be checkpointed
+        """
+        self.is_checkpointed = True
+        self._jdstream.checkpoint(self._ssc._jduration(interval))
+        return self
+
+    def groupByKey(
+        self: "DStream[Tuple[K, V]]", numPartitions: Optional[int] = None
+    ) -> "DStream[Tuple[K, Iterable[V]]]":
+        """
+        Return a new DStream by applying groupByKey on each RDD.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+        return self.transform(lambda rdd: rdd.groupByKey(numPartitions))
+
+    def countByValue(self: "DStream[K]") -> "DStream[Tuple[K, int]]":
+        """
+        Return a new DStream in which each RDD contains the counts of each
+        distinct value in each RDD of this DStream.
+        """
+        return self.map(lambda x: (x, 1)).reduceByKey(lambda x, y: x + y)
+
+    def saveAsTextFiles(self, prefix: str, suffix: Optional[str] = None) -> None:
+        """
+        Save each RDD in this DStream as at text file, using string
+        representation of elements.
+        """
+
+        def saveAsTextFile(t: Optional[datetime], rdd: RDD[T]) -> None:
+            path = rddToFileName(prefix, suffix, t)
+            try:
+                rdd.saveAsTextFile(path)
+            except Py4JJavaError as e:
+                # after recovered from checkpointing, the foreachRDD may
+                # be called twice
+                if "FileAlreadyExistsException" not in str(e):
+                    raise
+
+        return self.foreachRDD(saveAsTextFile)
+
+    # TODO: uncomment this until we have ssc.pickleFileStream()
+    # def saveAsPickleFiles(self, prefix, suffix=None):
+    #     """
+    #     Save each RDD in this DStream as at binary file, the elements are
+    #     serialized by pickle.
+    #     """
+    #     def saveAsPickleFile(t, rdd):
+    #         path = rddToFileName(prefix, suffix, t)
+    #         try:
+    #             rdd.saveAsPickleFile(path)
+    #         except Py4JJavaError as e:
+    #             # after recovered from checkpointing, the foreachRDD may
+    #             # be called twice
+    #             if 'FileAlreadyExistsException' not in str(e):
+    #                 raise
+    #     return self.foreachRDD(saveAsPickleFile)
+
+    @overload
+    def transform(self: "DStream[T]", func: Callable[[RDD[T]], RDD[U]]) -> "TransformedDStream[U]":
+        ...
+
+    @overload
+    def transform(
+        self: "DStream[T]", func: Callable[[datetime, RDD[T]], RDD[U]]
+    ) -> "TransformedDStream[U]":
+        ...
+
+    def transform(
+        self: "DStream[T]",
+        func: Union[Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]]],
+    ) -> "TransformedDStream[U]":
+        """
+        Return a new DStream in which each RDD is generated by applying a function
+        on each RDD of this DStream.
+
+        `func` can have one argument of `rdd`, or have two arguments of
+        (`time`, `rdd`)
+        """
+        if func.__code__.co_argcount == 1:
+            oldfunc = func
+
+            def func(_: datetime, rdd: RDD[T]) -> RDD[U]:
+                return oldfunc(rdd)  # type: ignore[arg-type, call-arg]
+
+        assert func.__code__.co_argcount == 2, "func should take one or two arguments"
+        return TransformedDStream(self, func)
+
+    @overload
+    def transformWith(
+        self: "DStream[T]",
+        func: Callable[[RDD[T], RDD[U]], RDD[V]],
+        other: "DStream[U]",
+        keepSerializer: bool = ...,
+    ) -> "DStream[V]":
+        ...
+
+    @overload
+    def transformWith(
+        self: "DStream[T]",
+        func: Callable[[datetime, RDD[T], RDD[U]], RDD[V]],
+        other: "DStream[U]",
+        keepSerializer: bool = ...,
+    ) -> "DStream[V]":
+        ...
+
+    def transformWith(
+        self: "DStream[T]",
+        func: Union[
+            Callable[[RDD[T], RDD[U]], RDD[V]],
+            Callable[[datetime, RDD[T], RDD[U]], RDD[V]],
+        ],
+        other: "DStream[U]",
+        keepSerializer: bool = False,
+    ) -> "DStream[V]":
+        """
+        Return a new DStream in which each RDD is generated by applying a function
+        on each RDD of this DStream and 'other' DStream.
+
+        `func` can have two arguments of (`rdd_a`, `rdd_b`) or have three
+        arguments of (`time`, `rdd_a`, `rdd_b`)
+        """
+        if func.__code__.co_argcount == 2:
+            oldfunc = func
+
+            def func(_: datetime, a: RDD[T], b: RDD[U]) -> RDD[V]:
+                return oldfunc(a, b)  # type: ignore[call-arg, arg-type]
+
+        assert func.__code__.co_argcount == 3, "func should take two or three arguments"
+        jfunc = TransformFunction(
+            self._sc,
+            func,
+            self._jrdd_deserializer,
+            other._jrdd_deserializer,
+        )
+        assert self._sc._jvm is not None
+        dstream = self._sc._jvm.PythonTransformed2DStream(
+            self._jdstream.dstream(), other._jdstream.dstream(), jfunc
+        )
+        jrdd_serializer = self._jrdd_deserializer if keepSerializer else self._sc.serializer
+        return DStream(dstream.asJavaDStream(), self._ssc, jrdd_serializer)
+
+    def repartition(self: "DStream[T]", numPartitions: int) -> "DStream[T]":
+        """
+        Return a new DStream with an increased or decreased level of parallelism.
+        """
+        return self.transform(lambda rdd: rdd.repartition(numPartitions))
+
+    @property
+    def _slideDuration(self) -> None:
+        """
+        Return the slideDuration in seconds of this DStream
+        """
+        return self._jdstream.dstream().slideDuration().milliseconds() / 1000.0
+
+    def union(self: "DStream[T]", other: "DStream[U]") -> "DStream[Union[T, U]]":
+        """
+        Return a new DStream by unifying data of another DStream with this DStream.
+
+        Parameters
+        ----------
+        other : :class:`DStream`
+            Another DStream having the same interval (i.e., slideDuration)
+            as this DStream.
+        """
+        if self._slideDuration != other._slideDuration:
+            raise ValueError("the two DStream should have same slide duration")
+        return self.transformWith(lambda a, b: a.union(b), other, True)
+
+    def cogroup(
+        self: "DStream[Tuple[K, V]]",
+        other: "DStream[Tuple[K, U]]",
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[K, Tuple[ResultIterable[V], ResultIterable[U]]]]":
+        """
+        Return a new DStream by applying 'cogroup' between RDDs of this
+        DStream and `other` DStream.
+
+        Hash partitioning is used to generate the RDDs with `numPartitions` partitions.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+        return self.transformWith(
+            lambda a, b: a.cogroup(b, numPartitions),
+            other,
+        )
+
+    def join(
+        self: "DStream[Tuple[K, V]]",
+        other: "DStream[Tuple[K, U]]",
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[K, Tuple[V, U]]]":
+        """
+        Return a new DStream by applying 'join' between RDDs of this DStream and
+        `other` DStream.
+
+        Hash partitioning is used to generate the RDDs with `numPartitions`
+        partitions.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+        return self.transformWith(lambda a, b: a.join(b, numPartitions), other)
+
+    def leftOuterJoin(
+        self: "DStream[Tuple[K, V]]",
+        other: "DStream[Tuple[K, U]]",
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[K, Tuple[V, Optional[U]]]]":
+        """
+        Return a new DStream by applying 'left outer join' between RDDs of this DStream and
+        `other` DStream.
+
+        Hash partitioning is used to generate the RDDs with `numPartitions`
+        partitions.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+        return self.transformWith(lambda a, b: a.leftOuterJoin(b, numPartitions), other)
+
+    def rightOuterJoin(
+        self: "DStream[Tuple[K, V]]",
+        other: "DStream[Tuple[K, U]]",
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[K, Tuple[Optional[V], U]]]":
+        """
+        Return a new DStream by applying 'right outer join' between RDDs of this DStream and
+        `other` DStream.
+
+        Hash partitioning is used to generate the RDDs with `numPartitions`
+        partitions.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+        return self.transformWith(lambda a, b: a.rightOuterJoin(b, numPartitions), other)
+
+    def fullOuterJoin(
+        self: "DStream[Tuple[K, V]]",
+        other: "DStream[Tuple[K, U]]",
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[K, Tuple[Optional[V], Optional[U]]]]":
+        """
+        Return a new DStream by applying 'full outer join' between RDDs of this DStream and
+        `other` DStream.
+
+        Hash partitioning is used to generate the RDDs with `numPartitions`
+        partitions.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+        return self.transformWith(lambda a, b: a.fullOuterJoin(b, numPartitions), other)
+
+    def _jtime(self, timestamp: Union[datetime, int, float]) -> JavaObject:
+        """Convert datetime or unix_timestamp into Time"""
+        if isinstance(timestamp, datetime):
+            timestamp = time.mktime(timestamp.timetuple())
+        assert self._sc._jvm is not None
+        return self._sc._jvm.Time(int(timestamp * 1000))
+
+    def slice(self, begin: Union[datetime, int], end: Union[datetime, int]) -> List[RDD[T]]:
+        """
+        Return all the RDDs between 'begin' to 'end' (both included)
+
+        `begin`, `end` could be datetime.datetime() or unix_timestamp
+        """
+        jrdds = self._jdstream.slice(self._jtime(begin), self._jtime(end))
+        return [RDD(jrdd, self._sc, self._jrdd_deserializer) for jrdd in jrdds]
+
+    def _validate_window_param(self, window: int, slide: Optional[int]) -> None:
+        duration = self._jdstream.dstream().slideDuration().milliseconds()
+        if int(window * 1000) % duration != 0:
+            raise ValueError(
+                "windowDuration must be multiple of the parent "
+                "dstream's slide (batch) duration (%d ms)" % duration
+            )
+        if slide and int(slide * 1000) % duration != 0:
+            raise ValueError(
+                "slideDuration must be multiple of the parent "
+                "dstream's slide (batch) duration (%d ms)" % duration
+            )
+
+    def window(self, windowDuration: int, slideDuration: Optional[int] = None) -> "DStream[T]":
+        """
+        Return a new DStream in which each RDD contains all the elements in seen in a
+        sliding window of time over this DStream.
+
+        Parameters
+        ----------
+        windowDuration : int
+            width of the window; must be a multiple of this DStream's
+            batching interval
+        slideDuration : int, optional
+            sliding interval of the window (i.e., the interval after which
+            the new DStream will generate RDDs); must be a multiple of this
+            DStream's batching interval
+        """
+        self._validate_window_param(windowDuration, slideDuration)
+        d = self._ssc._jduration(windowDuration)
+        if slideDuration is None:
+            return DStream(self._jdstream.window(d), self._ssc, self._jrdd_deserializer)
+        s = self._ssc._jduration(slideDuration)
+        return DStream(self._jdstream.window(d, s), self._ssc, self._jrdd_deserializer)
+
+    def reduceByWindow(
+        self: "DStream[T]",
+        reduceFunc: Callable[[T, T], T],
+        invReduceFunc: Optional[Callable[[T, T], T]],
+        windowDuration: int,
+        slideDuration: int,
+    ) -> "DStream[T]":
+        """
+        Return a new DStream in which each RDD has a single element generated by reducing all
+        elements in a sliding window over this DStream.
+
+        if `invReduceFunc` is not None, the reduction is done incrementally
+        using the old window's reduced value :
+
+        1. reduce the new values that entered the window (e.g., adding new counts)
+
+        2. "inverse reduce" the old values that left the window (e.g., subtracting old counts)
+        This is more efficient than `invReduceFunc` is None.
+
+        Parameters
+        ----------
+        reduceFunc : function
+            associative and commutative reduce function
+        invReduceFunc : function
+            inverse reduce function of `reduceFunc`; such that for all y,
+            and invertible x:
+            `invReduceFunc(reduceFunc(x, y), x) = y`
+        windowDuration : int
+            width of the window; must be a multiple of this DStream's
+            batching interval
+        slideDuration : int
+            sliding interval of the window (i.e., the interval after which
+            the new DStream will generate RDDs); must be a multiple of this
+            DStream's batching interval
+        """
+        keyed = self.map(lambda x: (1, x))
+        reduced = keyed.reduceByKeyAndWindow(
+            reduceFunc, invReduceFunc, windowDuration, slideDuration, 1
+        )
+        return reduced.map(lambda kv: kv[1])
+
+    def countByWindow(
+        self: "DStream[T]", windowDuration: int, slideDuration: int
+    ) -> "DStream[int]":
+        """
+        Return a new DStream in which each RDD has a single element generated
+        by counting the number of elements in a window over this DStream.
+        windowDuration and slideDuration are as defined in the window() operation.
+
+        This is equivalent to window(windowDuration, slideDuration).count(),
+        but will be more efficient if window is large.
+        """
+        return self.map(lambda x: 1).reduceByWindow(
+            operator.add, operator.sub, windowDuration, slideDuration
+        )
+
+    def countByValueAndWindow(
+        self: "DStream[T]",
+        windowDuration: int,
+        slideDuration: int,
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[T, int]]":
+        """
+        Return a new DStream in which each RDD contains the count of distinct elements in
+        RDDs in a sliding window over this DStream.
+
+        Parameters
+        ----------
+        windowDuration : int
+            width of the window; must be a multiple of this DStream's
+            batching interval
+        slideDuration : int
+            sliding interval of the window (i.e., the interval after which
+            the new DStream will generate RDDs); must be a multiple of this
+            DStream's batching interval
+        numPartitions : int, optional
+            number of partitions of each RDD in the new DStream.
+        """
+        keyed = self.map(lambda x: (x, 1))
+        counted = keyed.reduceByKeyAndWindow(
+            operator.add, operator.sub, windowDuration, slideDuration, numPartitions
+        )
+        return counted.filter(lambda kv: kv[1] > 0)
+
+    def groupByKeyAndWindow(
+        self: "DStream[Tuple[K, V]]",
+        windowDuration: int,
+        slideDuration: int,
+        numPartitions: Optional[int] = None,
+    ) -> "DStream[Tuple[K, Iterable[V]]]":
+        """
+        Return a new DStream by applying `groupByKey` over a sliding window.
+        Similar to `DStream.groupByKey()`, but applies it over a sliding window.
+
+        Parameters
+        ----------
+        windowDuration : int
+            width of the window; must be a multiple of this DStream's
+            batching interval
+        slideDuration : int
+            sliding interval of the window (i.e., the interval after which
+            the new DStream will generate RDDs); must be a multiple of this
+            DStream's batching interval
+        numPartitions : int, optional
+            Number of partitions of each RDD in the new DStream.
+        """
+        ls = self.mapValues(lambda x: [x])
+        grouped = ls.reduceByKeyAndWindow(
+            lambda a, b: a.extend(b) or a,  # type: ignore[func-returns-value]
+            lambda a, b: a[len(b) :],
+            windowDuration,
+            slideDuration,
+            numPartitions,
+        )
+        return grouped.mapValues(ResultIterable)
+
+    def reduceByKeyAndWindow(
+        self: "DStream[Tuple[K, V]]",
+        func: Callable[[V, V], V],
+        invFunc: Optional[Callable[[V, V], V]],
+        windowDuration: int,
+        slideDuration: Optional[int] = None,
+        numPartitions: Optional[int] = None,
+        filterFunc: Optional[Callable[[Tuple[K, V]], bool]] = None,
+    ) -> "DStream[Tuple[K, V]]":
+        """
+        Return a new DStream by applying incremental `reduceByKey` over a sliding window.
+
+        The reduced value of over a new window is calculated using the old window's reduce value :
+         1. reduce the new values that entered the window (e.g., adding new counts)
+         2. "inverse reduce" the old values that left the window (e.g., subtracting old counts)
+
+        `invFunc` can be None, then it will reduce all the RDDs in window, could be slower
+        than having `invFunc`.
+
+        Parameters
+        ----------
+        func : function
+            associative and commutative reduce function
+        invFunc : function
+            inverse function of `reduceFunc`
+        windowDuration : int
+            width of the window; must be a multiple of this DStream's
+            batching interval
+        slideDuration : int, optional
+            sliding interval of the window (i.e., the interval after which
+            the new DStream will generate RDDs); must be a multiple of this
+            DStream's batching interval
+        numPartitions : int, optional
+            number of partitions of each RDD in the new DStream.
+        filterFunc : function, optional
+            function to filter expired key-value pairs;
+            only pairs that satisfy the function are retained
+            set this to null if you do not want to filter
+        """
+        self._validate_window_param(windowDuration, slideDuration)
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+
+        reduced = self.reduceByKey(func, numPartitions)
+
+        if invFunc:
+
+            def reduceFunc(t: datetime, a: Any, b: Any) -> Any:
+                b = b.reduceByKey(func, numPartitions)
+                r = a.union(b).reduceByKey(func, numPartitions) if a else b
+                if filterFunc:
+                    r = r.filter(filterFunc)
+                return r
+
+            def invReduceFunc(t: datetime, a: Any, b: Any) -> Any:
+                b = b.reduceByKey(func, numPartitions)
+                joined = a.leftOuterJoin(b, numPartitions)
+                return joined.mapValues(
+                    lambda kv: invFunc(kv[0], kv[1])  # type: ignore[misc]
+                    if kv[1] is not None
+                    else kv[0]
+                )
+
+            jreduceFunc = TransformFunction(self._sc, reduceFunc, reduced._jrdd_deserializer)
+            jinvReduceFunc = TransformFunction(self._sc, invReduceFunc, reduced._jrdd_deserializer)
+            if slideDuration is None:
+                slideDuration = self._slideDuration
+            assert self._sc._jvm is not None
+            dstream = self._sc._jvm.PythonReducedWindowedDStream(
+                reduced._jdstream.dstream(),
+                jreduceFunc,
+                jinvReduceFunc,
+                self._ssc._jduration(windowDuration),
+                self._ssc._jduration(slideDuration),  # type: ignore[arg-type]
+            )
+            return DStream(dstream.asJavaDStream(), self._ssc, self._sc.serializer)
+        else:
+            return reduced.window(windowDuration, slideDuration).reduceByKey(
+                func, numPartitions  # type: ignore[arg-type]
+            )
+
+    def updateStateByKey(
+        self: "DStream[Tuple[K, V]]",
+        updateFunc: Callable[[Iterable[V], Optional[S]], S],
+        numPartitions: Optional[int] = None,
+        initialRDD: Optional[Union[RDD[Tuple[K, S]], Iterable[Tuple[K, S]]]] = None,
+    ) -> "DStream[Tuple[K, S]]":
+        """
+        Return a new "state" DStream where the state for each key is updated by applying
+        the given function on the previous state of the key and the new values of the key.
+
+        Parameters
+        ----------
+        updateFunc : function
+            State update function. If this function returns None, then
+            corresponding state key-value pair will be eliminated.
+        """
+        if numPartitions is None:
+            numPartitions = self._sc.defaultParallelism
+
+        if initialRDD and not isinstance(initialRDD, RDD):
+            initialRDD = self._sc.parallelize(initialRDD)
+
+        def reduceFunc(t: datetime, a: Any, b: Any) -> Any:
+            if a is None:
+                g = b.groupByKey(numPartitions).mapValues(lambda vs: (list(vs), None))
+            else:
+                g = a.cogroup(b.partitionBy(cast(int, numPartitions)), numPartitions)
+                g = g.mapValues(lambda ab: (list(ab[1]), list(ab[0])[0] if len(ab[0]) else None))
+            state = g.mapValues(lambda vs_s: updateFunc(vs_s[0], vs_s[1]))
+            return state.filter(lambda k_v: k_v[1] is not None)
+
+        jreduceFunc = TransformFunction(
+            self._sc,
+            reduceFunc,
+            self._sc.serializer,
+            self._jrdd_deserializer,
+        )
+        if initialRDD:
+            initialRDD = cast(RDD[Tuple[K, S]], initialRDD)._reserialize(self._jrdd_deserializer)
+            assert self._sc._jvm is not None
+            dstream = self._sc._jvm.PythonStateDStream(
+                self._jdstream.dstream(),
+                jreduceFunc,
+                initialRDD._jrdd,
+            )
+        else:
+            assert self._sc._jvm is not None
+            dstream = self._sc._jvm.PythonStateDStream(self._jdstream.dstream(), jreduceFunc)
+
+        return DStream(dstream.asJavaDStream(), self._ssc, self._sc.serializer)
+
+
+class TransformedDStream(DStream[U]):
+    """
+    TransformedDStream is a DStream generated by an Python function
+    transforming each RDD of a DStream to another RDDs.
+
+    Multiple continuous transformations of DStream can be combined into
+    one transformation.
+    """
+
+    @overload
+    def __init__(self: DStream[U], prev: DStream[T], func: Callable[[RDD[T]], RDD[U]]):
+        ...
+
+    @overload
+    def __init__(
+        self: DStream[U],
+        prev: DStream[T],
+        func: Callable[[datetime, RDD[T]], RDD[U]],
+    ):
+        ...
+
+    def __init__(
+        self,
+        prev: DStream[T],
+        func: Union[Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]]],
+    ):
+        self._ssc = prev._ssc
+        self._sc = self._ssc._sc
+        self._jrdd_deserializer = self._sc.serializer
+        self.is_cached = False
+        self.is_checkpointed = False
+        self._jdstream_val = None
+
+        # Using type() to avoid folding the functions and compacting the DStreams which is not
+        # not strictly an object of TransformedDStream.
+        if type(prev) is TransformedDStream and not prev.is_cached and not prev.is_checkpointed:
+            prev_func: Callable = prev.func
+            func = cast(Callable[[datetime, RDD[T]], RDD[U]], func)
+            self.func: Union[
+                Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]]
+            ] = lambda t, rdd: func(t, prev_func(t, rdd))
+            self.prev: DStream[T] = prev.prev
+        else:
+            self.prev = prev
+            self.func = func
+
+    @property
+    def _jdstream(self) -> JavaObject:
+        if self._jdstream_val is not None:
+            return self._jdstream_val
+
+        jfunc = TransformFunction(self._sc, self.func, self.prev._jrdd_deserializer)
+        assert self._sc._jvm is not None
+        dstream = self._sc._jvm.PythonTransformedDStream(self.prev._jdstream.dstream(), jfunc)
+        self._jdstream_val = dstream.asJavaDStream()
+        return self._jdstream_val