################################################################################
#  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 typing
import uuid
from enum import Enum
from typing import Callable, Union, List, cast, Optional, overload

from pyflink.util.java_utils import get_j_env_configuration

from pyflink.common import typeinfo, ExecutionConfig, Row
from pyflink.common.typeinfo import RowTypeInfo, Types, TypeInformation, _from_java_type
from pyflink.common.watermark_strategy import WatermarkStrategy, TimestampAssigner
from pyflink.datastream.connectors import Sink
from pyflink.datastream.functions import (_get_python_env, FlatMapFunction, MapFunction, Function,
                                          FunctionWrapper, SinkFunction, FilterFunction,
                                          KeySelector, ReduceFunction, CoMapFunction,
                                          CoFlatMapFunction, Partitioner, RuntimeContext,
                                          ProcessFunction, KeyedProcessFunction,
                                          KeyedCoProcessFunction, WindowFunction,
                                          ProcessWindowFunction, InternalWindowFunction,
                                          InternalIterableWindowFunction,
                                          InternalIterableProcessWindowFunction, CoProcessFunction,
                                          InternalSingleValueWindowFunction,
                                          InternalSingleValueProcessWindowFunction,
                                          PassThroughWindowFunction, AggregateFunction,
                                          NullByteKeySelector, AllWindowFunction,
                                          InternalIterableAllWindowFunction,
                                          ProcessAllWindowFunction,
                                          InternalIterableProcessAllWindowFunction,
                                          BroadcastProcessFunction,
                                          KeyedBroadcastProcessFunction,
                                          InternalSingleValueAllWindowFunction,
                                          PassThroughAllWindowFunction,
                                          InternalSingleValueProcessAllWindowFunction,
                                          AsyncFunctionDescriptor)
from pyflink.datastream.output_tag import OutputTag
from pyflink.datastream.slot_sharing_group import SlotSharingGroup
from pyflink.datastream.state import (ListStateDescriptor, StateDescriptor, ReducingStateDescriptor,
                                      AggregatingStateDescriptor, MapStateDescriptor, ReducingState)
from pyflink.datastream.utils import convert_to_python_obj
from pyflink.datastream.window import (CountTumblingWindowAssigner, CountSlidingWindowAssigner,
                                       CountWindowSerializer, TimeWindowSerializer, Trigger,
                                       WindowAssigner, WindowOperationDescriptor,
                                       GlobalWindowSerializer, MergingWindowAssigner)
from pyflink.java_gateway import get_gateway
from pyflink.util.java_utils import to_jarray

__all__ = ['CloseableIterator', 'DataStream', 'KeyedStream', 'ConnectedStreams', 'WindowedStream',
           'DataStreamSink', 'CloseableIterator', 'BroadcastStream', 'BroadcastConnectedStream']

WINDOW_STATE_NAME = 'window-contents'


class DataStream(object):
    """
    A DataStream represents a stream of elements of the same type. A DataStream can be transformed
    into another DataStream by applying a transformation as for example:

    ::
        >>> DataStream.map(MapFunctionImpl())
        >>> DataStream.filter(FilterFunctionImpl())
    """

    def __init__(self, j_data_stream):
        self._j_data_stream = j_data_stream

    def get_name(self) -> str:
        """
        Gets the name of the current data stream. This name is used by the visualization and logging
        during runtime.

        :return: Name of the stream.
        """
        return self._j_data_stream.getName()

    def name(self, name: str) -> 'DataStream':
        """
        Sets the name of the current data stream. This name is used by the visualization and logging
        during runtime.

        :param name: Name of the stream.
        :return: The named operator.
        """
        self._j_data_stream.name(name)
        return self

    def uid(self, uid: str) -> 'DataStream':
        """
        Sets an ID for this operator. The specified ID is used to assign the same operator ID across
        job submissions (for example when starting a job from a savepoint).

        Important: this ID needs to be unique per transformation and job. Otherwise, job submission
        will fail.

        :param uid: The unique user-specified ID of this transformation.
        :return: The operator with the specified ID.
        """
        self._j_data_stream.uid(uid)
        return self

    def set_uid_hash(self, uid_hash: str) -> 'DataStream':
        """
        Sets an user provided hash for this operator. This will be used AS IS the create the
        JobVertexID. The user provided hash is an alternative to the generated hashed, that is
        considered when identifying an operator through the default hash mechanics fails (e.g.
        because of changes between Flink versions).

        Important: this should be used as a workaround or for trouble shooting. The provided hash
        needs to be unique per transformation and job. Otherwise, job submission will fail.
        Furthermore, you cannot assign user-specified hash to intermediate nodes in an operator
        chain and trying so will let your job fail.

        A use case for this is in migration between Flink versions or changing the jobs in a way
        that changes the automatically generated hashes. In this case, providing the previous hashes
        directly through this method (e.g. obtained from old logs) can help to reestablish a lost
        mapping from states to their target operator.

        :param uid_hash: The user provided hash for this operator. This will become the jobVertexID,
                         which is shown in the logs and web ui.
        :return: The operator with the user provided hash.
        """
        self._j_data_stream.setUidHash(uid_hash)
        return self

    def set_parallelism(self, parallelism: int) -> 'DataStream':
        """
        Sets the parallelism for this operator.

        :param parallelism: THe parallelism for this operator.
        :return: The operator with set parallelism.
        """
        self._j_data_stream.setParallelism(parallelism)
        return self

    def set_max_parallelism(self, max_parallelism: int) -> 'DataStream':
        """
        Sets the maximum parallelism of this operator.

        The maximum parallelism specifies the upper bound for dynamic scaling. It also defines the
        number of key groups used for partitioned state.

        :param max_parallelism: Maximum parallelism.
        :return: The operator with set maximum parallelism.
        """
        self._j_data_stream.setMaxParallelism(max_parallelism)
        return self

    def get_type(self) -> TypeInformation:
        """
        Gets the type of the stream.

        :return: The type of the DataStream.
        """
        return typeinfo._from_java_type(self._j_data_stream.getType())

    def get_execution_environment(self):
        """
        Returns the StreamExecutionEnvironment that was used to create this DataStream.

        :return: The Execution Environment.
        """
        from pyflink.datastream import StreamExecutionEnvironment
        return StreamExecutionEnvironment(
            j_stream_execution_environment=self._j_data_stream.getExecutionEnvironment())

    def get_execution_config(self) -> ExecutionConfig:
        return ExecutionConfig(j_execution_config=self._j_data_stream.getExecutionConfig())

    def force_non_parallel(self) -> 'DataStream':
        """
        Sets the parallelism and maximum parallelism of this operator to one. And mark this operator
        cannot set a non-1 degree of parallelism.

        :return: The operator with only one parallelism.
        """
        self._j_data_stream.forceNonParallel()
        return self

    def set_buffer_timeout(self, timeout_millis: int) -> 'DataStream':
        """
        Sets the buffering timeout for data produced by this operation. The timeout defines how long
        data may linger ina partially full buffer before being sent over the network.

        Lower timeouts lead to lower tail latencies, but may affect throughput. Timeouts of 1 ms
        still sustain high throughput, even for jobs with high parallelism.

        A value of '-1' means that the default buffer timeout should be used. A value of '0'
        indicates that no buffering should happen, and all records/events should be immediately sent
        through the network, without additional buffering.

        :param timeout_millis: The maximum time between two output flushes.
        :return: The operator with buffer timeout set.
        """
        self._j_data_stream.setBufferTimeout(timeout_millis)
        return self

    def start_new_chain(self) -> 'DataStream':
        """
        Starts a new task chain beginning at this operator. This operator will be chained (thread
        co-located for increased performance) to any previous tasks even if possible.

        :return: The operator with chaining set.
        """
        self._j_data_stream.startNewChain()
        return self

    def disable_chaining(self) -> 'DataStream':
        """
        Turns off chaining for this operator so thread co-location will not be used as an
        optimization.
        Chaining can be turned off for the whole job by
        StreamExecutionEnvironment.disableOperatorChaining() however it is not advised for
        performance consideration.

        :return: The operator with chaining disabled.
        """
        self._j_data_stream.disableChaining()
        return self

    def slot_sharing_group(self, slot_sharing_group: Union[str, SlotSharingGroup]) -> 'DataStream':
        """
        Sets the slot sharing group of this operation. Parallel instances of operations that are in
        the same slot sharing group will be co-located in the same TaskManager slot, if possible.

        Operations inherit the slot sharing group of input operations if all input operations are in
        the same slot sharing group and no slot sharing group was explicitly specified.

        Initially an operation is in the default slot sharing group. An operation can be put into
        the default group explicitly by setting the slot sharing group to 'default'.

        :param slot_sharing_group: The slot sharing group name or which contains name and its
                        resource spec.
        :return: This operator.
        """
        if isinstance(slot_sharing_group, SlotSharingGroup):
            self._j_data_stream.slotSharingGroup(slot_sharing_group.get_java_slot_sharing_group())
        else:
            self._j_data_stream.slotSharingGroup(slot_sharing_group)
        return self

    def set_description(self, description: str) -> 'DataStream':
        """
        Sets the description for this operator.

        Description is used in json plan and web ui, but not in logging and metrics where only
        name is available. Description is expected to provide detailed information about the
        operator, while name is expected to be more simple, providing summary information only,
        so that we can have more user-friendly logging messages and metric tags without losing
        useful messages for debugging.

        :param description: The description for this operator.
        :return: The operator with new description.

        .. versionadded:: 1.15.0
        """
        self._j_data_stream.setDescription(description)
        return self

    def map(self, func: Union[Callable, MapFunction], output_type: TypeInformation = None) \
            -> 'DataStream':
        """
        Applies a Map transformation on a DataStream. The transformation calls a MapFunction for
        each element of the DataStream. Each MapFunction call returns exactly one element.

        Note that If user does not specify the output data type, the output data will be serialized
        as pickle primitive byte array.

        :param func: The MapFunction that is called for each element of the DataStream.
        :param output_type: The type information of the MapFunction output data.
        :return: The transformed DataStream.
        """
        if not isinstance(func, MapFunction) and not callable(func):
            raise TypeError("The input must be a MapFunction or a callable function")

        class MapProcessFunctionAdapter(ProcessFunction):

            def __init__(self, map_func):
                if isinstance(map_func, MapFunction):
                    self._open_func = map_func.open
                    self._close_func = map_func.close
                    self._map_func = map_func.map
                else:
                    self._open_func = None
                    self._close_func = None
                    self._map_func = map_func

            def open(self, runtime_context: RuntimeContext):
                if self._open_func:
                    self._open_func(runtime_context)

            def close(self):
                if self._close_func:
                    self._close_func()

            def process_element(self, value, ctx: 'ProcessFunction.Context'):
                yield self._map_func(value)

        return self.process(MapProcessFunctionAdapter(func), output_type) \
            .name("Map")

    def flat_map(self,
                 func: Union[Callable, FlatMapFunction],
                 output_type: TypeInformation = None) -> 'DataStream':
        """
        Applies a FlatMap transformation on a DataStream. The transformation calls a FlatMapFunction
        for each element of the DataStream. Each FlatMapFunction call can return any number of
        elements including none.

        :param func: The FlatMapFunction that is called for each element of the DataStream.
        :param output_type: The type information of output data.
        :return: The transformed DataStream.
        """

        if not isinstance(func, FlatMapFunction) and not callable(func):
            raise TypeError("The input must be a FlatMapFunction or a callable function")

        class FlatMapProcessFunctionAdapter(ProcessFunction):

            def __init__(self, flat_map_func):
                if isinstance(flat_map_func, FlatMapFunction):
                    self._open_func = flat_map_func.open
                    self._close_func = flat_map_func.close
                    self._flat_map_func = flat_map_func.flat_map
                else:
                    self._open_func = None
                    self._close_func = None
                    self._flat_map_func = flat_map_func

            def open(self, runtime_context: RuntimeContext):
                if self._open_func:
                    self._open_func(runtime_context)

            def close(self):
                if self._close_func:
                    self._close_func()

            def process_element(self, value, ctx: 'ProcessFunction.Context'):
                yield from self._flat_map_func(value)

        return self.process(FlatMapProcessFunctionAdapter(func), output_type) \
            .name("FlatMap")

    def key_by(self,
               key_selector: Union[Callable, KeySelector],
               key_type: TypeInformation = None) -> 'KeyedStream':
        """
        Creates a new KeyedStream that uses the provided key for partitioning its operator states.

        :param key_selector: The KeySelector to be used for extracting the key for partitioning.
        :param key_type: The type information describing the key type.
        :return: The DataStream with partitioned state(i.e. KeyedStream).
        """

        if not isinstance(key_selector, KeySelector) and not callable(key_selector):
            raise TypeError("Parameter key_selector should be type of KeySelector or a callable "
                            "function.")

        class AddKey(ProcessFunction):

            def __init__(self, key_selector):
                if isinstance(key_selector, KeySelector):
                    self._key_selector_open_func = key_selector.open
                    self._key_selector_close_func = key_selector.close
                    self._get_key_func = key_selector.get_key
                else:
                    self._key_selector_open_func = None
                    self._key_selector_close_func = None
                    self._get_key_func = key_selector

            def open(self, runtime_context: RuntimeContext):
                if self._key_selector_open_func:
                    self._key_selector_open_func(runtime_context)

            def close(self):
                if self._key_selector_close_func:
                    self._key_selector_close_func()

            def process_element(self, value, ctx: 'ProcessFunction.Context'):
                yield Row(self._get_key_func(value), value)

        output_type_info = typeinfo._from_java_type(
            self._j_data_stream.getTransformation().getOutputType())
        if key_type is None:
            key_type = Types.PICKLED_BYTE_ARRAY()

        gateway = get_gateway()
        stream_with_key_info = self.process(
            AddKey(key_selector),
            output_type=Types.ROW([key_type, output_type_info]))
        stream_with_key_info.name(gateway.jvm.org.apache.flink.python.util.PythonConfigUtil
                                  .STREAM_KEY_BY_MAP_OPERATOR_NAME)

        JKeyByKeySelector = gateway.jvm.KeyByKeySelector
        key_stream = KeyedStream(
            stream_with_key_info._j_data_stream.keyBy(
                JKeyByKeySelector(),
                Types.ROW([key_type]).get_java_type_info()), output_type_info,
            self)
        return key_stream

    def filter(self, func: Union[Callable, FilterFunction]) -> 'DataStream':
        """
        Applies a Filter transformation on a DataStream. The transformation calls a FilterFunction
        for each element of the DataStream and retains only those element for which the function
        returns true. Elements for which the function returns false are filtered.

        :param func: The FilterFunction that is called for each element of the DataStream.
        :return: The filtered DataStream.
        """

        if not isinstance(func, FilterFunction) and not callable(func):
            raise TypeError("The input must be a FilterFunction or a callable function")

        class FilterProcessFunctionAdapter(ProcessFunction):

            def __init__(self, filter_func):
                if isinstance(filter_func, FilterFunction):
                    self._open_func = filter_func.open
                    self._close_func = filter_func.close
                    self._filter_func = filter_func.filter
                else:
                    self._open_func = None
                    self._close_func = None
                    self._filter_func = filter_func

            def open(self, runtime_context: RuntimeContext):
                if self._open_func:
                    self._open_func(runtime_context)

            def close(self):
                if self._close_func:
                    self._close_func()

            def process_element(self, value, ctx: 'ProcessFunction.Context'):
                if self._filter_func(value):
                    yield value

        output_type = typeinfo._from_java_type(
            self._j_data_stream.getTransformation().getOutputType())
        return self.process(FilterProcessFunctionAdapter(func), output_type=output_type) \
            .name("Filter")

    def window_all(self, window_assigner: WindowAssigner) -> 'AllWindowedStream':
        """
        Windows this data stream to a AllWindowedStream, which evaluates windows over a non key
        grouped stream. Elements are put into windows by a WindowAssigner. The grouping of
        elements is done by window.

        A Trigger can be defined to specify when windows are evaluated. However, WindowAssigners
        have a default Trigger that is used if a Trigger is not specified.

        :param window_assigner: The WindowAssigner that assigns elements to windows.
        :return: The trigger windows data stream.

        .. versionadded:: 1.16.0
        """
        return AllWindowedStream(self, window_assigner)

    def union(self, *streams: 'DataStream') -> 'DataStream':
        """
        Creates a new DataStream by merging DataStream outputs of the same type with each other. The
        DataStreams merged using this operator will be transformed simultaneously.

        :param streams: The DataStream to union outputwith.
        :return: The DataStream.
        """
        j_data_streams = []
        for data_stream in streams:
            if isinstance(data_stream, KeyedStream):
                j_data_streams.append(data_stream._values()._j_data_stream)
            else:
                j_data_streams.append(data_stream._j_data_stream)
        gateway = get_gateway()
        JDataStream = gateway.jvm.org.apache.flink.streaming.api.datastream.DataStream
        j_data_stream_arr = get_gateway().new_array(JDataStream, len(j_data_streams))
        for i in range(len(j_data_streams)):
            j_data_stream_arr[i] = j_data_streams[i]
        j_united_stream = self._j_data_stream.union(j_data_stream_arr)
        return DataStream(j_data_stream=j_united_stream)

    @overload
    def connect(self, ds: 'DataStream') -> 'ConnectedStreams':
        pass

    @overload
    def connect(self, ds: 'BroadcastStream') -> 'BroadcastConnectedStream':
        pass

    def connect(self, ds: Union['DataStream', 'BroadcastStream']) \
            -> Union['ConnectedStreams', 'BroadcastConnectedStream']:
        """
        If ds is a :class:`DataStream`, creates a new :class:`ConnectedStreams` by connecting
        DataStream outputs of (possible) different types with each other. The DataStreams connected
        using this operator can be used with CoFunctions to apply joint transformations.

        If ds is a :class:`BroadcastStream`, creates a new :class:`BroadcastConnectedStream` by
        connecting the current :class:`DataStream` with a :class:`BroadcastStream`. The latter can
        be created using the :meth:`broadcast` method. The resulting stream can be further processed
        using the :meth:`BroadcastConnectedStream.process` method.

        :param ds: The DataStream or BroadcastStream with which this stream will be connected.
        :return: The ConnectedStreams or BroadcastConnectedStream.

        .. versionchanged:: 1.16.0
           Support connect BroadcastStream
        """
        if isinstance(ds, BroadcastStream):
            return BroadcastConnectedStream(
                self, ds, cast(BroadcastStream, ds).broadcast_state_descriptors
            )
        return ConnectedStreams(self, ds)

    def shuffle(self) -> 'DataStream':
        """
        Sets the partitioning of the DataStream so that the output elements are shuffled uniformly
        randomly to the next operation.

        :return: The DataStream with shuffle partitioning set.
        """
        return DataStream(self._j_data_stream.shuffle())

    def project(self, *field_indexes: int) -> 'DataStream':
        """
        Initiates a Project transformation on a Tuple DataStream.

        Note that only Tuple DataStreams can be projected.

        :param field_indexes: The field indexes of the input tuples that are retained. The order of
                              fields in the output tuple corresponds to the order of field indexes.
        :return: The projected DataStream.
        """
        if not isinstance(self.get_type(), typeinfo.TupleTypeInfo):
            raise Exception('Only Tuple DataStreams can be projected.')

        gateway = get_gateway()
        j_index_arr = gateway.new_array(gateway.jvm.int, len(field_indexes))
        for i in range(len(field_indexes)):
            j_index_arr[i] = field_indexes[i]
        return DataStream(self._j_data_stream.project(j_index_arr))

    def rescale(self) -> 'DataStream':
        """
        Sets the partitioning of the DataStream so that the output elements are distributed evenly
        to a subset of instances of the next operation in a round-robin fashion.

        The subset of downstream operations to which the upstream operation sends elements depends
        on the degree of parallelism of both the upstream and downstream operation. For example, if
        the upstream operation has parallelism 2 and the downstream operation has parallelism 4,
        then one upstream operation would distribute elements to two downstream operations. If, on
        the other hand, the downstream operation has parallelism 4 then two upstream operations will
        distribute to one downstream operation while the other two upstream operations will
        distribute to the other downstream operations.

        In cases where the different parallelisms are not multiples of each one or several
        downstream operations will have a differing number of inputs from upstream operations.

        :return: The DataStream with rescale partitioning set.
        """
        return DataStream(self._j_data_stream.rescale())

    def rebalance(self) -> 'DataStream':
        """
        Sets the partitioning of the DataStream so that the output elements are distributed evenly
        to instances of the next operation in a round-robin fashion.

        :return: The DataStream with rebalance partition set.
        """
        return DataStream(self._j_data_stream.rebalance())

    def forward(self) -> 'DataStream':
        """
        Sets the partitioning of the DataStream so that the output elements are forwarded to the
        local sub-task of the next operation.

        :return: The DataStream with forward partitioning set.
        """
        return DataStream(self._j_data_stream.forward())

    @overload
    def broadcast(self) -> 'DataStream':
        pass

    @overload
    def broadcast(self, broadcast_state_descriptor: MapStateDescriptor,
                  *other_broadcast_state_descriptors: MapStateDescriptor) -> 'BroadcastStream':
        pass

    def broadcast(self, broadcast_state_descriptor: Optional[MapStateDescriptor] = None,
                  *other_broadcast_state_descriptors: MapStateDescriptor) \
            -> Union['DataStream', 'BroadcastStream']:
        """
        Sets the partitioning of the DataStream so that the output elements are broadcasted to every
        parallel instance of the next operation.

        If :class:`~state.MapStateDescriptor` s are passed in, it returns a
        :class:`BroadcastStream` with :class:`~state.BroadcastState` s implicitly created as the
        descriptors specified.

        Example:
        ::

            >>> map_state_desc1 = MapStateDescriptor("state1", Types.INT(), Types.INT())
            >>> map_state_desc2 = MapStateDescriptor("state2", Types.INT(), Types.STRING())
            >>> broadcast_stream = ds1.broadcast(map_state_desc1, map_state_desc2)
            >>> broadcast_connected_stream = ds2.connect(broadcast_stream)

        :param broadcast_state_descriptor: the first MapStateDescriptor describing BroadcastState.
        :param other_broadcast_state_descriptors: the rest of MapStateDescriptors describing
            BroadcastStates, if any.
        :return: The DataStream with broadcast partitioning set or a BroadcastStream which can be
            used in :meth:`connect` to create a BroadcastConnectedStream for further processing of
            the elements.

        .. versionchanged:: 1.16.0
           Support return BroadcastStream
        """
        if broadcast_state_descriptor is not None:
            args = [broadcast_state_descriptor]
            args.extend(other_broadcast_state_descriptors)
            for arg in args:
                if not isinstance(arg, MapStateDescriptor):
                    raise TypeError("broadcast_state_descriptor must be MapStateDescriptor")
            broadcast_state_descriptors: List[MapStateDescriptor] = [arg for arg in args]
            return BroadcastStream(cast(DataStream, self.broadcast()), broadcast_state_descriptors)
        return DataStream(self._j_data_stream.broadcast())

    def process(self, func: ProcessFunction, output_type: TypeInformation = None) -> 'DataStream':
        """
        Applies the given ProcessFunction on the input stream, thereby creating a transformed output
        stream.

        The function will be called for every element in the input streams and can produce zero or
        more output elements.

        :param func: The ProcessFunction that is called for each element in the stream.
        :param output_type: TypeInformation for the result type of the function.
        :return: The transformed DataStream.
        """

        from pyflink.fn_execution import flink_fn_execution_pb2
        j_python_data_stream_function_operator, j_output_type_info = \
            _get_one_input_stream_operator(
                self,
                func,
                flink_fn_execution_pb2.UserDefinedDataStreamFunction.PROCESS,  # type: ignore
                output_type)
        return DataStream(self._j_data_stream.transform(
            "PROCESS",
            j_output_type_info,
            j_python_data_stream_function_operator))

    def assign_timestamps_and_watermarks(self, watermark_strategy: WatermarkStrategy) -> \
            'DataStream':
        """
        Assigns timestamps to the elements in the data stream and generates watermarks to signal
        event time progress. The given {@link WatermarkStrategy} is used to create a
        TimestampAssigner and WatermarkGenerator.

        :param watermark_strategy: The strategy to generate watermarks based on event timestamps.
        :return: The stream after the transformation, with assigned timestamps and watermarks.
        """
        if watermark_strategy._timestamp_assigner is not None:
            # in case users have specified custom TimestampAssigner, we need to extract and
            # generate watermark according to the specified TimestampAssigner.

            class TimestampAssignerProcessFunctionAdapter(ProcessFunction):

                def __init__(self, timestamp_assigner: TimestampAssigner):
                    self._extract_timestamp_func = timestamp_assigner.extract_timestamp

                def process_element(self, value, ctx: 'ProcessFunction.Context'):
                    yield value, self._extract_timestamp_func(value, ctx.timestamp())

            # step 1: extract the timestamp according to the specified TimestampAssigner
            timestamped_data_stream = self.process(
                TimestampAssignerProcessFunctionAdapter(watermark_strategy._timestamp_assigner),
                Types.TUPLE([self.get_type(), Types.LONG()]))
            timestamped_data_stream.name("Extract-Timestamp")

            # step 2: assign timestamp and watermark
            gateway = get_gateway()
            JCustomTimestampAssigner = gateway.jvm.org.apache.flink.streaming.api.functions.python \
                .eventtime.CustomTimestampAssigner
            j_watermarked_data_stream = (
                timestamped_data_stream._j_data_stream.assignTimestampsAndWatermarks(
                    watermark_strategy._j_watermark_strategy.withTimestampAssigner(
                        JCustomTimestampAssigner())))

            # step 3: remove the timestamp field which is added in step 1
            JRemoveTimestampMapFunction = gateway.jvm.org.apache.flink.streaming.api.functions \
                .python.eventtime.RemoveTimestampMapFunction
            result = DataStream(j_watermarked_data_stream.map(
                JRemoveTimestampMapFunction(), self._j_data_stream.getType()))
            result.name("Remove-Timestamp")
            return result
        else:
            # if user not specify a TimestampAssigner, then return directly assign the Java
            # watermark strategy.
            return DataStream(self._j_data_stream.assignTimestampsAndWatermarks(
                watermark_strategy._j_watermark_strategy))

    def partition_custom(self, partitioner: Union[Callable, Partitioner],
                         key_selector: Union[Callable, KeySelector]) -> 'DataStream':
        """
        Partitions a DataStream on the key returned by the selector, using a custom partitioner.
        This method takes the key selector to get the key to partition on, and a partitioner that
        accepts the key type.

        Note that this method works only on single field keys, i.e. the selector cannot return
        tuples of fields.

        :param partitioner: The partitioner to assign partitions to keys.
        :param key_selector: The KeySelector with which the DataStream is partitioned.
        :return: The partitioned DataStream.
        """

        if not isinstance(partitioner, Partitioner) and not callable(partitioner):
            raise TypeError("Parameter partitioner should be type of Partitioner or a callable "
                            "function.")

        if not isinstance(key_selector, KeySelector) and not callable(key_selector):
            raise TypeError("Parameter key_selector should be type of KeySelector or a callable "
                            "function.")

        gateway = get_gateway()

        class CustomPartitioner(ProcessFunction):
            """
            A wrapper class for partition_custom map function. It indicates that it is a partition
            custom operation that we need to apply PythonPartitionCustomOperator
            to run the map function.
            """

            def __init__(self, partitioner, key_selector):
                if isinstance(partitioner, Partitioner):
                    self._partitioner_open_func = partitioner.open
                    self._partitioner_close_func = partitioner.close
                    self._partition_func = partitioner.partition
                else:
                    self._partitioner_open_func = None
                    self._partitioner_close_func = None
                    self._partition_func = partitioner

                if isinstance(key_selector, KeySelector):
                    self._key_selector_open_func = key_selector.open
                    self._key_selector_close_func = key_selector.close
                    self._get_key_func = key_selector.get_key
                else:
                    self._key_selector_open_func = None
                    self._key_selector_close_func = None
                    self._get_key_func = key_selector

            def open(self, runtime_context: RuntimeContext):
                if self._partitioner_open_func:
                    self._partitioner_open_func(runtime_context)
                if self._key_selector_open_func:
                    self._key_selector_open_func(runtime_context)

                self.num_partitions = int(runtime_context.get_job_parameter(
                    "NUM_PARTITIONS", "-1"))
                if self.num_partitions <= 0:
                    raise ValueError(
                        "The partition number should be a positive value, got %s"
                        % self.num_partitions)

            def close(self):
                if self._partitioner_close_func:
                    self._partitioner_close_func()
                if self._key_selector_close_func:
                    self._key_selector_close_func()

            def process_element(self, value, ctx: 'ProcessFunction.Context'):
                partition = self._partition_func(self._get_key_func(value), self.num_partitions)
                yield Row(partition, value)

        original_type_info = self.get_type()
        stream_with_partition_info = self.process(
            CustomPartitioner(partitioner, key_selector),
            output_type=Types.ROW([Types.INT(), original_type_info]))

        stream_with_partition_info.name(
            gateway.jvm.org.apache.flink.python.util.PythonConfigUtil
            .STREAM_PARTITION_CUSTOM_MAP_OPERATOR_NAME)

        JPartitionCustomKeySelector = gateway.jvm.PartitionCustomKeySelector
        JIdParitioner = gateway.jvm.org.apache.flink.python.legacy.IdPartitioner
        partitioned_stream_with_partition_info = DataStream(
            stream_with_partition_info._j_data_stream.partitionCustom(
                JIdParitioner(), JPartitionCustomKeySelector()))

        partitioned_stream = partitioned_stream_with_partition_info.map(
            lambda x: x[1], original_type_info)
        partitioned_stream.name(gateway.jvm.org.apache.flink.python.util.PythonConfigUtil
                                .KEYED_STREAM_VALUE_OPERATOR_NAME)
        return DataStream(partitioned_stream._j_data_stream)

    def add_sink(self, sink_func: SinkFunction) -> 'DataStreamSink':
        """
        Adds the given sink to this DataStream. Only streams with sinks added will be executed once
        the StreamExecutionEnvironment.execute() method is called.

        :param sink_func: The SinkFunction object.
        :return: The closed DataStream.
        """
        return DataStreamSink(self._j_data_stream.addSink(sink_func.get_java_function()))

    def sink_to(self, sink: Sink) -> 'DataStreamSink':
        """
        Adds the given sink to this DataStream. Only streams with sinks added will be
        executed once the
        :func:`~pyflink.datastream.stream_execution_environment.StreamExecutionEnvironment.execute`
        method is called.

        :param sink: The user defined sink.
        :return: The closed DataStream.
        """
        ds = self

        from pyflink.datastream.connectors.base import SupportsPreprocessing
        if isinstance(sink, SupportsPreprocessing) and sink.get_transformer() is not None:
            ds = sink.get_transformer().apply(self)

        return DataStreamSink(ds._j_data_stream.sinkTo(sink.get_java_function()))

    def execute_and_collect(self, job_execution_name: str = None, limit: int = None) \
            -> Union['CloseableIterator', list]:
        """
        Triggers the distributed execution of the streaming dataflow and returns an iterator over
        the elements of the given DataStream.

        The DataStream application is executed in the regular distributed manner on the target
        environment, and the events from the stream are polled back to this application process and
        thread through Flink's REST API.

        The returned iterator must be closed to free all cluster resources.

        :param job_execution_name: The name of the job execution.
        :param limit: The limit for the collected elements.
        """
        JPythonConfigUtil = get_gateway().jvm.org.apache.flink.python.util.PythonConfigUtil
        JPythonConfigUtil.configPythonOperator(self._j_data_stream.getExecutionEnvironment())
        self._apply_chaining_optimization()
        if job_execution_name is None and limit is None:
            return CloseableIterator(self._j_data_stream.executeAndCollect(), self.get_type())
        elif job_execution_name is not None and limit is None:
            return CloseableIterator(self._j_data_stream.executeAndCollect(job_execution_name),
                                     self.get_type())
        if job_execution_name is None and limit is not None:
            return list(map(lambda data: convert_to_python_obj(data, self.get_type()),
                            self._j_data_stream.executeAndCollect(limit)))
        else:
            return list(map(lambda data: convert_to_python_obj(data, self.get_type()),
                            self._j_data_stream.executeAndCollect(job_execution_name, limit)))

    def print(self, sink_identifier: str = None) -> 'DataStreamSink':
        """
        Writes a DataStream to the standard output stream (stdout).
        For each element of the DataStream the object string is written.

        NOTE: This will print to stdout on the machine where the code is executed, i.e. the Flink
        worker, and is not fault tolerant.

        :param sink_identifier: The string to prefix the output with.
        :return: The closed DataStream.
        """
        if sink_identifier is not None:
            j_data_stream_sink = self._align_output_type()._j_data_stream.print(sink_identifier)
        else:
            j_data_stream_sink = self._align_output_type()._j_data_stream.print()
        return DataStreamSink(j_data_stream_sink)

    def get_side_output(self, output_tag: OutputTag) -> 'DataStream':
        """
        Gets the :class:`DataStream` that contains the elements that are emitted from an operation
        into the side output with the given :class:`OutputTag`.

        :param output_tag: output tag for the side stream
        :return: The DataStream with specified output tag

        .. versionadded:: 1.16.0
        """
        ds = DataStream(self._j_data_stream.getSideOutput(output_tag.get_java_output_tag()))
        return ds.map(lambda i: i, output_type=output_tag.type_info)

    def cache(self) -> 'CachedDataStream':
        """
        Cache the intermediate result of the transformation. Only support bounded streams and
        currently only block mode is supported. The cache is generated lazily at the first time the
        intermediate result is computed. The cache will be clear when the StreamExecutionEnvironment
        close.

        :return: The cached DataStream that can use in later job to reuse the cached intermediate
                 result.

        .. versionadded:: 1.16.0
        """
        return CachedDataStream(self._j_data_stream.cache())

    def _apply_chaining_optimization(self):
        """
        Chain the Python operators if possible.
        """
        gateway = get_gateway()
        JPythonOperatorChainingOptimizer = gateway.jvm.org.apache.flink.python.chain. \
            PythonOperatorChainingOptimizer
        j_transformation = JPythonOperatorChainingOptimizer.apply(
            self._j_data_stream.getExecutionEnvironment(),
            self._j_data_stream.getTransformation())
        self._j_data_stream = gateway.jvm.org.apache.flink.streaming.api.datastream.DataStream(
            self._j_data_stream.getExecutionEnvironment(), j_transformation)

    def _align_output_type(self) -> 'DataStream':
        """
        Transform the pickled python object into String if the output type is PickledByteArrayInfo.
        """
        from py4j.java_gateway import get_java_class

        gateway = get_gateway()
        ExternalTypeInfo_CLASS = get_java_class(
            gateway.jvm.org.apache.flink.table.runtime.typeutils.ExternalTypeInfo)
        RowTypeInfo_CLASS = get_java_class(
            gateway.jvm.org.apache.flink.api.java.typeutils.RowTypeInfo)
        output_type_info_class = self._j_data_stream.getTransformation().getOutputType().getClass()
        if output_type_info_class.isAssignableFrom(
                Types.PICKLED_BYTE_ARRAY().get_java_type_info()
                .getClass()):
            def python_obj_to_str_map_func(value):
                if not isinstance(value, (str, bytes)):
                    value = str(value)
                return value

            transformed_data_stream = DataStream(
                self.map(python_obj_to_str_map_func,
                         output_type=Types.STRING())._j_data_stream)
            return transformed_data_stream
        elif (output_type_info_class.isAssignableFrom(ExternalTypeInfo_CLASS) or
              output_type_info_class.isAssignableFrom(RowTypeInfo_CLASS)):
            def python_obj_to_str_map_func(value):
                assert isinstance(value, Row)
                return '{}[{}]'.format(value.get_row_kind(),
                                       ','.join([str(item) for item in value._values]))

            transformed_data_stream = DataStream(
                self.map(python_obj_to_str_map_func,
                         output_type=Types.STRING())._j_data_stream)
            return transformed_data_stream
        else:
            return self


class DataStreamSink(object):
    """
    A Stream Sink. This is used for emitting elements from a streaming topology.
    """

    def __init__(self, j_data_stream_sink):
        """
        The constructor of DataStreamSink.

        :param j_data_stream_sink: A DataStreamSink java object.
        """
        self._j_data_stream_sink = j_data_stream_sink

    def name(self, name: str) -> 'DataStreamSink':
        """
        Sets the name of this sink. THis name is used by the visualization and logging during
        runtime.

        :param name: The name of this sink.
        :return: The named sink.
        """
        self._j_data_stream_sink.name(name)
        return self

    def uid(self, uid: str) -> 'DataStreamSink':
        """
        Sets an ID for this operator. The specified ID is used to assign the same operator ID across
        job submissions (for example when starting a job from a savepoint).

        Important: this ID needs to be unique per transformation and job. Otherwise, job submission
        will fail.

        :param uid: The unique user-specified ID of this transformation.
        :return: The operator with the specified ID.
        """
        self._j_data_stream_sink.uid(uid)
        return self

    def set_uid_hash(self, uid_hash: str) -> 'DataStreamSink':
        """
        Sets an user provided hash for this operator. This will be used AS IS the create the
        JobVertexID. The user provided hash is an alternative to the generated hashed, that is
        considered when identifying an operator through the default hash mechanics fails (e.g.
        because of changes between Flink versions).

        Important: this should be used as a workaround or for trouble shooting. The provided hash
        needs to be unique per transformation and job. Otherwise, job submission will fail.
        Furthermore, you cannot assign user-specified hash to intermediate nodes in an operator
        chain and trying so will let your job fail.

        A use case for this is in migration between Flink versions or changing the jobs in a way
        that changes the automatically generated hashes. In this case, providing the previous hashes
        directly through this method (e.g. obtained from old logs) can help to reestablish a lost
        mapping from states to their target operator.

        :param uid_hash: The user provided hash for this operator. This will become the jobVertexID,
                         which is shown in the logs and web ui.
        :return: The operator with the user provided hash.
        """
        self._j_data_stream_sink.setUidHash(uid_hash)
        return self

    def set_parallelism(self, parallelism: int) -> 'DataStreamSink':
        """
        Sets the parallelism for this operator.

        :param parallelism: THe parallelism for this operator.
        :return: The operator with set parallelism.
        """
        self._j_data_stream_sink.setParallelism(parallelism)
        return self

    def set_description(self, description: str) -> 'DataStreamSink':
        """
        Sets the description for this sink.

        Description is used in json plan and web ui, but not in logging and metrics where only
        name is available. Description is expected to provide detailed information about the sink,
        while name is expected to be more simple, providing summary information only, so that we can
        have more user-friendly logging messages and metric tags without losing useful messages for
        debugging.

        :param description: The description for this sink.
        :return: The sink with new description.

        .. versionadded:: 1.15.0
        """
        self._j_data_stream_sink.setDescription(description)
        return self

    def disable_chaining(self) -> 'DataStreamSink':
        """
        Turns off chaining for this operator so thread co-location will not be used as an
        optimization.
        Chaining can be turned off for the whole job by
        StreamExecutionEnvironment.disableOperatorChaining() however it is not advised for
        performance consideration.

        :return: The operator with chaining disabled.
        """
        self._j_data_stream_sink.disableChaining()
        return self

    def slot_sharing_group(self, slot_sharing_group: Union[str, SlotSharingGroup]) \
            -> 'DataStreamSink':
        """
        Sets the slot sharing group of this operation. Parallel instances of operations that are in
        the same slot sharing group will be co-located in the same TaskManager slot, if possible.

        Operations inherit the slot sharing group of input operations if all input operations are in
        the same slot sharing group and no slot sharing group was explicitly specified.

        Initially an operation is in the default slot sharing group. An operation can be put into
        the default group explicitly by setting the slot sharing group to 'default'.

        :param slot_sharing_group: The slot sharing group name or which contains name and its
                        resource spec.
        :return: This operator.
        """
        if isinstance(slot_sharing_group, SlotSharingGroup):
            self._j_data_stream_sink.slotSharingGroup(
                slot_sharing_group.get_java_slot_sharing_group())
        else:
            self._j_data_stream_sink.slotSharingGroup(slot_sharing_group)
        return self


class KeyedStream(DataStream):
    """
    A KeyedStream represents a DataStream on which operator state is partitioned by key using a
    provided KeySelector. Typical operations supported by a DataStream are also possible on a
    KeyedStream, with the exception of partitioning methods such as shuffle, forward and keyBy.

    Reduce-style operations, such as reduce and sum work on elements that have the same key.
    """

    def __init__(self, j_keyed_stream, original_data_type_info, origin_stream: DataStream):
        """
        Constructor of KeyedStream.

        :param j_keyed_stream: A java KeyedStream object.
        :param original_data_type_info: Original data typeinfo.
        :param origin_stream: The DataStream before key by.
        """
        super(KeyedStream, self).__init__(j_data_stream=j_keyed_stream)
        self._original_data_type_info = original_data_type_info
        self._origin_stream = origin_stream

    def map(self, func: Union[Callable, MapFunction], output_type: TypeInformation = None) \
            -> 'DataStream':
        """
        Applies a Map transformation on a KeyedStream. The transformation calls a MapFunction for
        each element of the DataStream. Each MapFunction call returns exactly one element.

        Note that If user does not specify the output data type, the output data will be serialized
        as pickle primitive byte array.

        :param func: The MapFunction that is called for each element of the DataStream.
        :param output_type: The type information of the MapFunction output data.
        :return: The transformed DataStream.
        """
        if not isinstance(func, MapFunction) and not callable(func):
            raise TypeError("The input must be a MapFunction or a callable function")

        class MapKeyedProcessFunctionAdapter(KeyedProcessFunction):

            def __init__(self, map_func):
                if isinstance(map_func, MapFunction):
                    self._open_func = map_func.open
                    self._close_func = map_func.close
                    self._map_func = map_func.map
                else:
                    self._open_func = None
                    self._close_func = None
                    self._map_func = map_func

            def open(self, runtime_context: RuntimeContext):
                if self._open_func:
                    self._open_func(runtime_context)

            def close(self):
                if self._close_func:
                    self._close_func()

            def process_element(self, value, ctx: 'KeyedProcessFunction.Context'):
                yield self._map_func(value)

        return self.process(MapKeyedProcessFunctionAdapter(func), output_type) \
            .name("Map")  # type: ignore

    def flat_map(self,
                 func: Union[Callable, FlatMapFunction],
                 output_type: TypeInformation = None) -> 'DataStream':
        """
        Applies a FlatMap transformation on a KeyedStream. The transformation calls a
        FlatMapFunction for each element of the DataStream. Each FlatMapFunction call can return
        any number of elements including none.

        :param func: The FlatMapFunction that is called for each element of the DataStream.
        :param output_type: The type information of output data.
        :return: The transformed DataStream.
        """

        if not isinstance(func, FlatMapFunction) and not callable(func):
            raise TypeError("The input must be a FlatMapFunction or a callable function")

        class FlatMapKeyedProcessFunctionAdapter(KeyedProcessFunction):

            def __init__(self, flat_map_func):
                if isinstance(flat_map_func, FlatMapFunction):
                    self._open_func = flat_map_func.open
                    self._close_func = flat_map_func.close
                    self._flat_map_func = flat_map_func.flat_map
                else:
                    self._open_func = None
                    self._close_func = None
                    self._flat_map_func = flat_map_func

            def open(self, runtime_context: RuntimeContext):
                if self._open_func:
                    self._open_func(runtime_context)

            def close(self):
                if self._close_func:
                    self._close_func()

            def process_element(self, value, ctx: 'KeyedProcessFunction.Context'):
                yield from self._flat_map_func(value)

        return self.process(FlatMapKeyedProcessFunctionAdapter(func), output_type) \
            .name("FlatMap")

    def reduce(self, func: Union[Callable, ReduceFunction]) -> 'DataStream':
        """
        Applies a reduce transformation on the grouped data stream grouped on by the given
        key position. The `ReduceFunction` will receive input values based on the key value.
        Only input values with the same key will go to the same reducer.

        Example:
        ::

            >>> ds = env.from_collection([(1, 'a'), (2, 'a'), (3, 'a'), (4, 'b'])
            >>> ds.key_by(lambda x: x[1]).reduce(lambda a, b: a[0] + b[0], b[1])

        :param func: The ReduceFunction that is called for each element of the DataStream.
        :return: The transformed DataStream.
        """

        if not isinstance(func, ReduceFunction) and not callable(func):
            raise TypeError("The input must be a ReduceFunction or a callable function")

        output_type = _from_java_type(self._original_data_type_info.get_java_type_info())

        gateway = get_gateway()
        j_conf = get_j_env_configuration(self._j_data_stream.getExecutionEnvironment())
        python_execution_mode = (
            j_conf.get(
                gateway.jvm.org.apache.flink.python.PythonOptions.PYTHON_EXECUTION_MODE))

        class ReduceProcessKeyedProcessFunctionAdapter(KeyedProcessFunction):

            def __init__(self, reduce_function):
                if isinstance(reduce_function, ReduceFunction):
                    self._open_func = reduce_function.open
                    self._close_func = reduce_function.close
                    self._reduce_function = reduce_function.reduce
                else:
                    self._open_func = None
                    self._close_func = None
                    self._reduce_function = reduce_function
                self._reduce_state: ReducingState = None
                self._in_batch_execution_mode = True

            def open(self, runtime_context: RuntimeContext):
                if self._open_func:
                    self._open_func(runtime_context)

                self._reduce_state = runtime_context.get_reducing_state(
                    ReducingStateDescriptor(
                        "_reduce_state" + str(uuid.uuid4()),
                        self._reduce_function,
                        output_type))

                if python_execution_mode == "process":
                    from pyflink.fn_execution.datastream.process.runtime_context import (
                        StreamingRuntimeContext)
                    self._in_batch_execution_mode = (
                        cast(StreamingRuntimeContext, runtime_context)._in_batch_execution_mode)
                else:
                    self._in_batch_execution_mode = runtime_context.get_job_parameter(
                        "inBatchExecutionMode", "false") == "true"

            def close(self):
                if self._close_func:
                    self._close_func()

            def process_element(self, value, ctx: 'KeyedProcessFunction.Context'):
                if self._in_batch_execution_mode:
                    reduce_value = self._reduce_state.get()
                    if reduce_value is None:
                        # register a timer for emitting the result at the end when this is the
                        # first input for this key
                        ctx.timer_service().register_event_time_timer(0x7fffffffffffffff)
                    self._reduce_state.add(value)
                else:
                    self._reduce_state.add(value)
                    # only emitting the result when all the data for a key is received
                    yield self._reduce_state.get()

            def on_timer(self, timestamp: int, ctx: 'KeyedProcessFunction.OnTimerContext'):
                current_value = self._reduce_state.get()
                if current_value is not None:
                    yield current_value

        return self.process(ReduceProcessKeyedProcessFunctionAdapter(func), output_type) \
            .name("Reduce")

    def filter(self, func: Union[Callable, FilterFunction]) -> 'DataStream':
        if not isinstance(func, FilterFunction) and not callable(func):
            raise TypeError("The input must be a FilterFunction or a callable function")

        class FilterKeyedProcessFunctionAdapter(KeyedProcessFunction):

            def __init__(self, filter_func):
                if isinstance(filter_func, FilterFunction):
                    self._open_func = filter_func.open
                    self._close_func = filter_func.close
                    self._filter_func = filter_func.filter
                else:
                    self._open_func = None
                    self._close_func = None
                    self._filter_func = filter_func

            def open(self, runtime_context: RuntimeContext):
                if self._open_func:
                    self._open_func(runtime_context)

            def close(self):
                if self._close_func:
                    self._close_func()

            def process_element(self, value, ctx: 'KeyedProcessFunction.Context'):
                if self._filter_func(value):
                    yield value

        return self.process(FilterKeyedProcessFunctionAdapter(func), self._original_data_type_info)\
            .name("Filter")

    class AccumulateType(Enum):
        MIN = 1
        MAX = 2
        MIN_BY = 3
        MAX_BY = 4
        SUM = 5

    def _accumulate(self, position: Union[int, str], acc_type: AccumulateType):
        """
        The base method is used for operators such as min, max, min_by, max_by, sum.
        """
        if not isinstance(position, int) and not isinstance(position, str):
            raise TypeError("The field position must be of int or str type to locate the value to "
                            "calculate for min, max, min_by, max_by and sum."
                            "The given type is: %s" % type(position))

        class AccumulateReduceFunction(ReduceFunction):

            def __init__(self, position, agg_type):
                self._pos = position
                self._agg_type = agg_type
                self._reduce_func = None

            def reduce(self, value1, value2):

                def init_reduce_func(value_to_check):
                    if acc_type == KeyedStream.AccumulateType.MIN_BY:
                        # Logic for min_by operator.
                        def reduce_func(v1, v2):
                            if isinstance(value_to_check, (tuple, list, Row)):
                                return v2 if v2[self._pos] < v1[self._pos] else v1
                            else:
                                return v2 if v2 < v1 else v1
                        self._reduce_func = reduce_func

                    elif acc_type == KeyedStream.AccumulateType.MAX_BY:
                        # Logic for max_by operator.
                        def reduce_func(v1, v2):
                            if isinstance(value_to_check, (tuple, list, Row)):
                                return v2 if v2[self._pos] > v1[self._pos] else v1
                            else:
                                return v2 if v2 > v1 else v1
                        self._reduce_func = reduce_func

                    # for MIN / MAX / SUM
                    elif isinstance(value_to_check, tuple):
                        def reduce_func(v1, v2):
                            v1_list = list(v1)
                            if acc_type == KeyedStream.AccumulateType.MIN:
                                # Logic for min operator with tuple type input.
                                v1_list[self._pos] = v2[self._pos] \
                                    if v2[self._pos] < v1[self._pos] else v1[self._pos]
                            elif acc_type == KeyedStream.AccumulateType.MAX:
                                # Logic for max operator with tuple type input.
                                v1_list[self._pos] = v2[self._pos] \
                                    if v2[self._pos] > v1[self._pos] else v1[self._pos]
                            else:
                                # Logic for sum operator with tuple type input.
                                v1_list[self._pos] = v1[self._pos] + v2[self._pos]
                                return tuple(v1_list)
                            return tuple(v1_list)

                        self._reduce_func = reduce_func

                    elif isinstance(value_to_check, (list, Row)):
                        def reduce_func(v1, v2):
                            if acc_type == KeyedStream.AccumulateType.MIN:
                                # Logic for min operator with List and Row types input.
                                v1[self._pos] = v2[self._pos] \
                                    if v2[self._pos] < v1[self._pos] else v1[self._pos]
                            elif acc_type == KeyedStream.AccumulateType.MAX:
                                # Logic for max operator with List and Row types input.
                                v1[self._pos] = v2[self._pos] \
                                    if v2[self._pos] > v1[self._pos] else v1[self._pos]
                            else:
                                # Logic for sum operator with List and Row types input.
                                v1[self._pos] = v1[self._pos] + v2[self._pos]
                            return v1

                        self._reduce_func = reduce_func

                    else:
                        if self._pos != 0:
                            raise TypeError(
                                "The %s field selected on a basic type. A field expression "
                                "on a basic type can only select the 0th field (which means "
                                "selecting the entire basic type)." % self._pos)

                        def reduce_func(v1, v2):
                            if acc_type == KeyedStream.AccumulateType.MIN:
                                # Logic for min operator with basic type input.
                                return v2 if v2 < v1 else v1
                            elif acc_type == KeyedStream.AccumulateType.MAX:
                                # Logic for max operator with basic type input.
                                return v2 if v2 > v1 else v1
                            else:
                                # Logic for sum operator with basic type input.
                                return v1 + v2

                        self._reduce_func = reduce_func

                if not self._reduce_func:
                    init_reduce_func(value2)
                return self._reduce_func(value1, value2)

        return self.reduce(AccumulateReduceFunction(position, acc_type))

    def sum(self, position_to_sum: Union[int, str] = 0) -> 'DataStream':
        """
        Applies an aggregation that gives a rolling sum of the data stream at the given position
        grouped by the given key. An independent aggregate is kept per key.

        Example(Tuple data to sum):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('b', 1), ('b', 5)])
            >>> ds.key_by(lambda x: x[0]).sum(1)

        Example(Row data to sum):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...                          type_info=Types.ROW([Types.STRING(), Types.INT()]))
            >>> ds.key_by(lambda x: x[0]).sum(1)

        Example(Row data with fields name to sum):
        ::

            >>> ds = env.from_collection(
            ...     [('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...     type_info=Types.ROW_NAMED(["key", "value"], [Types.STRING(), Types.INT()])
            ... )
            >>> ds.key_by(lambda x: x[0]).sum("value")

        :param position_to_sum: The field position in the data points to sum, type can be int which
                                indicates the index of the column to operate on or str which
                                indicates the name of the column to operate on.
        :return: The transformed DataStream.

        .. versionadded:: 1.16.0
        """
        return self._accumulate(position_to_sum, KeyedStream.AccumulateType.SUM)

    def min(self, position_to_min: Union[int, str] = 0) -> 'DataStream':
        """
        Applies an aggregation that gives the current minimum of the data stream at the given
        position by the given key. An independent aggregate is kept per key.

        Example(Tuple data):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('b', 1), ('b', 5)])
            >>> ds.key_by(lambda x: x[0]).min(1)

        Example(Row data):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...                          type_info=Types.ROW([Types.STRING(), Types.INT()]))
            >>> ds.key_by(lambda x: x[0]).min(1)

        Example(Row data with fields name):
        ::

            >>> ds = env.from_collection(
            ...     [('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...     type_info=Types.ROW_NAMED(["key", "value"], [Types.STRING(), Types.INT()])
            ... )
            >>> ds.key_by(lambda x: x[0]).min("value")

        :param position_to_min: The field position in the data points to minimize. The type can be
                                int (field position) or str (field name). This is applicable to
                                Tuple types, List types, Row types, and basic types (which is
                                considered as having one field).
        :return: The transformed DataStream.

        .. versionadded:: 1.16.0
        """
        return self._accumulate(position_to_min, KeyedStream.AccumulateType.MIN)

    def max(self, position_to_max: Union[int, str] = 0) -> 'DataStream':
        """
        Applies an aggregation that gives the current maximize of the data stream at the given
        position by the given key. An independent aggregate is kept per key.

        Example(Tuple data):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('b', 1), ('b', 5)])
            >>> ds.key_by(lambda x: x[0]).max(1)

        Example(Row data):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...                          type_info=Types.ROW([Types.STRING(), Types.INT()]))
            >>> ds.key_by(lambda x: x[0]).max(1)

        Example(Row data with fields name):
        ::

            >>> ds = env.from_collection(
            ...     [('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...     type_info=Types.ROW_NAMED(["key", "value"], [Types.STRING(), Types.INT()])
            ... )
            >>> ds.key_by(lambda x: x[0]).max("value")

        :param position_to_max: The field position in the data points to maximize. The type can be
                                int (field position) or str (field name). This is applicable to
                                Tuple types, List types, Row types, and basic types (which is
                                considered as having one field).
        :return: The transformed DataStream.

        .. versionadded:: 1.16.0
        """
        return self._accumulate(position_to_max, KeyedStream.AccumulateType.MAX)

    def min_by(self, position_to_min_by: Union[int, str] = 0) -> 'DataStream':
        """
        Applies an aggregation that gives the current element with the minimum value at the
        given position by the given key. An independent aggregate is kept per key.
        If more elements have the minimum value at the given position,
        the operator returns the first one by default.

        Example(Tuple data):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('b', 1), ('b', 5)])
            >>> ds.key_by(lambda x: x[0]).min_by(1)

        Example(Row data):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...                          type_info=Types.ROW([Types.STRING(), Types.INT()]))
            >>> ds.key_by(lambda x: x[0]).min_by(1)

        Example(Row data with fields name):
        ::

            >>> ds = env.from_collection(
            ...     [('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...     type_info=Types.ROW_NAMED(["key", "value"], [Types.STRING(), Types.INT()])
            ... )
            >>> ds.key_by(lambda x: x[0]).min_by("value")

        :param position_to_min_by: The field position in the data points to minimize. The type can
                                   be int (field position) or str (field name). This is applicable
                                   to Tuple types, List types, Row types, and basic types (which is
                                   considered as having one field).
        :return: The transformed DataStream.

        .. versionadded:: 1.16.0
        """
        return self._accumulate(position_to_min_by, KeyedStream.AccumulateType.MIN_BY)

    def max_by(self, position_to_max_by: Union[int, str] = 0) -> 'DataStream':
        """
        Applies an aggregation that gives the current element with the maximize value at the
        given position by the given key. An independent aggregate is kept per key.
        If more elements have the maximize value at the given position,
        the operator returns the first one by default.


        Example(Tuple data):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('b', 1), ('b', 5)])
            >>> ds.key_by(lambda x: x[0]).max_by(1)

        Example(Row data):
        ::

            >>> ds = env.from_collection([('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...                          type_info=Types.ROW([Types.STRING(), Types.INT()]))
            >>> ds.key_by(lambda x: x[0]).max_by(1)

        Example(Row data with fields name):
        ::

            >>> ds = env.from_collection(
            ...     [('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2)],
            ...     type_info=Types.ROW_NAMED(["key", "value"], [Types.STRING(), Types.INT()])
            ... )
            >>> ds.key_by(lambda x: x[0]).max_by("value")

        :param position_to_max_by: The field position in the data points to maximize. The type can
                                   be int (field position) or str (field name). This is applicable
                                   to Tuple types, List types, Row types, and basic types (which is
                                   considered as having one field).
        :return: The transformed DataStream.

        .. versionadded:: 1.16.0
        """
        return self._accumulate(position_to_max_by, KeyedStream.AccumulateType.MAX_BY)

    def add_sink(self, sink_func: SinkFunction) -> 'DataStreamSink':
        return self._values().add_sink(sink_func)

    def key_by(self, key_selector: Union[Callable, KeySelector],
               key_type: TypeInformation = None) -> 'KeyedStream':
        return self._origin_stream.key_by(key_selector, key_type)

    def process(self, func: KeyedProcessFunction,  # type: ignore
                output_type: TypeInformation = None) -> 'DataStream':
        """
        Applies the given ProcessFunction on the input stream, thereby creating a transformed output
        stream.

        The function will be called for every element in the input streams and can produce zero or
        more output elements.

        :param func: The KeyedProcessFunction that is called for each element in the stream.
        :param output_type: TypeInformation for the result type of the function.
        :return: The transformed DataStream.
        """
        if not isinstance(func, KeyedProcessFunction):
            raise TypeError("KeyedProcessFunction is required for KeyedStream.")

        from pyflink.fn_execution import flink_fn_execution_pb2
        j_python_data_stream_function_operator, j_output_type_info = \
            _get_one_input_stream_operator(
                self,
                func,
                flink_fn_execution_pb2.UserDefinedDataStreamFunction.KEYED_PROCESS,  # type: ignore
                output_type)
        return DataStream(self._j_data_stream.transform(
            "KEYED PROCESS",
            j_output_type_info,
            j_python_data_stream_function_operator))

    def window(self, window_assigner: WindowAssigner) -> 'WindowedStream':
        """
        Windows this data stream to a WindowedStream, which evaluates windows over a key
        grouped stream. Elements are put into windows by a WindowAssigner. The grouping of
        elements is done both by key and by window.

        A Trigger can be defined to specify when windows are evaluated. However, WindowAssigners
        have a default Trigger that is used if a Trigger is not specified.

        :param window_assigner: The WindowAssigner that assigns elements to windows.
        :return: The trigger windows data stream.
        """
        return WindowedStream(self, window_assigner)

    def count_window(self, size: int, slide: int = 0):
        """
        Windows this KeyedStream into tumbling or sliding count windows.

        :param size: The size of the windows in number of elements.
        :param slide: The slide interval in number of elements.

        .. versionadded:: 1.16.0
        """
        if slide == 0:
            return WindowedStream(self, CountTumblingWindowAssigner(size))
        else:
            return WindowedStream(self, CountSlidingWindowAssigner(size, slide))

    def union(self, *streams) -> 'DataStream':
        return self._values().union(*streams)

    @overload
    def connect(self, ds: 'DataStream') -> 'ConnectedStreams':
        pass

    @overload
    def connect(self, ds: 'BroadcastStream') -> 'BroadcastConnectedStream':
        pass

    def connect(self, ds: Union['DataStream', 'BroadcastStream']) \
            -> Union['ConnectedStreams', 'BroadcastConnectedStream']:
        """
        If ds is a :class:`DataStream`, creates a new :class:`ConnectedStreams` by connecting
        DataStream outputs of (possible) different types with each other. The DataStreams connected
        using this operator can be used with CoFunctions to apply joint transformations.

        If ds is a :class:`BroadcastStream`, creates a new :class:`BroadcastConnectedStream` by
        connecting the current :class:`DataStream` with a :class:`BroadcastStream`. The latter can
        be created using the :meth:`broadcast` method. The resulting stream can be further processed
        using the :meth:`BroadcastConnectedStream.process` method.

        :param ds: The DataStream or BroadcastStream with which this stream will be connected.
        :return: The ConnectedStreams or BroadcastConnectedStream.

        .. versionchanged:: 1.16.0
           Support connect BroadcastStream
        """
        return super().connect(ds)

    def shuffle(self) -> 'DataStream':
        raise Exception('Cannot override partitioning for KeyedStream.')

    def project(self, *field_indexes) -> 'DataStream':
        return self._values().project(*field_indexes)

    def rescale(self) -> 'DataStream':
        raise Exception('Cannot override partitioning for KeyedStream.')

    def rebalance(self) -> 'DataStream':
        raise Exception('Cannot override partitioning for KeyedStream.')

    def forward(self) -> 'DataStream':
        raise Exception('Cannot override partitioning for KeyedStream.')

    def broadcast(self, *args):
        """
        Not supported, partitioning for KeyedStream cannot be overridden.
        """
        raise Exception('Cannot override partitioning for KeyedStream.')

    def partition_custom(self, partitioner: Union[Callable, Partitioner],
                         key_selector: Union[Callable, KeySelector]) -> 'DataStream':
        raise Exception('Cannot override partitioning for KeyedStream.')

    def print(self, sink_identifier=None):
        return self._values().print()

    def _values(self) -> 'DataStream':
        """
        Since python KeyedStream is in the format of Row(key_value, original_data), it is used for
        getting the original_data.
        """
        transformed_stream = self.map(lambda x: x, output_type=self._original_data_type_info)
        transformed_stream.name(get_gateway().jvm.org.apache.flink.python.util.PythonConfigUtil
                                .KEYED_STREAM_VALUE_OPERATOR_NAME)
        return DataStream(transformed_stream._j_data_stream)

    def set_parallelism(self, parallelism: int):
        raise Exception("Set parallelism for KeyedStream is not supported.")

    def name(self, name: str):
        raise Exception("Set name for KeyedStream is not supported.")

    def get_name(self) -> str:
        raise Exception("Get name of KeyedStream is not supported.")

    def uid(self, uid: str):
        raise Exception("Set uid for KeyedStream is not supported.")

    def set_uid_hash(self, uid_hash: str):
        raise Exception("Set uid hash for KeyedStream is not supported.")

    def set_max_parallelism(self, max_parallelism: int):
        raise Exception("Set max parallelism for KeyedStream is not supported.")

    def force_non_parallel(self):
        raise Exception("Set force non-parallel for KeyedStream is not supported.")

    def set_buffer_timeout(self, timeout_millis: int):
        raise Exception("Set buffer timeout for KeyedStream is not supported.")

    def start_new_chain(self) -> 'DataStream':
        raise Exception("Start new chain for KeyedStream is not supported.")

    def disable_chaining(self) -> 'DataStream':
        raise Exception("Disable chaining for KeyedStream is not supported.")

    def slot_sharing_group(self, slot_sharing_group: Union[str, SlotSharingGroup]) -> 'DataStream':
        raise Exception("Setting slot sharing group for KeyedStream is not supported.")

    def cache(self) -> 'CachedDataStream':
        raise Exception("Cache for KeyedStream is not supported.")


class CachedDataStream(DataStream):
    """
    CachedDataStream represents a DataStream whose intermediate result will be cached at the first
    time when it is computed. And the cached intermediate result can be used in later job that using
    the same CachedDataStream to avoid re-computing the intermediate result.
    """

    def __init__(self, j_data_stream):
        super(CachedDataStream, self).__init__(j_data_stream)

    def invalidate(self):
        """
        Invalidate the cache intermediate result of this DataStream to release the physical
        resources. Users are not required to invoke this method to release physical resources unless
        they want to. Cache will be recreated if it is used after invalidated.

        .. versionadded:: 1.16.0
        """
        self._j_data_stream.invalidate()

    def set_parallelism(self, parallelism: int):
        raise Exception("Set parallelism for CachedDataStream is not supported.")

    def name(self, name: str):
        raise Exception("Set name for CachedDataStream is not supported.")

    def get_name(self) -> str:
        raise Exception("Get name of CachedDataStream is not supported.")

    def uid(self, uid: str):
        raise Exception("Set uid for CachedDataStream is not supported.")

    def set_uid_hash(self, uid_hash: str):
        raise Exception("Set uid hash for CachedDataStream is not supported.")

    def set_max_parallelism(self, max_parallelism: int):
        raise Exception("Set max parallelism for CachedDataStream is not supported.")

    def force_non_parallel(self):
        raise Exception("Set force non-parallel for CachedDataStream is not supported.")

    def set_buffer_timeout(self, timeout_millis: int):
        raise Exception("Set buffer timeout for CachedDataStream is not supported.")

    def start_new_chain(self) -> 'DataStream':
        raise Exception("Start new chain for CachedDataStream is not supported.")

    def disable_chaining(self) -> 'DataStream':
        raise Exception("Disable chaining for CachedDataStream is not supported.")

    def slot_sharing_group(self, slot_sharing_group: Union[str, SlotSharingGroup]) -> 'DataStream':
        raise Exception("Setting slot sharing group for CachedDataStream is not supported.")


class WindowedStream(object):
    """
    A WindowedStream represents a data stream where elements are grouped by key, and for each
    key, the stream of elements is split into windows based on a WindowAssigner. Window emission
    is triggered based on a Trigger.

    The windows are conceptually evaluated for each key individually, meaning windows can trigger
    at different points for each key.

    Note that the WindowedStream is purely an API construct, during runtime the WindowedStream will
    be collapsed together with the KeyedStream and the operation over the window into one single
    operation.
    """

    def __init__(self, keyed_stream: KeyedStream, window_assigner: WindowAssigner):
        self._keyed_stream = keyed_stream
        self._window_assigner = window_assigner
        self._allowed_lateness = 0
        self._late_data_output_tag: Optional[OutputTag] = None
        self._window_trigger: Trigger = None

    def get_execution_environment(self):
        return self._keyed_stream.get_execution_environment()

    def get_input_type(self):
        return _from_java_type(self._keyed_stream._original_data_type_info.get_java_type_info())

    def trigger(self, trigger: Trigger) -> 'WindowedStream':
        """
        Sets the Trigger that should be used to trigger window emission.
        """
        self._window_trigger = trigger
        return self

    def allowed_lateness(self, time_ms: int) -> 'WindowedStream':
        """
        Sets the time by which elements are allowed to be late. Elements that arrive behind the
        watermark by more than the specified time will be dropped. By default, the allowed lateness
        is 0.

        Setting an allowed lateness is only valid for event-time windows.
        """
        self._allowed_lateness = time_ms
        return self

    def side_output_late_data(self, output_tag: OutputTag) -> 'WindowedStream':
        """
        Send late arriving data to the side output identified by the given :class:`OutputTag`. Data
        is considered late after the watermark has passed the end of the window plus the allowed
        lateness set using :func:`allowed_lateness`.

        You can get the stream of late data using :func:`~DataStream.get_side_output` on the
        :class:`DataStream` resulting from the windowed operation with the same :class:`OutputTag`.

        Example:
        ::

            >>> tag = OutputTag("late-data", Types.TUPLE([Types.INT(), Types.STRING()]))
            >>> main_stream = ds.key_by(lambda x: x[1]) \\
            ...                 .window(TumblingEventTimeWindows.of(Time.seconds(5))) \\
            ...                 .side_output_late_data(tag) \\
            ...                 .reduce(lambda a, b: a[0] + b[0], b[1])
            >>> late_stream = main_stream.get_side_output(tag)

        .. versionadded:: 1.16.0
        """
        self._late_data_output_tag = output_tag
        return self

    def reduce(self,
               reduce_function: Union[Callable, ReduceFunction],
               window_function: Union[WindowFunction, ProcessWindowFunction] = None,
               output_type: TypeInformation = None) -> DataStream:
        """
        Applies a reduce function to the window. The window function is called for each evaluation
        of the window for each key individually. The output of the reduce function is interpreted as
        a regular non-windowed stream.

        This window will try and incrementally aggregate data as much as the window policies
        permit. For example, tumbling time windows can aggregate the data, meaning that only one
        element per key is stored. Sliding time windows will aggregate on the granularity of the
        slide interval, so a few elements are stored per key (one per slide interval). Custom
        windows may not be able to incrementally aggregate, or may need to store extra values in an
        aggregation tree.

        Example:
        ::

            >>> ds.key_by(lambda x: x[1]) \\
            ...     .window(TumblingEventTimeWindows.of(Time.seconds(5))) \\
            ...     .reduce(lambda a, b: a[0] + b[0], b[1])

        :param reduce_function: The reduce function.
        :param window_function: The window function.
        :param output_type: Type information for the result type of the window function.
        :return: The data stream that is the result of applying the reduce function to the window.

        .. versionadded:: 1.16.0
        """
        if window_function is None:
            internal_window_function: InternalWindowFunction = InternalSingleValueWindowFunction(
                PassThroughWindowFunction())
            if output_type is None:
                output_type = self.get_input_type()
        elif isinstance(window_function, WindowFunction):
            internal_window_function = InternalSingleValueWindowFunction(window_function)
        elif isinstance(window_function, ProcessWindowFunction):
            internal_window_function = InternalSingleValueProcessWindowFunction(window_function)
        else:
            raise TypeError("window_function should be a WindowFunction or ProcessWindowFunction")

        reducing_state_descriptor = ReducingStateDescriptor(WINDOW_STATE_NAME,
                                                            reduce_function,
                                                            self.get_input_type())
        func_desc = type(reduce_function).__name__
        if window_function is not None:
            func_desc = "%s, %s" % (func_desc, type(window_function).__name__)

        return self._get_result_data_stream(internal_window_function,
                                            reducing_state_descriptor,
                                            func_desc,
                                            output_type)

    def aggregate(self,
                  aggregate_function: AggregateFunction,
                  window_function: Union[WindowFunction, ProcessWindowFunction] = None,
                  accumulator_type: TypeInformation = None,
                  output_type: TypeInformation = None) -> DataStream:
        """
        Applies the given window function to each window. The window function is called for each
        evaluation of the window for each key individually. The output of the window function is
        interpreted as a regular non-windowed stream.

        Arriving data is incrementally aggregated using the given aggregate function. This means
        that the window function typically has only a single value to process when called.

        Example:
        ::

            >>> class AverageAggregate(AggregateFunction):
            ...     def create_accumulator(self) -> Tuple[int, int]:
            ...         return 0, 0
            ...
            ...     def add(self, value: Tuple[str, int], accumulator: Tuple[int, int]) \\
            ...             -> Tuple[int, int]:
            ...         return accumulator[0] + value[1], accumulator[1] + 1
            ...
            ...     def get_result(self, accumulator: Tuple[int, int]) -> float:
            ...         return accumulator[0] / accumulator[1]
            ...
            ...     def merge(self, a: Tuple[int, int], b: Tuple[int, int]) -> Tuple[int, int]:
            ...         return a[0] + b[0], a[1] + b[1]
            >>> ds.key_by(lambda x: x[1]) \\
            ...     .window(TumblingEventTimeWindows.of(Time.seconds(5))) \\
            ...     .aggregate(AverageAggregate(),
            ...                accumulator_type=Types.TUPLE([Types.LONG(), Types.LONG()]),
            ...                output_type=Types.DOUBLE())

        :param aggregate_function: The aggregation function that is used for incremental
                                   aggregation.
        :param window_function: The window function.
        :param accumulator_type: Type information for the internal accumulator type of the
                                 aggregation function.
        :param output_type: Type information for the result type of the window function.
        :return: The data stream that is the result of applying the window function to the window.

        .. versionadded:: 1.16.0
        """
        if window_function is None:
            internal_window_function: InternalWindowFunction = InternalSingleValueWindowFunction(
                PassThroughWindowFunction())
        elif isinstance(window_function, WindowFunction):
            internal_window_function = InternalSingleValueWindowFunction(window_function)
        elif isinstance(window_function, ProcessWindowFunction):
            internal_window_function = InternalSingleValueProcessWindowFunction(window_function)
        else:
            raise TypeError("window_function should be a WindowFunction or ProcessWindowFunction")

        if accumulator_type is None:
            accumulator_type = Types.PICKLED_BYTE_ARRAY()
        elif isinstance(accumulator_type, list):
            accumulator_type = RowTypeInfo(accumulator_type)

        aggregating_state_descriptor = AggregatingStateDescriptor(WINDOW_STATE_NAME,
                                                                  aggregate_function,
                                                                  accumulator_type)
        func_desc = type(aggregate_function).__name__
        if window_function is not None:
            func_desc = "%s, %s" % (func_desc, type(window_function).__name__)
        return self._get_result_data_stream(internal_window_function,
                                            aggregating_state_descriptor,
                                            func_desc,
                                            output_type)

    def apply(self,
              window_function: WindowFunction, output_type: TypeInformation = None) -> DataStream:
        """
        Applies the given window function to each window. The window function is called for each
        evaluation of the window for each key individually. The output of the window function is
        interpreted as a regular non-windowed stream.

        Note that this function requires that all data in the windows is buffered until the window
        is evaluated, as the function provides no means of incremental aggregation.

        :param window_function: The window function.
        :param output_type: Type information for the result type of the window function.
        :return: The data stream that is the result of applying the window function to the window.
        """
        internal_window_function: InternalWindowFunction = InternalIterableWindowFunction(
            window_function)
        list_state_descriptor = ListStateDescriptor(WINDOW_STATE_NAME, self.get_input_type())
        func_desc = type(window_function).__name__
        return self._get_result_data_stream(internal_window_function,
                                            list_state_descriptor,
                                            func_desc,
                                            output_type)

    def process(self,
                process_window_function: ProcessWindowFunction,
                output_type: TypeInformation = None) -> DataStream:
        """
        Applies the given window function to each window. The window function is called for each
        evaluation of the window for each key individually. The output of the window function is
        interpreted as a regular non-windowed stream.

        Note that this function requires that all data in the windows is buffered until the window
        is evaluated, as the function provides no means of incremental aggregation.

        :param process_window_function: The window function.
        :param output_type: Type information for the result type of the window function.
        :return: The data stream that is the result of applying the window function to the window.
        """
        internal_window_function: InternalWindowFunction = InternalIterableProcessWindowFunction(
            process_window_function)
        list_state_descriptor = ListStateDescriptor(WINDOW_STATE_NAME, self.get_input_type())
        func_desc = type(process_window_function).__name__
        return self._get_result_data_stream(internal_window_function,
                                            list_state_descriptor,
                                            func_desc,
                                            output_type)

    def _get_result_data_stream(self,
                                internal_window_function: InternalWindowFunction,
                                window_state_descriptor: StateDescriptor,
                                func_desc: str,
                                output_type: TypeInformation):
        if self._window_trigger is None:
            self._window_trigger = self._window_assigner.get_default_trigger(
                self.get_execution_environment())
        window_serializer = self._window_assigner.get_window_serializer()
        window_operation_descriptor = WindowOperationDescriptor(
            self._window_assigner,
            self._window_trigger,
            self._allowed_lateness,
            self._late_data_output_tag,
            window_state_descriptor,
            window_serializer,
            internal_window_function)

        from pyflink.fn_execution import flink_fn_execution_pb2
        j_python_data_stream_function_operator, j_output_type_info = \
            _get_one_input_stream_operator(
                self._keyed_stream,
                window_operation_descriptor,
                flink_fn_execution_pb2.UserDefinedDataStreamFunction.WINDOW,  # type: ignore
                output_type)

        op_name = window_operation_descriptor.generate_op_name()
        op_desc = window_operation_descriptor.generate_op_desc("Window", func_desc)
        return DataStream(self._keyed_stream._j_data_stream.transform(
            op_name,
            j_output_type_info,
            j_python_data_stream_function_operator)).set_description(op_desc)


class AllWindowedStream(object):
    """
    A AllWindowedStream represents a data stream where the stream of elements is split into windows
    based on a WindowAssigner. Window emission is triggered based on a Trigger.

    If an Evictor is specified it will be used to evict elements from the window after evaluation
    was triggered by the Trigger but before the actual evaluation of the window.
    When using an evictor, window performance will degrade significantly, since pre-aggregation of
    window results cannot be used.

    Note that the AllWindowedStream is purely an API construct, during runtime the AllWindowedStream
    will be collapsed together with the operation over the window into one single operation.
    """

    def __init__(self, data_stream: DataStream, window_assigner: WindowAssigner):
        self._keyed_stream = data_stream.key_by(NullByteKeySelector())
        self._window_assigner = window_assigner
        self._allowed_lateness = 0
        self._late_data_output_tag: Optional[OutputTag] = None
        self._window_trigger: Trigger = None

    def get_execution_environment(self):
        return self._keyed_stream.get_execution_environment()

    def get_input_type(self):
        return _from_java_type(self._keyed_stream._original_data_type_info.get_java_type_info())

    def trigger(self, trigger: Trigger) -> 'AllWindowedStream':
        """
        Sets the Trigger that should be used to trigger window emission.
        """
        if isinstance(self._window_assigner, MergingWindowAssigner) \
                and (trigger.can_merge() is not True):
            raise TypeError("A merging window assigner cannot be used with a trigger that does "
                            "not support merging.")

        self._window_trigger = trigger
        return self

    def allowed_lateness(self, time_ms: int) -> 'AllWindowedStream':
        """
        Sets the time by which elements are allowed to be late. Elements that arrive behind the
        watermark by more than the specified time will be dropped. By default, the allowed lateness
        is 0.

        Setting an allowed lateness is only valid for event-time windows.
        """
        self._allowed_lateness = time_ms
        return self

    def side_output_late_data(self, output_tag: OutputTag) -> 'AllWindowedStream':
        """
        Send late arriving data to the side output identified by the given :class:`OutputTag`. Data
        is considered late after the watermark has passed the end of the window plus the allowed
        lateness set using :func:`allowed_lateness`.

        You can get the stream of late data using :func:`~DataStream.get_side_output` on the
        :class:`DataStream` resulting from the windowed operation with the same :class:`OutputTag`.

        Example:
        ::

            >>> tag = OutputTag("late-data", Types.TUPLE([Types.INT(), Types.STRING()]))
            >>> main_stream = ds.window_all(TumblingEventTimeWindows.of(Time.seconds(5))) \\
            ...                 .side_output_late_data(tag) \\
            ...                 .process(MyProcessAllWindowFunction(),
            ...                          Types.TUPLE([Types.LONG(), Types.LONG(), Types.INT()]))
            >>> late_stream = main_stream.get_side_output(tag)
        """
        self._late_data_output_tag = output_tag
        return self

    def reduce(self,
               reduce_function: Union[Callable, ReduceFunction],
               window_function: Union[AllWindowFunction, ProcessAllWindowFunction] = None,
               output_type: TypeInformation = None) -> DataStream:
        """
        Applies the given window function to each window. The window function is called for each
        evaluation of the window for each key individually. The output of the window function is
        interpreted as a regular non-windowed stream.

        Arriving data is incrementally aggregated using the given reducer.

        Example:
        ::

            >>> ds.window_all(TumblingEventTimeWindows.of(Time.seconds(5))) \\
            ...   .reduce(lambda a, b: a[0] + b[0], b[1])

        :param reduce_function: The reduce function.
        :param window_function: The window function.
        :param output_type: Type information for the result type of the window function.
        :return: The data stream that is the result of applying the reduce function to the window.

        .. versionadded:: 1.16.0
        """
        if window_function is None:
            internal_window_function: InternalWindowFunction = InternalSingleValueAllWindowFunction(
                PassThroughAllWindowFunction())
            if output_type is None:
                output_type = self.get_input_type()
        elif isinstance(window_function, AllWindowFunction):
            internal_window_function = InternalSingleValueAllWindowFunction(window_function)
        elif isinstance(window_function, ProcessAllWindowFunction):
            internal_window_function = InternalSingleValueProcessAllWindowFunction(window_function)
        else:
            raise TypeError("window_function should be a AllWindowFunction or "
                            "ProcessAllWindowFunction")

        reducing_state_descriptor = ReducingStateDescriptor(WINDOW_STATE_NAME,
                                                            reduce_function,
                                                            self.get_input_type())
        func_desc = type(reduce_function).__name__
        if window_function is not None:
            func_desc = "%s, %s" % (func_desc, type(window_function).__name__)

        return self._get_result_data_stream(internal_window_function,
                                            reducing_state_descriptor,
                                            func_desc,
                                            output_type)

    def aggregate(self,
                  aggregate_function: AggregateFunction,
                  window_function: Union[AllWindowFunction, ProcessAllWindowFunction] = None,
                  accumulator_type: TypeInformation = None,
                  output_type: TypeInformation = None) -> DataStream:
        """
        Applies the given window function to each window. The window function is called for each
        evaluation of the window for each key individually. The output of the window function is
        interpreted as a regular non-windowed stream.

        Arriving data is incrementally aggregated using the given aggregate function. This means
        that the window function typically has only a single value to process when called.

        Example:
        ::

            >>> class AverageAggregate(AggregateFunction):
            ...     def create_accumulator(self) -> Tuple[int, int]:
            ...         return 0, 0
            ...
            ...     def add(self, value: Tuple[str, int], accumulator: Tuple[int, int]) \\
            ...             -> Tuple[int, int]:
            ...         return accumulator[0] + value[1], accumulator[1] + 1
            ...
            ...     def get_result(self, accumulator: Tuple[int, int]) -> float:
            ...         return accumulator[0] / accumulator[1]
            ...
            ...     def merge(self, a: Tuple[int, int], b: Tuple[int, int]) -> Tuple[int, int]:
            ...         return a[0] + b[0], a[1] + b[1]
            ...
            >>> ds.window_all(TumblingEventTimeWindows.of(Time.seconds(5))) \\
            ...   .aggregate(AverageAggregate(),
            ...              accumulator_type=Types.TUPLE([Types.LONG(), Types.LONG()]),
            ...              output_type=Types.DOUBLE())

        :param aggregate_function: The aggregation function that is used for incremental
                                   aggregation.
        :param window_function: The window function.
        :param accumulator_type: Type information for the internal accumulator type of the
                                 aggregation function.
        :param output_type: Type information for the result type of the window function.
        :return: The data stream that is the result of applying the window function to the window.

        .. versionadded:: 1.16.0
        """
        if window_function is None:
            internal_window_function: InternalWindowFunction = InternalSingleValueAllWindowFunction(
                PassThroughAllWindowFunction())
        elif isinstance(window_function, AllWindowFunction):
            internal_window_function = InternalSingleValueAllWindowFunction(window_function)
        elif isinstance(window_function, ProcessAllWindowFunction):
            internal_window_function = InternalSingleValueProcessAllWindowFunction(window_function)
        else:
            raise TypeError("window_function should be a AllWindowFunction or "
                            "ProcessAllWindowFunction")

        if accumulator_type is None:
            accumulator_type = Types.PICKLED_BYTE_ARRAY()
        elif isinstance(accumulator_type, list):
            accumulator_type = RowTypeInfo(accumulator_type)

        aggregating_state_descriptor = AggregatingStateDescriptor(WINDOW_STATE_NAME,
                                                                  aggregate_function,
                                                                  accumulator_type)
        func_desc = type(aggregate_function).__name__
        if window_function is not None:
            func_desc = "%s, %s" % (func_desc, type(window_function).__name__)
        return self._get_result_data_stream(internal_window_function,
                                            aggregating_state_descriptor,
                                            func_desc,
                                            output_type)

    def apply(self,
              window_function: AllWindowFunction,
              output_type: TypeInformation = None) -> DataStream:
        """
        Applies the given window function to each window. The window function is called for each
        evaluation of the window. The output of the window function is interpreted as a regular
        non-windowed stream.

        Note that this function requires that all data in the windows is buffered until the window
        is evaluated, as the function provides no means of incremental aggregation.

        :param window_function: The window function.
        :param output_type: Type information for the result type of the window function.
        :return: The data stream that is the result of applying the window function to the window.
        """
        internal_window_function: InternalWindowFunction = InternalIterableAllWindowFunction(
            window_function)
        list_state_descriptor = ListStateDescriptor(WINDOW_STATE_NAME, self.get_input_type())
        func_desc = type(window_function).__name__
        return self._get_result_data_stream(internal_window_function,
                                            list_state_descriptor,
                                            func_desc,
                                            output_type)

    def process(self,
                process_window_function: ProcessAllWindowFunction,
                output_type: TypeInformation = None) -> DataStream:
        """
        Applies the given window function to each window. The window function is called for each
        evaluation of the window for each key individually. The output of the window function is
        interpreted as a regular non-windowed stream.

        Note that this function requires that all data in the windows is buffered until the window
        is evaluated, as the function provides no means of incremental aggregation.

        :param process_window_function: The window function.
        :param output_type: Type information for the result type of the window function.
        :return: The data stream that is the result of applying the window function to the window.
        """
        internal_window_function: InternalWindowFunction = InternalIterableProcessAllWindowFunction(
            process_window_function)
        list_state_descriptor = ListStateDescriptor(WINDOW_STATE_NAME, self.get_input_type())
        func_desc = type(process_window_function).__name__
        return self._get_result_data_stream(internal_window_function,
                                            list_state_descriptor,
                                            func_desc,
                                            output_type)

    def _get_result_data_stream(self,
                                internal_window_function: InternalWindowFunction,
                                window_state_descriptor: StateDescriptor,
                                func_desc: str,
                                output_type: TypeInformation):
        if self._window_trigger is None:
            self._window_trigger = self._window_assigner.get_default_trigger(
                self.get_execution_environment())
        window_serializer = self._window_assigner.get_window_serializer()
        window_operation_descriptor = WindowOperationDescriptor(
            self._window_assigner,
            self._window_trigger,
            self._allowed_lateness,
            self._late_data_output_tag,
            window_state_descriptor,
            window_serializer,
            internal_window_function)

        from pyflink.fn_execution import flink_fn_execution_pb2
        j_python_data_stream_function_operator, j_output_type_info = \
            _get_one_input_stream_operator(
                self._keyed_stream,
                window_operation_descriptor,
                flink_fn_execution_pb2.UserDefinedDataStreamFunction.WINDOW,  # type: ignore
                output_type)

        op_name = window_operation_descriptor.generate_op_name()
        op_desc = window_operation_descriptor.generate_op_desc("AllWindow", func_desc)
        return DataStream(self._keyed_stream._j_data_stream.transform(
            op_name,
            j_output_type_info,
            j_python_data_stream_function_operator)).set_description(op_desc)


class ConnectedStreams(object):
    """
    ConnectedStreams represent two connected streams of (possibly) different data types.
    Connected streams are useful for cases where operations on one stream directly
    affect the operations on the other stream, usually via shared state between the streams.

    An example for the use of connected streams would be to apply rules that change over time
    onto another stream. One of the connected streams has the rules, the other stream the
    elements to apply the rules to. The operation on the connected stream maintains the
    current set of rules in the state. It may receive either a rule update and update the state
    or a data element and apply the rules in the state to the element.

    The connected stream can be conceptually viewed as a union stream of an Either type, that
    holds either the first stream's type or the second stream's type.
    """

    def __init__(self, stream1: DataStream, stream2: DataStream):
        self.stream1 = stream1
        self.stream2 = stream2

    def key_by(self, key_selector1: Union[Callable, KeySelector],
               key_selector2: Union[Callable, KeySelector],
               key_type: TypeInformation = None) -> 'ConnectedStreams':
        """
        KeyBy operation for connected data stream. Assigns keys to the elements of
        input1 and input2 using keySelector1 and keySelector2 with explicit type information
        for the common key type.

        :param key_selector1: The `KeySelector` used for grouping the first input.
        :param key_selector2: The `KeySelector` used for grouping the second input.
        :param key_type: The type information of the common key type
        :return: The partitioned `ConnectedStreams`
        """

        ds1 = self.stream1
        ds2 = self.stream2
        if isinstance(self.stream1, KeyedStream):
            ds1 = self.stream1._origin_stream
        if isinstance(self.stream2, KeyedStream):
            ds2 = self.stream2._origin_stream
        return ConnectedStreams(
            ds1.key_by(key_selector1, key_type),
            ds2.key_by(key_selector2, key_type))

    def map(self, func: CoMapFunction, output_type: TypeInformation = None) -> 'DataStream':
        """
        Applies a CoMap transformation on a `ConnectedStreams` and maps the output to a common
        type. The transformation calls a `CoMapFunction.map1` for each element of the first
        input and `CoMapFunction.map2` for each element of the second input. Each CoMapFunction
        call returns exactly one element.

        :param func: The CoMapFunction used to jointly transform the two input DataStreams
        :param output_type: `TypeInformation` for the result type of the function.
        :return: The transformed `DataStream`
        """
        if not isinstance(func, CoMapFunction):
            raise TypeError("The input function must be a CoMapFunction!")

        if self._is_keyed_stream():
            class CoMapKeyedCoProcessFunctionAdapter(KeyedCoProcessFunction):
                def __init__(self, co_map_func: CoMapFunction):
                    self._open_func = co_map_func.open
                    self._close_func = co_map_func.close
                    self._map1_func = co_map_func.map1
                    self._map2_func = co_map_func.map2

                def open(self, runtime_context: RuntimeContext):
                    self._open_func(runtime_context)

                def close(self):
                    self._close_func()

                def process_element1(self, value, ctx: 'KeyedCoProcessFunction.Context'):
                    result = self._map1_func(value)
                    if result is not None:
                        yield result

                def process_element2(self, value, ctx: 'KeyedCoProcessFunction.Context'):
                    result = self._map2_func(value)
                    if result is not None:
                        yield result

            return self.process(CoMapKeyedCoProcessFunctionAdapter(func), output_type) \
                .name("Co-Map")
        else:
            class CoMapCoProcessFunctionAdapter(CoProcessFunction):
                def __init__(self, co_map_func: CoMapFunction):
                    self._open_func = co_map_func.open
                    self._close_func = co_map_func.close
                    self._map1_func = co_map_func.map1
                    self._map2_func = co_map_func.map2

                def open(self, runtime_context: RuntimeContext):
                    self._open_func(runtime_context)

                def close(self):
                    self._close_func()

                def process_element1(self, value, ctx: 'CoProcessFunction.Context'):
                    result = self._map1_func(value)
                    if result is not None:
                        yield result

                def process_element2(self, value, ctx: 'CoProcessFunction.Context'):
                    result = self._map2_func(value)
                    if result is not None:
                        yield result

            return self.process(CoMapCoProcessFunctionAdapter(func), output_type) \
                .name("Co-Map")

    def flat_map(self, func: CoFlatMapFunction, output_type: TypeInformation = None) \
            -> 'DataStream':
        """
        Applies a CoFlatMap transformation on a `ConnectedStreams` and maps the output to a
        common type. The transformation calls a `CoFlatMapFunction.flatMap1` for each element
        of the first input and `CoFlatMapFunction.flatMap2` for each element of the second
        input. Each CoFlatMapFunction call returns any number of elements including none.

        :param func: The CoFlatMapFunction used to jointly transform the two input DataStreams
        :param output_type: `TypeInformation` for the result type of the function.
        :return: The transformed `DataStream`
        """

        if not isinstance(func, CoFlatMapFunction):
            raise TypeError("The input must be a CoFlatMapFunction!")

        if self._is_keyed_stream():
            class FlatMapKeyedCoProcessFunctionAdapter(KeyedCoProcessFunction):

                def __init__(self, co_flat_map_func: CoFlatMapFunction):
                    self._open_func = co_flat_map_func.open
                    self._close_func = co_flat_map_func.close
                    self._flat_map1_func = co_flat_map_func.flat_map1
                    self._flat_map2_func = co_flat_map_func.flat_map2

                def open(self, runtime_context: RuntimeContext):
                    self._open_func(runtime_context)

                def close(self):
                    self._close_func()

                def process_element1(self, value, ctx: 'KeyedCoProcessFunction.Context'):
                    result = self._flat_map1_func(value)
                    if result:
                        yield from result

                def process_element2(self, value, ctx: 'KeyedCoProcessFunction.Context'):
                    result = self._flat_map2_func(value)
                    if result:
                        yield from result

            return self.process(FlatMapKeyedCoProcessFunctionAdapter(func), output_type) \
                .name("Co-Flat Map")
        else:
            class FlatMapCoProcessFunctionAdapter(CoProcessFunction):

                def __init__(self, co_flat_map_func: CoFlatMapFunction):
                    self._open_func = co_flat_map_func.open
                    self._close_func = co_flat_map_func.close
                    self._flat_map1_func = co_flat_map_func.flat_map1
                    self._flat_map2_func = co_flat_map_func.flat_map2

                def open(self, runtime_context: RuntimeContext):
                    self._open_func(runtime_context)

                def close(self):
                    self._close_func()

                def process_element1(self, value, ctx: 'CoProcessFunction.Context'):
                    result = self._flat_map1_func(value)
                    if result:
                        yield from result

                def process_element2(self, value, ctx: 'CoProcessFunction.Context'):
                    result = self._flat_map2_func(value)
                    if result:
                        yield from result

            return self.process(FlatMapCoProcessFunctionAdapter(func), output_type) \
                .name("Co-Flat Map")

    def process(self,
                func: Union[CoProcessFunction, KeyedCoProcessFunction],
                output_type: TypeInformation = None) -> 'DataStream':
        if not isinstance(func, CoProcessFunction) and not isinstance(func, KeyedCoProcessFunction):
            raise TypeError("The input must be a CoProcessFunction or KeyedCoProcessFunction!")

        from pyflink.fn_execution.flink_fn_execution_pb2 import UserDefinedDataStreamFunction
        if self._is_keyed_stream():
            func_type = UserDefinedDataStreamFunction.KEYED_CO_PROCESS  # type: ignore
            func_name = "Keyed Co-Process"
        else:
            func_type = UserDefinedDataStreamFunction.CO_PROCESS  # type: ignore
            func_name = "Co-Process"

        j_connected_stream = self.stream1._j_data_stream.connect(self.stream2._j_data_stream)
        j_operator, j_output_type = _get_two_input_stream_operator(
            self,
            func,
            func_type,
            output_type)
        return DataStream(j_connected_stream.transform(func_name, j_output_type, j_operator))

    def _is_keyed_stream(self):
        return isinstance(self.stream1, KeyedStream) and isinstance(self.stream2, KeyedStream)


class BroadcastStream(object):
    """
    A BroadcastStream is a stream with :class:`state.BroadcastState` (s). This can be created by any
    stream using the :meth:`DataStream.broadcast` method and implicitly creates states where the
    user can store elements of the created :class:`BroadcastStream`. (see
    :class:`BroadcastConnectedStream`).

    Note that no further operation can be applied to these streams. The only available option is
    to connect them with a keyed or non-keyed stream, using the :meth:`KeyedStream.connect` and the
    :meth:`DataStream.connect` respectively. Applying these methods will result it a
    :class:`BroadcastConnectedStream` for further processing.

    .. versionadded:: 1.16.0
    """

    def __init__(
        self,
        input_stream: Union['DataStream', 'KeyedStream'],
        broadcast_state_descriptors: List[MapStateDescriptor],
    ):
        self.input_stream = input_stream
        self.broadcast_state_descriptors = broadcast_state_descriptors


class BroadcastConnectedStream(object):
    """
    A BroadcastConnectedStream represents the result of connecting a keyed or non-keyed stream, with
    a :class:`BroadcastStream` with :class:`~state.BroadcastState` (s). As in the case of
    :class:`ConnectedStreams` these streams are useful for cases where operations on one stream
    directly affect the operations on the other stream, usually via shared state between the
    streams.

    An example for the use of such connected streams would be to apply rules that change over time
    onto another, possibly keyed stream. The stream with the broadcast state has the rules, and will
    store them in the broadcast state, while the other stream will contain the elements to apply the
    rules to. By broadcasting the rules, these will be available in all parallel instances, and can
    be applied to all partitions of the other stream.

    .. versionadded:: 1.16.0
    """

    def __init__(
        self,
        non_broadcast_stream: Union['DataStream', 'KeyedStream'],
        broadcast_stream: 'BroadcastStream',
        broadcast_state_descriptors: List[MapStateDescriptor],
    ):
        self.non_broadcast_stream = non_broadcast_stream
        self.broadcast_stream = broadcast_stream
        self.broadcast_state_descriptors = broadcast_state_descriptors

    @overload
    def process(
        self,
        func: BroadcastProcessFunction,
        output_type: TypeInformation = None,
    ) -> 'DataStream':
        pass

    @overload
    def process(
        self,
        func: KeyedBroadcastProcessFunction,
        output_type: TypeInformation = None
    ) -> 'DataStream':
        pass

    def process(
        self,
        func: Union[BroadcastProcessFunction, KeyedBroadcastProcessFunction],
        output_type: TypeInformation = None,
    ) -> 'DataStream':
        """
        Assumes as inputs a :class:`BroadcastStream` and a :class:`DataStream` or
        :class:`KeyedStream` and applies the given :class:`BroadcastProcessFunction` or
        :class:`KeyedBroadcastProcessFunction` on them, thereby creating a transformed output
        stream.

        :param func: The :class:`BroadcastProcessFunction` that is called for each element in the
            non-broadcasted :class:`DataStream`, or the :class:`KeyedBroadcastProcessFunction` that
            is called for each element in the non-broadcasted :class:`KeyedStream`.
        :param output_type: The type of the output elements, should be
            :class:`common.TypeInformation` or list (implicit :class:`RowTypeInfo`) or None (
            implicit :meth:`Types.PICKLED_BYTE_ARRAY`).
        :return: The transformed :class:`DataStream`.
        """
        if isinstance(func, BroadcastProcessFunction) and self._is_keyed_stream():
            raise TypeError("BroadcastProcessFunction should be applied to non-keyed DataStream")
        if isinstance(func, KeyedBroadcastProcessFunction) and (not self._is_keyed_stream()):
            raise TypeError("KeyedBroadcastProcessFunction should be applied to keyed DataStream")

        j_input_transformation1 = self.non_broadcast_stream._j_data_stream.getTransformation()
        j_input_transformation2 = (
            self.broadcast_stream.input_stream._j_data_stream.getTransformation()
        )

        if output_type is None:
            output_type_info: TypeInformation = Types.PICKLED_BYTE_ARRAY()
        elif isinstance(output_type, list):
            output_type_info = RowTypeInfo(output_type)
        elif isinstance(output_type, TypeInformation):
            output_type_info = output_type
        else:
            raise TypeError("output_type must be None, list or TypeInformation")
        j_output_type = output_type_info.get_java_type_info()

        from pyflink.fn_execution.flink_fn_execution_pb2 import UserDefinedDataStreamFunction
        jvm = get_gateway().jvm
        JPythonConfigUtil = jvm.org.apache.flink.python.util.PythonConfigUtil

        if self._is_keyed_stream():
            func_type = UserDefinedDataStreamFunction.KEYED_CO_BROADCAST_PROCESS  # type: ignore
            func_name = "Keyed-Co-Process-Broadcast"
        else:
            func_type = UserDefinedDataStreamFunction.CO_BROADCAST_PROCESS  # type: ignore
            func_name = "Co-Process-Broadcast"

        j_state_names = to_jarray(
            jvm.String, [i.get_name() for i in self.broadcast_state_descriptors]
        )
        j_state_descriptors = JPythonConfigUtil.convertStateNamesToStateDescriptors(j_state_names)
        j_conf = get_j_env_configuration(
            self.broadcast_stream.input_stream._j_data_stream.getExecutionEnvironment())
        j_data_stream_python_function_info = _create_j_data_stream_python_function_info(
            func, func_type
        )
        j_env = (
            self.non_broadcast_stream.get_execution_environment()._j_stream_execution_environment
        )

        if self._is_keyed_stream():
            JTransformation = jvm.org.apache.flink.streaming.api.transformations.python \
                .PythonKeyedBroadcastStateTransformation
            j_transformation = JTransformation(
                func_name,
                j_conf,
                j_data_stream_python_function_info,
                j_input_transformation1,
                j_input_transformation2,
                j_state_descriptors,
                self.non_broadcast_stream._j_data_stream.getKeyType(),
                self.non_broadcast_stream._j_data_stream.getKeySelector(),
                j_output_type,
                j_env.getParallelism(),
            )
        else:
            JTransformation = jvm.org.apache.flink.streaming.api.transformations.python \
                .PythonBroadcastStateTransformation
            j_transformation = JTransformation(
                func_name,
                j_conf,
                j_data_stream_python_function_info,
                j_input_transformation1,
                j_input_transformation2,
                j_state_descriptors,
                j_output_type,
                j_env.getParallelism(),
            )
        j_env.addOperator(j_transformation)
        j_data_stream = JPythonConfigUtil.createSingleOutputStreamOperator(j_env, j_transformation)

        return DataStream(j_data_stream)

    def _is_keyed_stream(self):
        return isinstance(self.non_broadcast_stream, KeyedStream)


def _get_one_input_stream_operator(data_stream: DataStream,
                                   func: Union[Function,
                                               FunctionWrapper,
                                               WindowOperationDescriptor,
                                               AsyncFunctionDescriptor],
                                   func_type: int,
                                   output_type: Union[TypeInformation, List] = None):
    """
    Create a Java one input stream operator.

    :param func: a function object that implements the Function interface.
    :param func_type: function type, supports MAP, FLAT_MAP, etc.
    :param output_type: the data type of the function output data.
    :return: A Java operator which is responsible for execution user defined python function.
    """

    gateway = get_gateway()

    j_input_types = data_stream._j_data_stream.getTransformation().getOutputType()
    if output_type is None:
        output_type_info: TypeInformation = Types.PICKLED_BYTE_ARRAY()
    elif isinstance(output_type, list):
        output_type_info = RowTypeInfo(output_type)
    else:
        output_type_info = output_type

    j_data_stream_python_function_info = _create_j_data_stream_python_function_info(func, func_type)
    j_output_type_info = output_type_info.get_java_type_info()
    j_conf = get_j_env_configuration(data_stream._j_data_stream.getExecutionEnvironment())
    python_execution_mode = (
        j_conf.get(gateway.jvm.org.apache.flink.python.PythonOptions.PYTHON_EXECUTION_MODE))

    from pyflink.fn_execution.flink_fn_execution_pb2 import UserDefinedDataStreamFunction
    if func_type == UserDefinedDataStreamFunction.PROCESS:  # type: ignore
        if python_execution_mode == 'thread':
            JDataStreamPythonFunctionOperator = gateway.jvm.EmbeddedPythonProcessOperator
        else:
            JDataStreamPythonFunctionOperator = gateway.jvm.ExternalPythonProcessOperator
    elif func_type == UserDefinedDataStreamFunction.KEYED_PROCESS:  # type: ignore
        if python_execution_mode == 'thread':
            JDataStreamPythonFunctionOperator = gateway.jvm.EmbeddedPythonKeyedProcessOperator
        else:
            JDataStreamPythonFunctionOperator = gateway.jvm.ExternalPythonKeyedProcessOperator
    elif func_type == UserDefinedDataStreamFunction.WINDOW:  # type: ignore
        window_serializer = typing.cast(WindowOperationDescriptor, func).window_serializer
        if isinstance(window_serializer, TimeWindowSerializer):
            j_namespace_serializer = \
                gateway.jvm.org.apache.flink.table.runtime.operators.window.TimeWindow.Serializer()
        elif isinstance(window_serializer, CountWindowSerializer):
            j_namespace_serializer = \
                gateway.jvm.org.apache.flink.table.runtime.operators.window.CountWindow.Serializer()
        elif isinstance(window_serializer, GlobalWindowSerializer):
            j_namespace_serializer = \
                gateway.jvm.org.apache.flink.table.runtime.operators.window.GlobalWindow \
                .Serializer()
        else:
            j_namespace_serializer = \
                gateway.jvm.org.apache.flink.streaming.api.utils.ByteArrayWrapperSerializer()
        if python_execution_mode == 'thread':
            JDataStreamPythonWindowFunctionOperator = gateway.jvm.EmbeddedPythonWindowOperator
        else:
            JDataStreamPythonWindowFunctionOperator = gateway.jvm.ExternalPythonKeyedProcessOperator

        j_python_function_operator = JDataStreamPythonWindowFunctionOperator(
            j_conf,
            j_data_stream_python_function_info,
            j_input_types,
            j_output_type_info,
            j_namespace_serializer)
        return j_python_function_operator, j_output_type_info
    else:
        raise TypeError("Unsupported function type: %s" % func_type)

    j_python_function_operator = JDataStreamPythonFunctionOperator(
        j_conf,
        j_data_stream_python_function_info,
        j_input_types,
        j_output_type_info)

    return j_python_function_operator, j_output_type_info


def _get_two_input_stream_operator(connected_streams: ConnectedStreams,
                                   func: Union[Function, FunctionWrapper],
                                   func_type: int,
                                   type_info: TypeInformation):
    """
    Create a Java two input stream operator.

    :param func: a function object that implements the Function interface.
    :param func_type: function type, supports MAP, FLAT_MAP, etc.
    :param type_info: the data type of the function output data.
    :return: A Java operator which is responsible for execution user defined python function.
    """
    gateway = get_gateway()

    j_input_types1 = connected_streams.stream1._j_data_stream.getTransformation().getOutputType()
    j_input_types2 = connected_streams.stream2._j_data_stream.getTransformation().getOutputType()

    if type_info is None:
        output_type_info: TypeInformation = Types.PICKLED_BYTE_ARRAY()
    elif isinstance(type_info, list):
        output_type_info = RowTypeInfo(type_info)
    else:
        output_type_info = type_info

    j_data_stream_python_function_info = _create_j_data_stream_python_function_info(func, func_type)
    j_output_type_info = output_type_info.get_java_type_info()
    j_conf = get_j_env_configuration(
        connected_streams.stream1._j_data_stream.getExecutionEnvironment())
    python_execution_mode = (
        j_conf.get(gateway.jvm.org.apache.flink.python.PythonOptions.PYTHON_EXECUTION_MODE))

    from pyflink.fn_execution.flink_fn_execution_pb2 import UserDefinedDataStreamFunction
    if func_type == UserDefinedDataStreamFunction.CO_PROCESS:  # type: ignore
        if python_execution_mode == 'thread':
            JTwoInputPythonFunctionOperator = gateway.jvm.EmbeddedPythonCoProcessOperator
        else:
            JTwoInputPythonFunctionOperator = gateway.jvm.ExternalPythonCoProcessOperator
    elif func_type == UserDefinedDataStreamFunction.KEYED_CO_PROCESS:  # type: ignore
        if python_execution_mode == 'thread':
            JTwoInputPythonFunctionOperator = gateway.jvm.EmbeddedPythonKeyedCoProcessOperator
        else:
            JTwoInputPythonFunctionOperator = gateway.jvm.ExternalPythonKeyedCoProcessOperator
    else:
        raise TypeError("Unsupported function type: %s" % func_type)

    j_python_data_stream_function_operator = JTwoInputPythonFunctionOperator(
        j_conf,
        j_data_stream_python_function_info,
        j_input_types1,
        j_input_types2,
        j_output_type_info)

    return j_python_data_stream_function_operator, j_output_type_info


def _create_j_data_stream_python_function_info(
    func: Union[Function, FunctionWrapper, WindowOperationDescriptor, AsyncFunctionDescriptor],
    func_type: int
) -> bytes:
    gateway = get_gateway()

    import cloudpickle
    serialized_func = cloudpickle.dumps(func)

    j_data_stream_python_function = gateway.jvm.DataStreamPythonFunction(
        bytearray(serialized_func), _get_python_env()
    )
    return gateway.jvm.DataStreamPythonFunctionInfo(j_data_stream_python_function, func_type)


class CloseableIterator(object):
    """
    Representing an Iterator that is also auto closeable.
    """

    def __init__(self, j_closeable_iterator, type_info: TypeInformation = None):
        self._j_closeable_iterator = j_closeable_iterator
        self._type_info = type_info

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def next(self):
        if not self._j_closeable_iterator.hasNext():
            raise StopIteration('No more data.')
        return convert_to_python_obj(self._j_closeable_iterator.next(), self._type_info)

    def close(self):
        self._j_closeable_iterator.close()