You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by di...@apache.org on 2021/08/12 05:19:01 UTC
[flink] branch master updated (c54084a -> 54beaaf)
This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git.
from c54084a [FLINK-23645][table-sql] Fix SqlClient doesn't response correctly to Signal.INT (#16743)
new 1b8bcf2 [hotfix][python][tests] Refactor the test cases in DataStreamTests to reduce the number of itcases
new 54beaaf [FLINK-23616][python] Support to chain the Python DataStream operators as much as possible
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../shortcodes/generated/python_configuration.html | 6 +
flink-python/pyflink/datastream/data_stream.py | 2 +-
.../datastream/stream_execution_environment.py | 10 +-
.../pyflink/datastream/tests/test_data_stream.py | 608 ++++++++-------------
.../fn_execution/beam/beam_operations_fast.pxd | 30 +-
.../fn_execution/beam/beam_operations_fast.pyx | 123 ++++-
.../fn_execution/beam/beam_operations_slow.py | 62 ++-
.../fn_execution/datastream/input_handler.py | 12 +-
.../pyflink/fn_execution/datastream/operations.py | 93 ++--
.../pyflink/fn_execution/flink_fn_execution_pb2.py | 38 +-
.../pyflink/proto/flink-fn-execution.proto | 1 +
.../java/org/apache/flink/python/Constants.java | 1 +
.../org/apache/flink/python/PythonOptions.java | 9 +
.../chain/PythonOperatorChainingOptimizer.java | 448 +++++++++++++++
.../apache/flink/python/util/PythonConfigUtil.java | 35 +-
.../python/DataStreamPythonFunctionInfo.java | 18 +-
.../AbstractDataStreamPythonFunctionOperator.java | 71 +++
.../AbstractOneInputPythonFunctionOperator.java | 164 +++++-
.../AbstractTwoInputPythonFunctionOperator.java | 173 +++++-
.../python/OneInputPythonFunctionOperator.java | 214 --------
.../operators/python/PythonCoProcessOperator.java | 17 +-
.../python/PythonKeyedCoProcessOperator.java | 16 +-
.../python/PythonKeyedProcessOperator.java | 16 +-
.../operators/python/PythonProcessOperator.java | 14 +-
.../python/TwoInputPythonFunctionOperator.java | 219 --------
.../beam/BeamDataStreamPythonFunctionRunner.java | 149 +++--
.../python/beam/BeamPythonFunctionRunner.java | 16 +-
.../flink/streaming/api/utils/ProtoUtils.java | 54 +-
.../AbstractOneInputPythonFunctionOperator.java | 3 +-
.../python/AbstractStatelessFunctionOperator.java | 6 +-
.../AbstractPythonStreamAggregateOperator.java | 4 +-
.../python/beam/BeamTablePythonFunctionRunner.java | 67 ++-
.../chain/PythonOperatorChainingOptimizerTest.java | 544 ++++++++++++++++++
33 files changed, 2207 insertions(+), 1036 deletions(-)
create mode 100644 flink-python/src/main/java/org/apache/flink/python/chain/PythonOperatorChainingOptimizer.java
create mode 100644 flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractDataStreamPythonFunctionOperator.java
delete mode 100644 flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/OneInputPythonFunctionOperator.java
delete mode 100644 flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/TwoInputPythonFunctionOperator.java
copy flink-python/src/main/java/org/apache/flink/{streaming/api => table/runtime}/operators/python/AbstractOneInputPythonFunctionOperator.java (91%)
create mode 100644 flink-python/src/test/java/org/apache/flink/python/chain/PythonOperatorChainingOptimizerTest.java
[flink] 02/02: [FLINK-23616][python] Support to chain the Python
DataStream operators as much as possible
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 54beaafff580bb3a65ffe58796b2ca8306631079
Author: Dian Fu <di...@apache.org>
AuthorDate: Mon May 10 19:16:39 2021 +0800
[FLINK-23616][python] Support to chain the Python DataStream operators as much as possible
This closes #16777.
---
.../shortcodes/generated/python_configuration.html | 6 +
flink-python/pyflink/datastream/data_stream.py | 2 +-
.../datastream/stream_execution_environment.py | 10 +-
.../pyflink/datastream/tests/test_data_stream.py | 41 +-
.../fn_execution/beam/beam_operations_fast.pxd | 30 +-
.../fn_execution/beam/beam_operations_fast.pyx | 123 ++++-
.../fn_execution/beam/beam_operations_slow.py | 62 ++-
.../fn_execution/datastream/input_handler.py | 12 +-
.../pyflink/fn_execution/datastream/operations.py | 93 ++--
.../pyflink/fn_execution/flink_fn_execution_pb2.py | 38 +-
.../pyflink/proto/flink-fn-execution.proto | 1 +
.../java/org/apache/flink/python/Constants.java | 1 +
.../org/apache/flink/python/PythonOptions.java | 9 +
.../chain/PythonOperatorChainingOptimizer.java | 448 +++++++++++++++++
.../apache/flink/python/util/PythonConfigUtil.java | 35 +-
.../python/DataStreamPythonFunctionInfo.java | 18 +-
.../AbstractDataStreamPythonFunctionOperator.java | 71 +++
.../AbstractOneInputPythonFunctionOperator.java | 164 ++++++-
.../AbstractTwoInputPythonFunctionOperator.java | 173 ++++++-
.../python/OneInputPythonFunctionOperator.java | 214 --------
.../operators/python/PythonCoProcessOperator.java | 17 +-
.../python/PythonKeyedCoProcessOperator.java | 16 +-
.../python/PythonKeyedProcessOperator.java | 16 +-
.../operators/python/PythonProcessOperator.java | 14 +-
.../python/TwoInputPythonFunctionOperator.java | 219 ---------
.../beam/BeamDataStreamPythonFunctionRunner.java | 149 ++++--
.../python/beam/BeamPythonFunctionRunner.java | 16 +-
.../flink/streaming/api/utils/ProtoUtils.java | 54 +-
.../AbstractOneInputPythonFunctionOperator.java | 3 +-
.../python/AbstractStatelessFunctionOperator.java | 6 +-
.../AbstractPythonStreamAggregateOperator.java | 4 +-
.../python/beam/BeamTablePythonFunctionRunner.java | 67 ++-
.../chain/PythonOperatorChainingOptimizerTest.java | 544 +++++++++++++++++++++
33 files changed, 2004 insertions(+), 672 deletions(-)
diff --git a/docs/layouts/shortcodes/generated/python_configuration.html b/docs/layouts/shortcodes/generated/python_configuration.html
index ccf4e54..b600141 100644
--- a/docs/layouts/shortcodes/generated/python_configuration.html
+++ b/docs/layouts/shortcodes/generated/python_configuration.html
@@ -81,6 +81,12 @@
<td>When it is false, metric for Python will be disabled. You can disable the metric to achieve better performance at some circumstance.</td>
</tr>
<tr>
+ <td><h5>python.operator-chaining.enabled</h5></td>
+ <td style="word-wrap: break-word;">true</td>
+ <td>Boolean</td>
+ <td>Python operator chaining allows non-shuffle operations to be co-located in the same thread fully avoiding serialization and de-serialization.</td>
+ </tr>
+ <tr>
<td><h5>python.profile.enabled</h5></td>
<td style="word-wrap: break-word;">false</td>
<td>Boolean</td>
diff --git a/flink-python/pyflink/datastream/data_stream.py b/flink-python/pyflink/datastream/data_stream.py
index 4676eca..d1156f2 100644
--- a/flink-python/pyflink/datastream/data_stream.py
+++ b/flink-python/pyflink/datastream/data_stream.py
@@ -68,7 +68,7 @@ class DataStream(object):
def name(self, name: str) -> 'DataStream':
"""
Sets the name of the current data stream. This name is used by the visualization and logging
- during runting.
+ during runtime.
:param name: Name of the stream.
:return: The named operator.
diff --git a/flink-python/pyflink/datastream/stream_execution_environment.py b/flink-python/pyflink/datastream/stream_execution_environment.py
index a83796e..0bb331d 100644
--- a/flink-python/pyflink/datastream/stream_execution_environment.py
+++ b/flink-python/pyflink/datastream/stream_execution_environment.py
@@ -737,7 +737,6 @@ class StreamExecutionEnvironment(object):
"""
j_stream_graph = self._generate_stream_graph(clear_transformations=True, job_name=job_name)
-
return JobExecutionResult(self._j_stream_execution_environment.execute(j_stream_graph))
def execute_async(self, job_name: str = 'Flink Streaming Job') -> JobClient:
@@ -752,7 +751,6 @@ class StreamExecutionEnvironment(object):
submission succeeded.
"""
j_stream_graph = self._generate_stream_graph(clear_transformations=True, job_name=job_name)
-
j_job_client = self._j_stream_execution_environment.executeAsync(j_stream_graph)
return JobClient(j_job_client=j_job_client)
@@ -769,7 +767,6 @@ class StreamExecutionEnvironment(object):
:return: The execution plan of the program, as a JSON String.
"""
j_stream_graph = self._generate_stream_graph(False)
-
return j_stream_graph.getStreamingPlanAsJSON()
@staticmethod
@@ -920,9 +917,12 @@ class StreamExecutionEnvironment(object):
def _generate_stream_graph(self, clear_transformations: bool = False, job_name: str = None) \
-> JavaObject:
- j_stream_graph = get_gateway().jvm \
+ gateway = get_gateway()
+ gateway.jvm.org.apache.flink.python.chain.PythonOperatorChainingOptimizer.apply(
+ self._j_stream_execution_environment)
+ j_stream_graph = gateway.jvm \
.org.apache.flink.python.util.PythonConfigUtil.generateStreamGraphWithDependencies(
- self._j_stream_execution_environment, clear_transformations)
+ self._j_stream_execution_environment, clear_transformations)
if job_name is not None:
j_stream_graph.setJobName(job_name)
diff --git a/flink-python/pyflink/datastream/tests/test_data_stream.py b/flink-python/pyflink/datastream/tests/test_data_stream.py
index 57dc80d..093e578 100644
--- a/flink-python/pyflink/datastream/tests/test_data_stream.py
+++ b/flink-python/pyflink/datastream/tests/test_data_stream.py
@@ -197,16 +197,17 @@ class DataStreamTests(object):
ds1.connect(ds2) \
.key_by(lambda x: x[0], lambda x: x[0]) \
.process(MyKeyedCoProcessFunction()) \
+ .map(lambda x: Row(x[0], x[1] + 1)) \
.add_sink(self.test_sink)
self.env.execute('test_keyed_co_process_function')
results = self.test_sink.get_results(True)
- expected = ["<Row('a', 1)>",
- "<Row('b', 1)>",
+ expected = ["<Row('a', 2)>",
"<Row('b', 2)>",
- "<Row('c', 1)>",
+ "<Row('b', 3)>",
"<Row('c', 2)>",
- "<Row('d', 1)>",
- "<Row('on_timer', 3)>"]
+ "<Row('c', 3)>",
+ "<Row('d', 2)>",
+ "<Row('on_timer', 4)>"]
self.assert_equals_sorted(expected, results)
def test_keyed_co_map(self):
@@ -246,14 +247,13 @@ class DataStreamTests(object):
ds1.connect(ds2)\
.key_by(MyKeySelector(), MyKeySelector(), key_type=Types.INT())\
.map(AssertKeyCoMapFunction())\
+ .map(lambda x: (x[0], x[1] + 1)) \
.add_sink(self.test_sink)
self.env.execute()
results = self.test_sink.get_results(True)
- expected = ["Row(f0='e', f1=2)", "Row(f0='a', f1=0)", "Row(f0='b', f1=0)",
- "Row(f0='c', f1=1)", "Row(f0='d', f1=1)", "Row(f0='e', f1=2)",
- "Row(f0='a', f1=0)", "Row(f0='b', f1=0)", "Row(f0='c', f1=1)",
- "Row(f0='d', f1=1)"]
+ expected = ["('e', 3)", "('a', 1)", "('b', 1)", "('c', 2)", "('d', 2)", "('e', 3)",
+ "('a', 1)", "('b', 1)", "('c', 2)", "('d', 2)"]
self.assert_equals_sorted(expected, results)
def test_keyed_co_flat_map(self):
@@ -263,10 +263,11 @@ class DataStreamTests(object):
type_info=Types.ROW([Types.STRING(), Types.STRING()]))
ds1.connect(ds2).key_by(lambda x: 1, lambda x: 1) \
.flat_map(MyRichCoFlatMapFunction(), output_type=Types.STRING()) \
+ .filter(lambda x: x != '4') \
.add_sink(self.test_sink)
self.env.execute('test_keyed_co_flat_map')
results = self.test_sink.get_results(False)
- expected = ['2', '2', '3', '3', '4', '4', 'a', 'b', 'c']
+ expected = ['2', '2', '3', '3', 'a', 'b', 'c']
self.assert_equals_sorted(expected, results)
def test_execute_and_collect(self):
@@ -361,11 +362,12 @@ class DataStreamTests(object):
state_value += 1
return state_value
- keyed_stream.map(AssertKeyMapFunction()).add_sink(self.test_sink)
+ keyed_stream.map(AssertKeyMapFunction())\
+ .map(lambda x: (x[0], x[1] + 1))\
+ .add_sink(self.test_sink)
self.env.execute('test_keyed_map')
results = self.test_sink.get_results(True)
- expected = ["Row(f0='e', f1=2)", "Row(f0='a', f1=0)", "Row(f0='b', f1=1)",
- "Row(f0='c', f1=0)", "Row(f0='d', f1=1)"]
+ expected = ["('e', 3)", "('a', 1)", "('b', 2)", "('c', 1)", "('d', 2)"]
self.assert_equals_sorted(expected, results)
def test_keyed_flat_map(self):
@@ -403,11 +405,12 @@ class DataStreamTests(object):
self.state.update(state_value)
yield value
- keyed_stream.flat_map(AssertKeyMapFunction()).add_sink(self.test_sink)
+ keyed_stream.flat_map(AssertKeyMapFunction())\
+ .map(lambda x: (x[0], x[1] + 1))\
+ .add_sink(self.test_sink)
self.env.execute('test_keyed_flat_map')
results = self.test_sink.get_results(True)
- expected = ["Row(f0='e', f1=2)", "Row(f0='a', f1=0)", "Row(f0='b', f1=0)",
- "Row(f0='c', f1=1)", "Row(f0='d', f1=1)"]
+ expected = ["('e', 3)", "('a', 1)", "('b', 1)", "('c', 2)", "('d', 2)"]
self.assert_equals_sorted(expected, results)
def test_keyed_filter(self):
@@ -447,10 +450,12 @@ class DataStreamTests(object):
self.state.update(state_value)
return True
- keyed_stream.filter(AssertKeyFilterFunction()).add_sink(self.test_sink)
+ keyed_stream.filter(AssertKeyFilterFunction())\
+ .filter(lambda x: x[1] > 0)\
+ .add_sink(self.test_sink)
self.env.execute('key_by_test')
results = self.test_sink.get_results(False)
- expected = ['+I[a, 0]', '+I[c, 1]', '+I[e, 2]']
+ expected = ['+I[c, 1]', '+I[e, 2]']
self.assert_equals_sorted(expected, results)
def test_multi_key_by(self):
diff --git a/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pxd b/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pxd
index 73dae74..dd5fd71 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pxd
+++ b/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pxd
@@ -19,11 +19,34 @@
from apache_beam.coders.coder_impl cimport StreamCoderImpl
from apache_beam.runners.worker.operations cimport Operation
+from apache_beam.utils.windowed_value cimport WindowedValue
+
+from pyflink.fn_execution.coder_impl_fast cimport InputStreamWrapper
+
+cdef class InputProcessor:
+ cpdef has_next(self)
+ cpdef next(self)
+
+cdef class NetworkInputProcessor(InputProcessor):
+ cdef InputStreamWrapper _input_stream_wrapper
+
+cdef class IntermediateInputProcessor(InputProcessor):
+ cdef object _input_values
+ cdef object _next_value
+
+cdef class OutputProcessor:
+ cdef Operation _consumer
+ cpdef process_outputs(self, WindowedValue windowed_value, results)
+
+cdef class NetworkOutputProcessor(OutputProcessor):
+ cdef StreamCoderImpl _value_coder_impl
+
+cdef class IntermediateOutputProcessor(OutputProcessor):
+ pass
cdef class FunctionOperation(Operation):
- cdef Operation consumer
+ cdef OutputProcessor _output_processor
cdef bint _is_python_coder
- cdef StreamCoderImpl _value_coder_impl
cdef object process_element
cdef object operation
cdef object operation_cls
@@ -34,6 +57,7 @@ cdef class StatelessFunctionOperation(FunctionOperation):
pass
cdef class StatefulFunctionOperation(FunctionOperation):
- cdef object keyed_state_backend
+ cdef object _keyed_state_backend
+ cdef WindowedValue _reusable_windowed_value
cpdef void add_timer_info(self, timer_family_id, timer_info)
cpdef process_timer(self, tag, timer_data)
diff --git a/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx b/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx
index 19ab6fe..06fbce3 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx
+++ b/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx
@@ -21,13 +21,83 @@
# cython: boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True
from libc.stdint cimport *
from apache_beam.coders.coder_impl cimport OutputStream as BOutputStream
+from apache_beam.utils cimport windowed_value
from apache_beam.utils.windowed_value cimport WindowedValue
from pyflink.fn_execution.coder_impl_fast cimport InputStreamWrapper
+from apache_beam.runners.worker.bundle_processor import DataOutputOperation
from pyflink.fn_execution.beam.beam_coder_impl_fast import FlinkLengthPrefixCoderBeamWrapper
from pyflink.fn_execution.table.operations import BundleOperation
from pyflink.fn_execution.profiler import Profiler
+
+cdef class InputProcessor:
+
+ cpdef has_next(self):
+ pass
+
+ cpdef next(self):
+ pass
+
+
+cdef class NetworkInputProcessor(InputProcessor):
+
+ def __init__(self, InputStreamWrapper input_stream_wrapper):
+ self._input_stream_wrapper = input_stream_wrapper
+
+ cpdef has_next(self):
+ return self._input_stream_wrapper.has_next()
+
+ cpdef next(self):
+ return self._input_stream_wrapper.next()
+
+
+cdef class IntermediateInputProcessor(InputProcessor):
+
+ def __init__(self, input_values):
+ self._input_values = input_values
+ self._next_value = None
+
+ cpdef has_next(self):
+ try:
+ self._next_value = next(self._input_values)
+ except StopIteration:
+ self._next_value = None
+
+ return self._next_value is not None
+
+ cpdef next(self):
+ return self._next_value
+
+
+cdef class OutputProcessor:
+
+ cpdef process_outputs(self, WindowedValue windowed_value, results):
+ pass
+
+
+cdef class NetworkOutputProcessor(OutputProcessor):
+
+ def __init__(self, consumer):
+ assert isinstance(consumer, DataOutputOperation)
+ self._consumer = consumer
+ self._value_coder_impl = consumer.windowed_coder.wrapped_value_coder.get_impl()
+
+ cpdef process_outputs(self, WindowedValue windowed_value, results):
+ output_stream = self._consumer.output_stream
+ self._value_coder_impl.encode_to_stream(results, output_stream, True)
+ output_stream.maybe_flush()
+
+
+cdef class IntermediateOutputProcessor(OutputProcessor):
+
+ def __init__(self, consumer):
+ self._consumer = consumer
+
+ cpdef process_outputs(self, WindowedValue windowed_value, results):
+ self._consumer.process(windowed_value.with_value(results))
+
+
cdef class FunctionOperation(Operation):
"""
Base class of function operation that will execute StatelessFunction or StatefulFunction for
@@ -36,13 +106,18 @@ cdef class FunctionOperation(Operation):
def __init__(self, name, spec, counter_factory, sampler, consumers, operation_cls):
super(FunctionOperation, self).__init__(name, spec, counter_factory, sampler)
- self.consumer = consumers['output'][0]
- self._value_coder_impl = self.consumer.windowed_coder.wrapped_value_coder.get_impl()._value_coder
+ consumer = consumers['output'][0]
+ if isinstance(consumer, DataOutputOperation):
+ self._output_processor = NetworkOutputProcessor(consumer)
- if isinstance(self._value_coder_impl, FlinkLengthPrefixCoderBeamWrapper):
- self._is_python_coder = False
+ _value_coder_impl = consumer.windowed_coder.wrapped_value_coder.get_impl()._value_coder
+ if isinstance(_value_coder_impl, FlinkLengthPrefixCoderBeamWrapper):
+ self._is_python_coder = False
+ else:
+ self._is_python_coder = True
else:
- self._is_python_coder = True
+ self._output_processor = IntermediateOutputProcessor(consumer)
+ self._is_python_coder = False
self.operation_cls = operation_cls
self.operation = self.generate_operation()
@@ -72,27 +147,24 @@ cdef class FunctionOperation(Operation):
cpdef process(self, WindowedValue o):
cdef InputStreamWrapper input_stream_wrapper
- cdef BOutputStream output_stream
- output_stream = self.consumer.output_stream
+ cdef InputProcessor input_processor
with self.scoped_process_state:
if self._is_python_coder:
for value in o.value:
- self._value_coder_impl.encode_to_stream(
- self.process_element(value), output_stream, True)
- output_stream.maybe_flush()
+ self._output_processor.process_outputs(o, self.process_element(value))
else:
- input_stream_wrapper = o.value
+ if isinstance(o.value, InputStreamWrapper):
+ input_processor = NetworkInputProcessor(o.value)
+ else:
+ input_processor = IntermediateInputProcessor(o.value)
if isinstance(self.operation, BundleOperation):
- while input_stream_wrapper.has_next():
- self.process_element(input_stream_wrapper.next())
- result = self.operation.finish_bundle()
- self._value_coder_impl.encode_to_stream(
- result, output_stream, True)
+ while input_processor.has_next():
+ self.process_element(input_processor.next())
+ self._output_processor.process_outputs(o, self.operation.finish_bundle())
else:
- while input_stream_wrapper.has_next():
- result = self.process_element(input_stream_wrapper.next())
- self._value_coder_impl.encode_to_stream(
- result, output_stream, True)
+ while input_processor.has_next():
+ result = self.process_element(input_processor.next())
+ self._output_processor.process_outputs(o, result)
def progress_metrics(self):
metrics = super(FunctionOperation, self).progress_metrics()
@@ -126,12 +198,13 @@ cdef class StatelessFunctionOperation(FunctionOperation):
cdef class StatefulFunctionOperation(FunctionOperation):
def __init__(self, name, spec, counter_factory, sampler, consumers, operation_cls,
keyed_state_backend):
- self.keyed_state_backend = keyed_state_backend
+ self._keyed_state_backend = keyed_state_backend
+ self._reusable_windowed_value = windowed_value.create(None, -1, None, None)
super(StatefulFunctionOperation, self).__init__(
name, spec, counter_factory, sampler, consumers, operation_cls)
cdef object generate_operation(self):
- return self.operation_cls(self.spec, self.keyed_state_backend)
+ return self.operation_cls(self.spec, self._keyed_state_backend)
cpdef void add_timer_info(self, timer_family_id, timer_info):
# ignore timer_family_id
@@ -139,7 +212,7 @@ cdef class StatefulFunctionOperation(FunctionOperation):
cpdef process_timer(self, tag, timer_data):
cdef BOutputStream output_stream
- output_stream = self.consumer.output_stream
- self._value_coder_impl.encode_to_stream(
+ self._output_processor.process_outputs(
+ self._reusable_windowed_value,
# the field user_key holds the timer data
- self.operation.process_timer(timer_data.user_key), output_stream, True)
+ self.operation.process_timer(timer_data.user_key))
diff --git a/flink-python/pyflink/fn_execution/beam/beam_operations_slow.py b/flink-python/pyflink/fn_execution/beam/beam_operations_slow.py
index 3065c4a..580e3f8 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_operations_slow.py
+++ b/flink-python/pyflink/fn_execution/beam/beam_operations_slow.py
@@ -15,16 +15,48 @@
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
+import abc
from abc import abstractmethod
+from typing import Iterable, Any
-from apache_beam.runners.worker.bundle_processor import TimerInfo
+from apache_beam.runners.worker.bundle_processor import TimerInfo, DataOutputOperation
from apache_beam.runners.worker.operations import Operation
+from apache_beam.utils import windowed_value
from apache_beam.utils.windowed_value import WindowedValue
from pyflink.fn_execution.table.operations import BundleOperation
from pyflink.fn_execution.profiler import Profiler
+class OutputProcessor(abc.ABC):
+
+ @abstractmethod
+ def process_outputs(self, windowed_value: WindowedValue, results: Iterable[Any]):
+ pass
+
+
+class NetworkOutputProcessor(OutputProcessor):
+
+ def __init__(self, consumer):
+ assert isinstance(consumer, DataOutputOperation)
+ self._consumer = consumer
+ self._value_coder_impl = consumer.windowed_coder.wrapped_value_coder.get_impl()
+
+ def process_outputs(self, windowed_value: WindowedValue, results: Iterable[Any]):
+ output_stream = self._consumer.output_stream
+ self._value_coder_impl.encode_to_stream(results, output_stream, True)
+ output_stream.maybe_flush()
+
+
+class IntermediateOutputProcessor(OutputProcessor):
+
+ def __init__(self, consumer):
+ self._consumer = consumer
+
+ def process_outputs(self, windowed_value: WindowedValue, results: Iterable[Any]):
+ self._consumer.process(windowed_value.with_value(results))
+
+
class FunctionOperation(Operation):
"""
Base class of function operation that will execute StatelessFunction or StatefulFunction for
@@ -33,8 +65,11 @@ class FunctionOperation(Operation):
def __init__(self, name, spec, counter_factory, sampler, consumers, operation_cls):
super(FunctionOperation, self).__init__(name, spec, counter_factory, sampler)
- self.consumer = consumers['output'][0]
- self._value_coder_impl = self.consumer.windowed_coder.wrapped_value_coder.get_impl()
+ consumer = consumers['output'][0]
+ if isinstance(consumer, DataOutputOperation):
+ self._output_processor = NetworkOutputProcessor(consumer)
+ else:
+ self._output_processor = IntermediateOutputProcessor(consumer)
self.operation_cls = operation_cls
self.operation = self.generate_operation()
self.process_element = self.operation.process_element
@@ -81,18 +116,13 @@ class FunctionOperation(Operation):
def process(self, o: WindowedValue):
with self.scoped_process_state:
- output_stream = self.consumer.output_stream
if isinstance(self.operation, BundleOperation):
for value in o.value:
self.process_element(value)
- self._value_coder_impl.encode_to_stream(
- self.operation.finish_bundle(), output_stream, True)
- output_stream.maybe_flush()
+ self._output_processor.process_outputs(o, self.operation.finish_bundle())
else:
for value in o.value:
- self._value_coder_impl.encode_to_stream(
- self.process_element(value), output_stream, True)
- output_stream.maybe_flush()
+ self._output_processor.process_outputs(o, self.process_element(value))
def monitoring_infos(self, transform_id, tag_to_pcollection_id):
"""
@@ -118,20 +148,20 @@ class StatelessFunctionOperation(FunctionOperation):
class StatefulFunctionOperation(FunctionOperation):
def __init__(self, name, spec, counter_factory, sampler, consumers, operation_cls,
keyed_state_backend):
- self.keyed_state_backend = keyed_state_backend
+ self._keyed_state_backend = keyed_state_backend
+ self._reusable_windowed_value = windowed_value.create(None, -1, None, None)
super(StatefulFunctionOperation, self).__init__(
name, spec, counter_factory, sampler, consumers, operation_cls)
def generate_operation(self):
- return self.operation_cls(self.spec, self.keyed_state_backend)
+ return self.operation_cls(self.spec, self._keyed_state_backend)
def add_timer_info(self, timer_family_id: str, timer_info: TimerInfo):
# ignore timer_family_id
self.operation.add_timer_info(timer_info)
def process_timer(self, tag, timer_data):
- output_stream = self.consumer.output_stream
- self._value_coder_impl.encode_to_stream(
+ self._output_processor.process_outputs(
+ self._reusable_windowed_value,
# the field user_key holds the timer data
- self.operation.process_timer(timer_data.user_key), output_stream, True)
- output_stream.maybe_flush()
+ self.operation.process_timer(timer_data.user_key))
diff --git a/flink-python/pyflink/fn_execution/datastream/input_handler.py b/flink-python/pyflink/fn_execution/datastream/input_handler.py
index f450f24..6722ba6 100644
--- a/flink-python/pyflink/fn_execution/datastream/input_handler.py
+++ b/flink-python/pyflink/fn_execution/datastream/input_handler.py
@@ -44,7 +44,7 @@ class RunnerInputHandler(ABC):
watermark = value[1]
data = value[2]
self._advance_watermark(watermark)
- yield from _emit_results(timestamp, self._process_element_func(data, timestamp))
+ yield from _emit_results(timestamp, watermark, self._process_element_func(data, timestamp))
def _advance_watermark(self, watermark: int) -> None:
self._internal_timer_service.advance_watermark(watermark)
@@ -78,9 +78,11 @@ class TimerHandler(ABC):
else:
namespace = None
if timer_type == TimerType.EVENT_TIME.value:
- yield from _emit_results(timestamp, self._on_event_time(timestamp, key, namespace))
+ yield from _emit_results(
+ timestamp, watermark, self._on_event_time(timestamp, key, namespace))
elif timer_type == TimerType.PROCESSING_TIME.value:
- yield from _emit_results(timestamp, self._on_processing_time(timestamp, key, namespace))
+ yield from _emit_results(
+ timestamp, watermark, self._on_processing_time(timestamp, key, namespace))
else:
raise Exception("Unsupported timer type: %d" % timer_type)
@@ -94,7 +96,7 @@ class TimerHandler(ABC):
self._internal_timer_service.advance_watermark(watermark)
-def _emit_results(timestamp, results):
+def _emit_results(timestamp, watermark, results):
if results:
for result in results:
- yield Row(timestamp, result)
+ yield Row(timestamp, watermark, result)
diff --git a/flink-python/pyflink/fn_execution/datastream/operations.py b/flink-python/pyflink/fn_execution/datastream/operations.py
index 8c077ab..3d36b30 100644
--- a/flink-python/pyflink/fn_execution/datastream/operations.py
+++ b/flink-python/pyflink/fn_execution/datastream/operations.py
@@ -137,56 +137,73 @@ def extract_stateless_function(user_defined_function_proto, runtime_context: Run
:param runtime_context: the streaming runtime context
"""
func_type = user_defined_function_proto.function_type
- user_defined_func = pickle.loads(user_defined_function_proto.payload)
-
- def open_func():
- if hasattr(user_defined_func, "open"):
- user_defined_func.open(runtime_context)
+ UserDefinedDataStreamFunction = flink_fn_execution_pb2.UserDefinedDataStreamFunction
- def close_func():
- if hasattr(user_defined_func, "close"):
- user_defined_func.close()
+ if func_type == UserDefinedDataStreamFunction.REVISE_OUTPUT:
+ def open_func():
+ pass
- UserDefinedDataStreamFunction = flink_fn_execution_pb2.UserDefinedDataStreamFunction
- if func_type == UserDefinedDataStreamFunction.PROCESS:
- process_element = user_defined_func.process_element
- ctx = InternalProcessFunctionContext(NonKeyedTimerServiceImpl())
+ def close_func():
+ pass
- def wrapped_func(value):
+ def revise_output(value):
# VALUE[CURRENT_TIMESTAMP, CURRENT_WATERMARK, NORMAL_DATA]
timestamp = value[0]
- watermark = value[1]
- ctx.set_timestamp(timestamp)
- ctx.timer_service().advance_watermark(watermark)
- results = process_element(value[2], ctx)
- yield from _emit_results(timestamp, results)
+ element = value[2]
+ yield Row(timestamp, element)
- process_element_func = wrapped_func
+ process_element_func = revise_output
- elif func_type == UserDefinedDataStreamFunction.CO_PROCESS:
- process_element1 = user_defined_func.process_element1
- process_element2 = user_defined_func.process_element2
- ctx = InternalProcessFunctionContext(NonKeyedTimerServiceImpl())
+ else:
+ user_defined_func = pickle.loads(user_defined_function_proto.payload)
- def wrapped_func(value):
- # VALUE[CURRENT_TIMESTAMP, CURRENT_WATERMARK, [isLeft, leftInput, rightInput]]
- timestamp = value[0]
- watermark = value[1]
- ctx.set_timestamp(timestamp)
- ctx.timer_service().advance_watermark(watermark)
+ def open_func():
+ if hasattr(user_defined_func, "open"):
+ user_defined_func.open(runtime_context)
+
+ def close_func():
+ if hasattr(user_defined_func, "close"):
+ user_defined_func.close()
- normal_data = value[2]
- if normal_data[0]:
- results = process_element1(normal_data[1], ctx)
- else:
- results = process_element2(normal_data[2], ctx)
+ if func_type == UserDefinedDataStreamFunction.PROCESS:
+ process_element = user_defined_func.process_element
+ ctx = InternalProcessFunctionContext(NonKeyedTimerServiceImpl())
- yield from _emit_results(timestamp, results)
+ def wrapped_func(value):
+ # VALUE[CURRENT_TIMESTAMP, CURRENT_WATERMARK, NORMAL_DATA]
+ timestamp = value[0]
+ watermark = value[1]
+ ctx.set_timestamp(timestamp)
+ ctx.timer_service().advance_watermark(watermark)
+ results = process_element(value[2], ctx)
+ yield from _emit_results(timestamp, watermark, results)
- process_element_func = wrapped_func
+ process_element_func = wrapped_func
- else:
- raise Exception("Unsupported function_type: " + str(func_type))
+ elif func_type == UserDefinedDataStreamFunction.CO_PROCESS:
+ process_element1 = user_defined_func.process_element1
+ process_element2 = user_defined_func.process_element2
+ ctx = InternalProcessFunctionContext(NonKeyedTimerServiceImpl())
+
+ def wrapped_func(value):
+ # VALUE[CURRENT_TIMESTAMP, CURRENT_WATERMARK, [isLeft, leftInput, rightInput]]
+ timestamp = value[0]
+ watermark = value[1]
+ ctx.set_timestamp(timestamp)
+ ctx.timer_service().advance_watermark(watermark)
+
+ normal_data = value[2]
+ if normal_data[0]:
+ results = process_element1(normal_data[1], ctx)
+ else:
+ results = process_element2(normal_data[2], ctx)
+
+ yield from _emit_results(timestamp, watermark, results)
+
+ process_element_func = wrapped_func
+
+ else:
+ raise Exception("Unsupported function_type: " + str(func_type))
return open_func, close_func, process_element_func
diff --git a/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.py b/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.py
index 4196765..894afaf 100644
--- a/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.py
+++ b/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.py
@@ -36,7 +36,7 @@ DESCRIPTOR = _descriptor.FileDescriptor(
name='flink-fn-execution.proto',
package='org.apache.flink.fn_execution.v1',
syntax='proto3',
- serialized_pb=_b('\n\x18\x66link-fn-execution.proto\x12 org.apache.flink.fn_execution.v1\"\x86\x01\n\x05Input\x12\x44\n\x03udf\x18\x01 \x01(\x0b\x32\x35.org.apache.flink.fn_execution.v1.UserDefinedFunctionH\x00\x12\x15\n\x0binputOffset\x18\x02 \x01(\x05H\x00\x12\x17\n\rinputConstant\x18\x03 \x01(\x0cH\x00\x42\x07\n\x05input\"\xa8\x01\n\x13UserDefinedFunction\x12\x0f\n\x07payload\x18\x01 \x01(\x0c\x12\x37\n\x06inputs\x18\x02 \x03(\x0b\x32\'.org.apache.flink.fn_execution.v1.Input\x12\x14 [...]
+ serialized_pb=_b('\n\x18\x66link-fn-execution.proto\x12 org.apache.flink.fn_execution.v1\"\x86\x01\n\x05Input\x12\x44\n\x03udf\x18\x01 \x01(\x0b\x32\x35.org.apache.flink.fn_execution.v1.UserDefinedFunctionH\x00\x12\x15\n\x0binputOffset\x18\x02 \x01(\x05H\x00\x12\x17\n\rinputConstant\x18\x03 \x01(\x0cH\x00\x42\x07\n\x05input\"\xa8\x01\n\x13UserDefinedFunction\x12\x0f\n\x07payload\x18\x01 \x01(\x0c\x12\x37\n\x06inputs\x18\x02 \x03(\x0b\x32\'.org.apache.flink.fn_execution.v1.Input\x12\x14 [...]
)
@@ -373,11 +373,15 @@ _USERDEFINEDDATASTREAMFUNCTION_FUNCTIONTYPE = _descriptor.EnumDescriptor(
name='WINDOW', index=4, number=4,
options=None,
type=None),
+ _descriptor.EnumValueDescriptor(
+ name='REVISE_OUTPUT', index=5, number=100,
+ options=None,
+ type=None),
],
containing_type=None,
options=None,
serialized_start=6739,
- serialized_end=6835,
+ serialized_end=6854,
)
_sym_db.RegisterEnumDescriptor(_USERDEFINEDDATASTREAMFUNCTION_FUNCTIONTYPE)
@@ -398,8 +402,8 @@ _CODERINFODESCRIPTOR_MODE = _descriptor.EnumDescriptor(
],
containing_type=None,
options=None,
- serialized_start=7802,
- serialized_end=7834,
+ serialized_start=7821,
+ serialized_end=7853,
)
_sym_db.RegisterEnumDescriptor(_CODERINFODESCRIPTOR_MODE)
@@ -1897,7 +1901,7 @@ _USERDEFINEDDATASTREAMFUNCTION = _descriptor.Descriptor(
oneofs=[
],
serialized_start=5984,
- serialized_end=6835,
+ serialized_end=6854,
)
@@ -1927,8 +1931,8 @@ _CODERINFODESCRIPTOR_FLATTENROWTYPE = _descriptor.Descriptor(
extension_ranges=[],
oneofs=[
],
- serialized_start=7431,
- serialized_end=7505,
+ serialized_start=7450,
+ serialized_end=7524,
)
_CODERINFODESCRIPTOR_ROWTYPE = _descriptor.Descriptor(
@@ -1957,8 +1961,8 @@ _CODERINFODESCRIPTOR_ROWTYPE = _descriptor.Descriptor(
extension_ranges=[],
oneofs=[
],
- serialized_start=7507,
- serialized_end=7574,
+ serialized_start=7526,
+ serialized_end=7593,
)
_CODERINFODESCRIPTOR_ARROWTYPE = _descriptor.Descriptor(
@@ -1987,8 +1991,8 @@ _CODERINFODESCRIPTOR_ARROWTYPE = _descriptor.Descriptor(
extension_ranges=[],
oneofs=[
],
- serialized_start=7576,
- serialized_end=7645,
+ serialized_start=7595,
+ serialized_end=7664,
)
_CODERINFODESCRIPTOR_OVERWINDOWARROWTYPE = _descriptor.Descriptor(
@@ -2017,8 +2021,8 @@ _CODERINFODESCRIPTOR_OVERWINDOWARROWTYPE = _descriptor.Descriptor(
extension_ranges=[],
oneofs=[
],
- serialized_start=7647,
- serialized_end=7726,
+ serialized_start=7666,
+ serialized_end=7745,
)
_CODERINFODESCRIPTOR_RAWTYPE = _descriptor.Descriptor(
@@ -2047,8 +2051,8 @@ _CODERINFODESCRIPTOR_RAWTYPE = _descriptor.Descriptor(
extension_ranges=[],
oneofs=[
],
- serialized_start=7728,
- serialized_end=7800,
+ serialized_start=7747,
+ serialized_end=7819,
)
_CODERINFODESCRIPTOR = _descriptor.Descriptor(
@@ -2123,8 +2127,8 @@ _CODERINFODESCRIPTOR = _descriptor.Descriptor(
name='data_type', full_name='org.apache.flink.fn_execution.v1.CoderInfoDescriptor.data_type',
index=0, containing_type=None, fields=[]),
],
- serialized_start=6838,
- serialized_end=7847,
+ serialized_start=6857,
+ serialized_end=7866,
)
_INPUT.fields_by_name['udf'].message_type = _USERDEFINEDFUNCTION
diff --git a/flink-python/pyflink/proto/flink-fn-execution.proto b/flink-python/pyflink/proto/flink-fn-execution.proto
index fc77342..d4c449f 100644
--- a/flink-python/pyflink/proto/flink-fn-execution.proto
+++ b/flink-python/pyflink/proto/flink-fn-execution.proto
@@ -346,6 +346,7 @@ message UserDefinedDataStreamFunction {
KEYED_PROCESS = 2;
KEYED_CO_PROCESS = 3;
WINDOW = 4;
+ REVISE_OUTPUT = 100;
}
message JobParameter {
diff --git a/flink-python/src/main/java/org/apache/flink/python/Constants.java b/flink-python/src/main/java/org/apache/flink/python/Constants.java
index 2bd9dd9..0827166 100644
--- a/flink-python/src/main/java/org/apache/flink/python/Constants.java
+++ b/flink-python/src/main/java/org/apache/flink/python/Constants.java
@@ -32,6 +32,7 @@ public class Constants {
public static final String TRANSFORM_ID = "transform";
public static final String MAIN_INPUT_NAME = "input";
public static final String MAIN_OUTPUT_NAME = "output";
+ public static final String WINDOW_STRATEGY = "windowing_strategy";
public static final String INPUT_COLLECTION_ID = "input";
public static final String OUTPUT_COLLECTION_ID = "output";
diff --git a/flink-python/src/main/java/org/apache/flink/python/PythonOptions.java b/flink-python/src/main/java/org/apache/flink/python/PythonOptions.java
index 7e57df7..8055efe 100644
--- a/flink-python/src/main/java/org/apache/flink/python/PythonOptions.java
+++ b/flink-python/src/main/java/org/apache/flink/python/PythonOptions.java
@@ -76,6 +76,15 @@ public class PythonOptions {
+ "The interval between each profiling is determined by the config options "
+ "python.fn-execution.bundle.size and python.fn-execution.bundle.time.");
+ /** The configuration to enable or disable python operator chaining. */
+ public static final ConfigOption<Boolean> PYTHON_OPERATOR_CHAINING_ENABLED =
+ ConfigOptions.key("python.operator-chaining.enabled")
+ .booleanType()
+ .defaultValue(true)
+ .withDescription(
+ "Python operator chaining allows non-shuffle operations to be co-located in the "
+ + "same thread fully avoiding serialization and de-serialization.");
+
public static final ConfigOption<String> PYTHON_FILES =
ConfigOptions.key("python.files")
.stringType()
diff --git a/flink-python/src/main/java/org/apache/flink/python/chain/PythonOperatorChainingOptimizer.java b/flink-python/src/main/java/org/apache/flink/python/chain/PythonOperatorChainingOptimizer.java
new file mode 100644
index 0000000..ca1e1b3
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/python/chain/PythonOperatorChainingOptimizer.java
@@ -0,0 +1,448 @@
+/*
+ * 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.
+ */
+
+package org.apache.flink.python.chain;
+
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.dag.Transformation;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.memory.ManagedMemoryUseCase;
+import org.apache.flink.python.PythonOptions;
+import org.apache.flink.python.util.PythonConfigUtil;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.operators.SimpleOperatorFactory;
+import org.apache.flink.streaming.api.operators.SourceOperatorFactory;
+import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
+import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
+import org.apache.flink.streaming.api.operators.python.AbstractDataStreamPythonFunctionOperator;
+import org.apache.flink.streaming.api.operators.python.AbstractOneInputPythonFunctionOperator;
+import org.apache.flink.streaming.api.operators.python.PythonCoProcessOperator;
+import org.apache.flink.streaming.api.operators.python.PythonKeyedCoProcessOperator;
+import org.apache.flink.streaming.api.operators.python.PythonKeyedProcessOperator;
+import org.apache.flink.streaming.api.operators.python.PythonProcessOperator;
+import org.apache.flink.streaming.api.transformations.AbstractBroadcastStateTransformation;
+import org.apache.flink.streaming.api.transformations.AbstractMultipleInputTransformation;
+import org.apache.flink.streaming.api.transformations.FeedbackTransformation;
+import org.apache.flink.streaming.api.transformations.LegacySinkTransformation;
+import org.apache.flink.streaming.api.transformations.OneInputTransformation;
+import org.apache.flink.streaming.api.transformations.PartitionTransformation;
+import org.apache.flink.streaming.api.transformations.PhysicalTransformation;
+import org.apache.flink.streaming.api.transformations.ReduceTransformation;
+import org.apache.flink.streaming.api.transformations.SideOutputTransformation;
+import org.apache.flink.streaming.api.transformations.SinkTransformation;
+import org.apache.flink.streaming.api.transformations.TimestampsAndWatermarksTransformation;
+import org.apache.flink.streaming.api.transformations.TwoInputTransformation;
+import org.apache.flink.streaming.api.transformations.UnionTransformation;
+import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner;
+
+import org.apache.flink.shaded.guava30.com.google.common.collect.Sets;
+
+import org.apache.commons.compress.utils.Lists;
+
+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * A util class which attempts to chain all available Python functions.
+ *
+ * <p>An operator could be chained to it's predecessor if all of the following conditions are met:
+ *
+ * <ul>
+ * <li>Both of them are Python operators
+ * <li>The parallelism, the maximum parallelism and the slot sharing group are all the same
+ * <li>The chaining strategy is ChainingStrategy.ALWAYS and the chaining strategy of the
+ * predecessor isn't ChainingStrategy.NEVER
+ * <li>This partitioner between them is ForwardPartitioner
+ * </ul>
+ *
+ * <p>The properties of the generated chained operator are as following:
+ *
+ * <ul>
+ * <li>The name is the concatenation of all the names of the chained operators
+ * <li>The parallelism, the maximum parallelism and the slot sharing group are from one of the
+ * chained operators as all of them are the same between the chained operators
+ * <li>The chaining strategy is the same as the head operator
+ * <li>The uid and the uidHash are the same as the head operator
+ * </ul>
+ */
+public class PythonOperatorChainingOptimizer {
+
+ private final List<Transformation<?>> transformations;
+
+ private final boolean isChainingEnabled;
+
+ // key-value pairs where the value is the output transformations of the key transformation
+ private final transient Map<Transformation<?>, Set<Transformation<?>>> outputMap =
+ new HashMap<>();
+
+ public PythonOperatorChainingOptimizer(
+ List<Transformation<?>> transformations, boolean isChainingEnabled) {
+ this.transformations = checkNotNull(transformations);
+ this.isChainingEnabled = isChainingEnabled;
+ buildOutputMap();
+ }
+
+ private void buildOutputMap() {
+ for (Transformation<?> transformation : transformations) {
+ for (Transformation<?> input : transformation.getInputs()) {
+ Set<Transformation<?>> outputs =
+ outputMap.computeIfAbsent(input, i -> Sets.newHashSet());
+ outputs.add(transformation);
+ }
+ }
+ }
+
+ @VisibleForTesting
+ public List<Transformation<?>> optimize() throws Exception {
+ List<Transformation<?>> optimizedTransformations = Lists.newArrayList();
+ Set<Transformation<?>> removedTransformations = Sets.newIdentityHashSet();
+ for (Transformation<?> transformation : transformations) {
+ if (!removedTransformations.contains(transformation)) {
+ ChainInfo chainInfo = chainWithInputIfPossible(transformation);
+ optimizedTransformations.add(chainInfo.newTransformation);
+ optimizedTransformations.removeAll(chainInfo.oldTransformations);
+ removedTransformations.addAll(chainInfo.oldTransformations);
+ }
+ }
+ return optimizedTransformations;
+ }
+
+ private ChainInfo chainWithInputIfPossible(Transformation<?> transform) throws Exception {
+ ChainInfo chainInfo = null;
+ if (transform instanceof OneInputTransformation
+ && PythonConfigUtil.isPythonDataStreamOperator(transform)) {
+ Transformation<?> input = transform.getInputs().get(0);
+ while (!PythonConfigUtil.isPythonDataStreamOperator(input)) {
+ if (input instanceof PartitionTransformation
+ && ((PartitionTransformation<?>) input).getPartitioner()
+ instanceof ForwardPartitioner) {
+ input = input.getInputs().get(0);
+ } else {
+ return ChainInfo.of(transform);
+ }
+ }
+
+ if (isChainable(input, transform)) {
+ Transformation<?> chainedTransformation =
+ createChainedTransformation(input, transform);
+ Set<Transformation<?>> outputTransformations = outputMap.get(transform);
+ if (outputTransformations != null) {
+ for (Transformation<?> output : outputTransformations) {
+ replaceInput(output, transform, chainedTransformation);
+ }
+ }
+ chainInfo = ChainInfo.of(chainedTransformation, Arrays.asList(input, transform));
+ }
+ }
+
+ if (chainInfo == null) {
+ chainInfo = ChainInfo.of(transform);
+ }
+
+ return chainInfo;
+ }
+
+ private Transformation<?> createChainedTransformation(
+ Transformation<?> upTransform, Transformation<?> downTransform) {
+ AbstractDataStreamPythonFunctionOperator<?> upOperator =
+ (AbstractDataStreamPythonFunctionOperator<?>)
+ ((SimpleOperatorFactory<?>) getOperatorFactory(upTransform)).getOperator();
+ PythonProcessOperator<?, ?> downOperator =
+ (PythonProcessOperator<?, ?>)
+ ((SimpleOperatorFactory<?>) getOperatorFactory(downTransform))
+ .getOperator();
+
+ DataStreamPythonFunctionInfo upPythonFunctionInfo = upOperator.getPythonFunctionInfo();
+ DataStreamPythonFunctionInfo downPythonFunctionInfo = downOperator.getPythonFunctionInfo();
+
+ DataStreamPythonFunctionInfo chainedPythonFunctionInfo =
+ new DataStreamPythonFunctionInfo(
+ downPythonFunctionInfo.getPythonFunction(),
+ upPythonFunctionInfo,
+ downPythonFunctionInfo.getFunctionType());
+
+ AbstractDataStreamPythonFunctionOperator<?> chainedOperator =
+ upOperator.copy(chainedPythonFunctionInfo, downOperator.getProducedType());
+
+ PhysicalTransformation<?> chainedTransformation;
+ if (upOperator instanceof AbstractOneInputPythonFunctionOperator) {
+ chainedTransformation =
+ new OneInputTransformation(
+ upTransform.getInputs().get(0),
+ upTransform.getName() + ", " + downTransform.getName(),
+ (OneInputStreamOperator<?, ?>) chainedOperator,
+ downTransform.getOutputType(),
+ upTransform.getParallelism());
+
+ ((OneInputTransformation<?, ?>) chainedTransformation)
+ .setStateKeySelector(
+ ((OneInputTransformation) upTransform).getStateKeySelector());
+ ((OneInputTransformation<?, ?>) chainedTransformation)
+ .setStateKeyType(
+ ((OneInputTransformation<?, ?>) upTransform).getStateKeyType());
+ } else {
+ chainedTransformation =
+ new TwoInputTransformation(
+ upTransform.getInputs().get(0),
+ upTransform.getInputs().get(1),
+ upTransform.getName() + ", " + downTransform.getName(),
+ (TwoInputStreamOperator<?, ?, ?>) chainedOperator,
+ downTransform.getOutputType(),
+ upTransform.getParallelism());
+
+ ((TwoInputTransformation<?, ?, ?>) chainedTransformation)
+ .setStateKeySelectors(
+ ((TwoInputTransformation) upTransform).getStateKeySelector1(),
+ ((TwoInputTransformation) upTransform).getStateKeySelector2());
+ ((TwoInputTransformation<?, ?, ?>) chainedTransformation)
+ .setStateKeyType(
+ ((TwoInputTransformation<?, ?, ?>) upTransform).getStateKeyType());
+ }
+
+ chainedTransformation.setUid(upTransform.getUid());
+ if (upTransform.getUserProvidedNodeHash() != null) {
+ chainedTransformation.setUidHash(upTransform.getUserProvidedNodeHash());
+ }
+
+ for (ManagedMemoryUseCase useCase : upTransform.getManagedMemorySlotScopeUseCases()) {
+ chainedTransformation.declareManagedMemoryUseCaseAtSlotScope(useCase);
+ }
+ for (ManagedMemoryUseCase useCase : downTransform.getManagedMemorySlotScopeUseCases()) {
+ chainedTransformation.declareManagedMemoryUseCaseAtSlotScope(useCase);
+ }
+
+ for (Map.Entry<ManagedMemoryUseCase, Integer> useCase :
+ upTransform.getManagedMemoryOperatorScopeUseCaseWeights().entrySet()) {
+ chainedTransformation.declareManagedMemoryUseCaseAtOperatorScope(
+ useCase.getKey(), useCase.getValue());
+ }
+ for (Map.Entry<ManagedMemoryUseCase, Integer> useCase :
+ downTransform.getManagedMemoryOperatorScopeUseCaseWeights().entrySet()) {
+ chainedTransformation.declareManagedMemoryUseCaseAtOperatorScope(
+ useCase.getKey(),
+ useCase.getValue()
+ + chainedTransformation
+ .getManagedMemoryOperatorScopeUseCaseWeights()
+ .getOrDefault(useCase.getKey(), 0));
+ }
+
+ chainedTransformation.setBufferTimeout(
+ Math.min(upTransform.getBufferTimeout(), downTransform.getBufferTimeout()));
+
+ if (upTransform.getMaxParallelism() > 0) {
+ chainedTransformation.setMaxParallelism(upTransform.getMaxParallelism());
+ }
+ chainedTransformation.setChainingStrategy(
+ getOperatorFactory(upTransform).getChainingStrategy());
+ chainedTransformation.setCoLocationGroupKey(upTransform.getCoLocationGroupKey());
+ chainedTransformation.setResources(
+ upTransform.getMinResources().merge(downTransform.getMinResources()),
+ upTransform.getPreferredResources().merge(downTransform.getPreferredResources()));
+
+ if (upTransform.getSlotSharingGroup().isPresent()) {
+ chainedTransformation.setSlotSharingGroup(upTransform.getSlotSharingGroup().get());
+ }
+ return chainedTransformation;
+ }
+
+ private boolean isChainable(Transformation<?> upTransform, Transformation<?> downTransform) {
+ return upTransform.getParallelism() == downTransform.getParallelism()
+ && upTransform.getMaxParallelism() == downTransform.getMaxParallelism()
+ && upTransform.getSlotSharingGroup().equals(downTransform.getSlotSharingGroup())
+ && areOperatorsChainable(upTransform, downTransform)
+ && isChainingEnabled;
+ }
+
+ // ----------------------- Utility Methods -----------------------
+
+ public static void apply(StreamExecutionEnvironment env) throws Exception {
+ Field transformationsField =
+ StreamExecutionEnvironment.class.getDeclaredField("transformations");
+ transformationsField.setAccessible(true);
+ List<Transformation<?>> transformations =
+ (List<Transformation<?>>) transformationsField.get(env);
+ Configuration config = PythonConfigUtil.getEnvironmentConfig(env);
+ PythonOperatorChainingOptimizer chainingOptimizer =
+ new PythonOperatorChainingOptimizer(
+ transformations,
+ config.getBoolean(PythonOptions.PYTHON_OPERATOR_CHAINING_ENABLED));
+ transformationsField.set(env, chainingOptimizer.optimize());
+ }
+
+ private static boolean areOperatorsChainable(
+ Transformation<?> upTransform, Transformation<?> downTransform) {
+ if (!areOperatorsChainableByChainingStrategy(upTransform, downTransform)) {
+ return false;
+ }
+
+ AbstractDataStreamPythonFunctionOperator<?> upOperator =
+ (AbstractDataStreamPythonFunctionOperator<?>)
+ ((SimpleOperatorFactory<?>) getOperatorFactory(upTransform)).getOperator();
+ AbstractDataStreamPythonFunctionOperator<?> downOperator =
+ (AbstractDataStreamPythonFunctionOperator<?>)
+ ((SimpleOperatorFactory<?>) getOperatorFactory(downTransform))
+ .getOperator();
+
+ return downOperator instanceof PythonProcessOperator
+ && (upOperator instanceof PythonKeyedProcessOperator
+ || upOperator instanceof PythonKeyedCoProcessOperator
+ || upOperator instanceof PythonProcessOperator
+ || upOperator instanceof PythonCoProcessOperator);
+ }
+
+ private static boolean areOperatorsChainableByChainingStrategy(
+ Transformation<?> upTransform, Transformation<?> downTransform) {
+ // we use switch/case here to make sure this is exhaustive if ever values are added to the
+ // ChainingStrategy enum
+ boolean isChainable;
+
+ StreamOperatorFactory<?> upStreamOperator = getOperatorFactory(upTransform);
+ StreamOperatorFactory<?> downStreamOperator = getOperatorFactory(downTransform);
+
+ switch (upStreamOperator.getChainingStrategy()) {
+ case NEVER:
+ isChainable = false;
+ break;
+ case ALWAYS:
+ case HEAD:
+ case HEAD_WITH_SOURCES:
+ isChainable = true;
+ break;
+ default:
+ throw new RuntimeException(
+ "Unknown chaining strategy: " + upStreamOperator.getChainingStrategy());
+ }
+
+ switch (downStreamOperator.getChainingStrategy()) {
+ case NEVER:
+ case HEAD:
+ isChainable = false;
+ break;
+ case ALWAYS:
+ // keep the value from upstream
+ break;
+ case HEAD_WITH_SOURCES:
+ // only if upstream is a source
+ isChainable &= (upStreamOperator instanceof SourceOperatorFactory);
+ break;
+ default:
+ throw new RuntimeException(
+ "Unknown chaining strategy: " + upStreamOperator.getChainingStrategy());
+ }
+
+ return isChainable;
+ }
+
+ private static StreamOperatorFactory<?> getOperatorFactory(Transformation<?> transform) {
+ if (transform instanceof OneInputTransformation) {
+ return ((OneInputTransformation<?, ?>) transform).getOperatorFactory();
+ } else if (transform instanceof TwoInputTransformation) {
+ return ((TwoInputTransformation<?, ?, ?>) transform).getOperatorFactory();
+ } else if (transform instanceof AbstractMultipleInputTransformation) {
+ return ((AbstractMultipleInputTransformation<?>) transform).getOperatorFactory();
+ } else {
+ return null;
+ }
+ }
+
+ private static void replaceInput(
+ Transformation<?> transformation,
+ Transformation<?> oldInput,
+ Transformation<?> newInput)
+ throws Exception {
+ if (transformation instanceof OneInputTransformation
+ || transformation instanceof FeedbackTransformation
+ || transformation instanceof SideOutputTransformation
+ || transformation instanceof ReduceTransformation
+ || transformation instanceof SinkTransformation
+ || transformation instanceof LegacySinkTransformation
+ || transformation instanceof TimestampsAndWatermarksTransformation) {
+ Field inputField = transformation.getClass().getDeclaredField("input");
+ inputField.setAccessible(true);
+ inputField.set(transformation, newInput);
+ } else if (transformation instanceof TwoInputTransformation) {
+ Field inputField;
+ if (((TwoInputTransformation<?, ?, ?>) transformation).getInput1() == oldInput) {
+ inputField = transformation.getClass().getDeclaredField("input1");
+ } else {
+ inputField = transformation.getClass().getDeclaredField("input2");
+ }
+ inputField.setAccessible(true);
+ inputField.set(transformation, newInput);
+ } else if (transformation instanceof UnionTransformation
+ || transformation instanceof AbstractMultipleInputTransformation) {
+ Field inputsField = transformation.getClass().getDeclaredField("inputs");
+ inputsField.setAccessible(true);
+ List<Transformation<?>> newInputs = Lists.newArrayList();
+ newInputs.addAll(transformation.getInputs());
+ newInputs.remove(oldInput);
+ newInputs.add(newInput);
+ inputsField.set(transformation, newInputs);
+ } else if (transformation instanceof AbstractBroadcastStateTransformation) {
+ Field inputField;
+ if (((AbstractBroadcastStateTransformation<?, ?, ?>) transformation).getRegularInput()
+ == oldInput) {
+ inputField = transformation.getClass().getDeclaredField("regularInput");
+ } else {
+ inputField = transformation.getClass().getDeclaredField("broadcastInput");
+ }
+ inputField.setAccessible(true);
+ inputField.set(transformation, newInput);
+ } else {
+ throw new RuntimeException("Unsupported transformation: " + transformation);
+ }
+ }
+
+ // ----------------------- Utility Classes -----------------------
+
+ private static class ChainInfo {
+ /** The transformation which represents the chaining of the {@link #oldTransformations}. */
+ public final Transformation<?> newTransformation;
+
+ /** The transformations which will be chained together. */
+ public final Collection<Transformation<?>> oldTransformations;
+
+ private ChainInfo(
+ Transformation<?> newTransformation,
+ Collection<Transformation<?>> oldTransformations) {
+ this.newTransformation = newTransformation;
+ this.oldTransformations = oldTransformations;
+ }
+
+ /** No chaining happens. */
+ public static ChainInfo of(Transformation<?> newTransformation) {
+ return new ChainInfo(newTransformation, Collections.emptyList());
+ }
+
+ public static ChainInfo of(
+ Transformation<?> newTransformation,
+ Collection<Transformation<?>> oldTransformations) {
+ return new ChainInfo(newTransformation, oldTransformations);
+ }
+ }
+}
diff --git a/flink-python/src/main/java/org/apache/flink/python/util/PythonConfigUtil.java b/flink-python/src/main/java/org/apache/flink/python/util/PythonConfigUtil.java
index 992dc38..2ef5600 100644
--- a/flink-python/src/main/java/org/apache/flink/python/util/PythonConfigUtil.java
+++ b/flink-python/src/main/java/org/apache/flink/python/util/PythonConfigUtil.java
@@ -28,8 +28,9 @@ import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.graph.StreamGraph;
import org.apache.flink.streaming.api.operators.SimpleOperatorFactory;
import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
+import org.apache.flink.streaming.api.operators.python.AbstractDataStreamPythonFunctionOperator;
+import org.apache.flink.streaming.api.operators.python.AbstractOneInputPythonFunctionOperator;
import org.apache.flink.streaming.api.operators.python.AbstractPythonFunctionOperator;
-import org.apache.flink.streaming.api.operators.python.OneInputPythonFunctionOperator;
import org.apache.flink.streaming.api.operators.python.PythonProcessOperator;
import org.apache.flink.streaming.api.transformations.AbstractMultipleInputTransformation;
import org.apache.flink.streaming.api.transformations.OneInputTransformation;
@@ -92,9 +93,9 @@ public class PythonConfigUtil {
/**
* Generate a {@link StreamGraph} for transformations maintained by current {@link
* StreamExecutionEnvironment}, and reset the merged env configurations with dependencies to
- * every {@link OneInputPythonFunctionOperator}. It is an idempotent operation that can be call
- * multiple times. Remember that only when need to execute the StreamGraph can we set the
- * clearTransformations to be True.
+ * every {@link AbstractOneInputPythonFunctionOperator}. It is an idempotent operation that can
+ * be call multiple times. Remember that only when need to execute the StreamGraph can we set
+ * the clearTransformations to be True.
*/
public static StreamGraph generateStreamGraphWithDependencies(
StreamExecutionEnvironment env, boolean clearTransformations)
@@ -169,7 +170,7 @@ public class PythonConfigUtil {
}
/**
- * Configure the {@link OneInputPythonFunctionOperator} to be chained with the
+ * Configure the {@link AbstractOneInputPythonFunctionOperator} to be chained with the
* upstream/downstream operator by setting their parallelism, slot sharing group, co-location
* group to be the same, and applying a {@link ForwardPartitioner}. 1. operator with name
* "_keyed_stream_values_operator" should align with its downstream operator. 2. operator with
@@ -234,7 +235,7 @@ public class PythonConfigUtil {
return null;
}
- private static boolean isPythonOperator(Transformation<?> transform) {
+ public static boolean isPythonOperator(Transformation<?> transform) {
if (transform instanceof OneInputTransformation) {
return isPythonOperator(
((OneInputTransformation<?, ?>) transform).getOperatorFactory());
@@ -258,6 +259,28 @@ public class PythonConfigUtil {
}
}
+ public static boolean isPythonDataStreamOperator(Transformation<?> transform) {
+ if (transform instanceof OneInputTransformation) {
+ return isPythonDataStreamOperator(
+ ((OneInputTransformation<?, ?>) transform).getOperatorFactory());
+ } else if (transform instanceof TwoInputTransformation) {
+ return isPythonDataStreamOperator(
+ ((TwoInputTransformation<?, ?, ?>) transform).getOperatorFactory());
+ } else {
+ return false;
+ }
+ }
+
+ private static boolean isPythonDataStreamOperator(
+ StreamOperatorFactory<?> streamOperatorFactory) {
+ if (streamOperatorFactory instanceof SimpleOperatorFactory) {
+ return ((SimpleOperatorFactory<?>) streamOperatorFactory).getOperator()
+ instanceof AbstractDataStreamPythonFunctionOperator;
+ } else {
+ return false;
+ }
+ }
+
/**
* Generator a new {@link PythonConfig} with the combined config which is derived from
* oldConfig.
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/functions/python/DataStreamPythonFunctionInfo.java b/flink-python/src/main/java/org/apache/flink/streaming/api/functions/python/DataStreamPythonFunctionInfo.java
index b3d5f1b..9aa1a05 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/functions/python/DataStreamPythonFunctionInfo.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/functions/python/DataStreamPythonFunctionInfo.java
@@ -19,24 +19,26 @@ package org.apache.flink.streaming.api.functions.python;
import org.apache.flink.annotation.Internal;
import org.apache.flink.table.functions.python.PythonFunction;
-
-import java.io.Serializable;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
/** {@link DataStreamPythonFunctionInfo} holds a PythonFunction and its function type. */
@Internal
-public class DataStreamPythonFunctionInfo implements Serializable {
- private static final long serialVersionUID = 1L;
+public class DataStreamPythonFunctionInfo extends PythonFunctionInfo {
+ private static final long serialVersionUID = 2L;
+
+ private static final Object[] EMPTY = new Object[0];
- private final PythonFunction pythonFunction;
private final int functionType;
public DataStreamPythonFunctionInfo(PythonFunction pythonFunction, int functionType) {
- this.pythonFunction = pythonFunction;
+ super(pythonFunction, EMPTY);
this.functionType = functionType;
}
- public PythonFunction getPythonFunction() {
- return pythonFunction;
+ public DataStreamPythonFunctionInfo(
+ PythonFunction pythonFunction, DataStreamPythonFunctionInfo input, int functionType) {
+ super(pythonFunction, new DataStreamPythonFunctionInfo[] {input});
+ this.functionType = functionType;
}
public int getFunctionType() {
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractDataStreamPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractDataStreamPythonFunctionOperator.java
new file mode 100644
index 0000000..25619b5
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractDataStreamPythonFunctionOperator.java
@@ -0,0 +1,71 @@
+/*
+ * 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.
+ */
+
+package org.apache.flink.streaming.api.operators.python;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
+import org.apache.flink.table.functions.python.PythonEnv;
+import org.apache.flink.util.Preconditions;
+
+/** Base class for all Python DataStream operators. */
+@Internal
+public abstract class AbstractDataStreamPythonFunctionOperator<OUT>
+ extends AbstractPythonFunctionOperator<OUT> implements ResultTypeQueryable<OUT> {
+
+ private static final long serialVersionUID = 1L;
+
+ /** The serialized python function to be executed. */
+ private final DataStreamPythonFunctionInfo pythonFunctionInfo;
+
+ /** The TypeInformation of output data. */
+ private final TypeInformation<OUT> outputTypeInfo;
+
+ public AbstractDataStreamPythonFunctionOperator(
+ Configuration config,
+ DataStreamPythonFunctionInfo pythonFunctionInfo,
+ TypeInformation<OUT> outputTypeInfo) {
+ super(config);
+ this.pythonFunctionInfo = Preconditions.checkNotNull(pythonFunctionInfo);
+ this.outputTypeInfo = Preconditions.checkNotNull(outputTypeInfo);
+ }
+
+ @Override
+ public PythonEnv getPythonEnv() {
+ return pythonFunctionInfo.getPythonFunction().getPythonEnv();
+ }
+
+ @Override
+ public TypeInformation<OUT> getProducedType() {
+ return outputTypeInfo;
+ }
+
+ public abstract <T> AbstractDataStreamPythonFunctionOperator<T> copy(
+ DataStreamPythonFunctionInfo pythonFunctionInfo, TypeInformation<T> outputTypeInfo);
+
+ // ----------------------------------------------------------------------
+ // Getters
+ // ----------------------------------------------------------------------
+
+ public DataStreamPythonFunctionInfo getPythonFunctionInfo() {
+ return pythonFunctionInfo;
+ }
+}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractOneInputPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractOneInputPythonFunctionOperator.java
index 5c8a35a..f27983b 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractOneInputPythonFunctionOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractOneInputPythonFunctionOperator.java
@@ -19,24 +19,180 @@
package org.apache.flink.streaming.api.operators.python;
import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeinfo.Types;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.fnexecution.v1.FlinkFnApi;
+import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
import org.apache.flink.streaming.api.operators.BoundedOneInput;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.operators.TimestampedCollector;
+import org.apache.flink.streaming.api.operators.python.collector.RunnerOutputCollector;
+import org.apache.flink.streaming.api.utils.PythonTypeUtils;
+import org.apache.flink.types.Row;
+import org.apache.flink.util.Preconditions;
-/** Base class for all one input stream operators to execute Python functions. */
+import java.util.Collections;
+import java.util.Map;
+
+import static org.apache.flink.streaming.api.utils.ProtoUtils.createRawTypeCoderInfoDescriptorProto;
+
+/**
+ * {@link AbstractOneInputPythonFunctionOperator} is responsible for launching beam runner which
+ * will start a python harness to execute user defined python function.
+ *
+ * <p>The operator will buffer the timestamp of input elements in a queue, and set into the produced
+ * output element.
+ */
@Internal
public abstract class AbstractOneInputPythonFunctionOperator<IN, OUT>
- extends AbstractPythonFunctionOperator<OUT>
+ extends AbstractDataStreamPythonFunctionOperator<OUT>
implements OneInputStreamOperator<IN, OUT>, BoundedOneInput {
private static final long serialVersionUID = 1L;
- public AbstractOneInputPythonFunctionOperator(Configuration config) {
- super(config);
+ /** The options used to configure the Python worker process. */
+ private final Map<String, String> jobOptions;
+
+ /** The TypeInformation of input data. */
+ private final TypeInformation<IN> inputTypeInfo;
+
+ /** The TypeInformation of runner input data. */
+ private transient TypeInformation<Row> runnerInputTypeInfo;
+
+ /** The TypeInformation of runner output data. */
+ private transient TypeInformation<Row> runnerOutputTypeInfo;
+
+ /** The TypeSerializer of python worker input data. */
+ private transient TypeSerializer<Row> runnerInputTypeSerializer;
+
+ /** The TypeSerializer of python worker output data. */
+ private transient TypeSerializer<Row> runnerOutputTypeSerializer;
+
+ /** Reusable InputStream used to holding the execution results to be deserialized. */
+ private transient ByteArrayInputStreamWithPos bais;
+
+ /** InputStream Wrapper. */
+ private transient DataInputViewStreamWrapper baisWrapper;
+
+ /** Reusable OutputStream used to holding the serialized input elements. */
+ protected transient ByteArrayOutputStreamWithPos baos;
+
+ /** OutputStream Wrapper. */
+ protected transient DataOutputViewStreamWrapper baosWrapper;
+
+ private transient RunnerInputHandler runnerInputHandler;
+
+ private transient RunnerOutputCollector<OUT> runnerOutputCollector;
+
+ public AbstractOneInputPythonFunctionOperator(
+ Configuration config,
+ DataStreamPythonFunctionInfo pythonFunctionInfo,
+ TypeInformation<IN> inputTypeInfo,
+ TypeInformation<OUT> outputTypeInfo) {
+ super(config, pythonFunctionInfo, outputTypeInfo);
+ this.jobOptions = config.toMap();
+ this.inputTypeInfo = Preconditions.checkNotNull(inputTypeInfo);
+ }
+
+ @Override
+ public void open() throws Exception {
+ bais = new ByteArrayInputStreamWithPos();
+ baisWrapper = new DataInputViewStreamWrapper(bais);
+ baos = new ByteArrayOutputStreamWithPos();
+ baosWrapper = new DataOutputViewStreamWrapper(baos);
+
+ runnerInputTypeInfo = RunnerInputHandler.getRunnerInputTypeInfo(inputTypeInfo);
+ runnerOutputTypeInfo = RunnerOutputCollector.getRunnerOutputTypeInfo(getProducedType());
+ runnerInputTypeSerializer =
+ PythonTypeUtils.TypeInfoToSerializerConverter.typeInfoSerializerConverter(
+ runnerInputTypeInfo);
+ runnerOutputTypeSerializer =
+ PythonTypeUtils.TypeInfoToSerializerConverter.typeInfoSerializerConverter(
+ runnerOutputTypeInfo);
+
+ runnerInputHandler = new RunnerInputHandler();
+ runnerOutputCollector = new RunnerOutputCollector<>(new TimestampedCollector<>(output));
+
+ super.open();
}
@Override
public void endInput() throws Exception {
invokeFinishBundle();
}
+
+ @Override
+ public void emitResult(Tuple2<byte[], Integer> resultTuple) throws Exception {
+ byte[] rawResult = resultTuple.f0;
+ int length = resultTuple.f1;
+ bais.setBuffer(rawResult, 0, length);
+ Row runnerOutput = runnerOutputTypeSerializer.deserialize(baisWrapper);
+ runnerOutputCollector.collect(runnerOutput);
+ }
+
+ public void processElement(long timestamp, long watermark, Object element) throws Exception {
+ Row row = runnerInputHandler.buildRunnerInputData(timestamp, watermark, element);
+ runnerInputTypeSerializer.serialize(row, baosWrapper);
+ pythonFunctionRunner.process(baos.toByteArray());
+ baos.reset();
+ elementCount++;
+ checkInvokeFinishBundleByCount();
+ emitResults();
+ }
+
+ public Map<String, String> getInternalParameters() {
+ return Collections.emptyMap();
+ }
+
+ public FlinkFnApi.CoderInfoDescriptor createInputCoderInfoDescriptor() {
+ return createRawTypeCoderInfoDescriptorProto(
+ runnerInputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
+ }
+
+ public FlinkFnApi.CoderInfoDescriptor createOutputCoderInfoDescriptor() {
+ return createRawTypeCoderInfoDescriptorProto(
+ runnerOutputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
+ }
+
+ // ----------------------------------------------------------------------
+ // Getters
+ // ----------------------------------------------------------------------
+
+ protected Map<String, String> getJobOptions() {
+ return jobOptions;
+ }
+
+ public TypeInformation<IN> getInputTypeInfo() {
+ return inputTypeInfo;
+ }
+
+ private static final class RunnerInputHandler {
+
+ /** Reusable row for element data. */
+ private final Row reusableRunnerInput;
+
+ public RunnerInputHandler() {
+ this.reusableRunnerInput = new Row(3);
+ }
+
+ public Row buildRunnerInputData(long timestamp, long watermark, Object elementData) {
+ reusableRunnerInput.setField(0, timestamp);
+ reusableRunnerInput.setField(1, watermark);
+ reusableRunnerInput.setField(2, elementData);
+ return reusableRunnerInput;
+ }
+
+ public static TypeInformation<Row> getRunnerInputTypeInfo(
+ TypeInformation<?> elementDataTypeInfo) {
+ // structure: [timestamp, watermark, data]
+ return Types.ROW(Types.LONG, Types.LONG, elementDataTypeInfo);
+ }
+ }
}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractTwoInputPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractTwoInputPythonFunctionOperator.java
index d13e380..5bf078d 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractTwoInputPythonFunctionOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractTwoInputPythonFunctionOperator.java
@@ -19,24 +19,189 @@
package org.apache.flink.streaming.api.operators.python;
import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeinfo.Types;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.fnexecution.v1.FlinkFnApi;
+import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
import org.apache.flink.streaming.api.operators.BoundedMultiInput;
+import org.apache.flink.streaming.api.operators.TimestampedCollector;
import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
+import org.apache.flink.streaming.api.operators.python.collector.RunnerOutputCollector;
+import org.apache.flink.types.Row;
+import org.apache.flink.util.Preconditions;
-/** Base class for all two input stream operators to execute Python functions. */
+import java.util.Map;
+
+import static org.apache.flink.streaming.api.utils.ProtoUtils.createRawTypeCoderInfoDescriptorProto;
+import static org.apache.flink.streaming.api.utils.PythonTypeUtils.TypeInfoToSerializerConverter.typeInfoSerializerConverter;
+
+/**
+ * {@link AbstractTwoInputPythonFunctionOperator} is responsible for launching beam runner which
+ * will start a python harness to execute two-input user defined python function.
+ */
@Internal
public abstract class AbstractTwoInputPythonFunctionOperator<IN1, IN2, OUT>
- extends AbstractPythonFunctionOperator<OUT>
+ extends AbstractDataStreamPythonFunctionOperator<OUT>
implements TwoInputStreamOperator<IN1, IN2, OUT>, BoundedMultiInput {
private static final long serialVersionUID = 1L;
- public AbstractTwoInputPythonFunctionOperator(Configuration config) {
- super(config);
+ /** The options used to configure the Python worker process. */
+ private final Map<String, String> jobOptions;
+
+ /** The left input type. */
+ private final TypeInformation<IN1> inputTypeInfo1;
+
+ /** The right input type. */
+ private final TypeInformation<IN2> inputTypeInfo2;
+
+ /** The TypeInformation of python worker input data. */
+ private transient TypeInformation<Row> runnerInputTypeInfo;
+
+ private transient TypeInformation<Row> runnerOutputTypeInfo;
+
+ /** The TypeSerializer of python worker input data. */
+ private transient TypeSerializer<Row> runnerInputTypeSerializer;
+
+ /** The TypeSerializer of the runner output. */
+ private transient TypeSerializer<Row> runnerOutputTypeSerializer;
+
+ private transient ByteArrayInputStreamWithPos bais;
+ private transient DataInputViewStreamWrapper baisWrapper;
+ protected transient ByteArrayOutputStreamWithPos baos;
+ protected transient DataOutputViewStreamWrapper baosWrapper;
+
+ private transient RunnerInputHandler runnerInputHandler;
+ private transient RunnerOutputCollector<OUT> runnerOutputCollector;
+
+ public AbstractTwoInputPythonFunctionOperator(
+ Configuration config,
+ DataStreamPythonFunctionInfo pythonFunctionInfo,
+ TypeInformation<IN1> inputTypeInfo1,
+ TypeInformation<IN2> inputTypeInfo2,
+ TypeInformation<OUT> outputTypeInfo) {
+ super(config, pythonFunctionInfo, outputTypeInfo);
+ this.jobOptions = config.toMap();
+ this.inputTypeInfo1 = Preconditions.checkNotNull(inputTypeInfo1);
+ this.inputTypeInfo2 = Preconditions.checkNotNull(inputTypeInfo2);
+ }
+
+ @Override
+ public void open() throws Exception {
+ bais = new ByteArrayInputStreamWithPos();
+ baisWrapper = new DataInputViewStreamWrapper(bais);
+ baos = new ByteArrayOutputStreamWithPos();
+ baosWrapper = new DataOutputViewStreamWrapper(baos);
+
+ runnerInputTypeInfo =
+ RunnerInputHandler.getRunnerInputTypeInfo(inputTypeInfo1, inputTypeInfo2);
+ runnerOutputTypeInfo = RunnerOutputCollector.getRunnerOutputTypeInfo(getProducedType());
+ runnerInputTypeSerializer = typeInfoSerializerConverter(runnerInputTypeInfo);
+ runnerOutputTypeSerializer = typeInfoSerializerConverter(runnerOutputTypeInfo);
+
+ runnerInputHandler = new RunnerInputHandler();
+ runnerOutputCollector = new RunnerOutputCollector<>(new TimestampedCollector<>(output));
+
+ super.open();
}
@Override
public void endInput(int inputId) throws Exception {
invokeFinishBundle();
}
+
+ @Override
+ public void emitResult(Tuple2<byte[], Integer> resultTuple) throws Exception {
+ byte[] rawResult = resultTuple.f0;
+ int length = resultTuple.f1;
+ bais.setBuffer(rawResult, 0, length);
+ Row runnerOutput = runnerOutputTypeSerializer.deserialize(baisWrapper);
+ runnerOutputCollector.collect(runnerOutput);
+ }
+
+ public void processElement(boolean isLeft, long timestamp, long watermark, Object element)
+ throws Exception {
+ Row row = runnerInputHandler.buildRunnerInputData(isLeft, timestamp, watermark, element);
+ runnerInputTypeSerializer.serialize(row, baosWrapper);
+ pythonFunctionRunner.process(baos.toByteArray());
+ baos.reset();
+ elementCount++;
+ checkInvokeFinishBundleByCount();
+ emitResults();
+ }
+
+ public FlinkFnApi.CoderInfoDescriptor createInputCoderInfoDescriptor() {
+ return createRawTypeCoderInfoDescriptorProto(
+ runnerInputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
+ }
+
+ public FlinkFnApi.CoderInfoDescriptor createOutputCoderInfoDescriptor() {
+ return createRawTypeCoderInfoDescriptorProto(
+ runnerOutputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
+ }
+
+ // ----------------------------------------------------------------------
+ // Getters
+ // ----------------------------------------------------------------------
+
+ protected Map<String, String> getJobOptions() {
+ return jobOptions;
+ }
+
+ protected TypeInformation<IN1> getLeftInputType() {
+ return inputTypeInfo1;
+ }
+
+ protected TypeInformation<IN2> getRightInputType() {
+ return inputTypeInfo2;
+ }
+
+ private static final class RunnerInputHandler {
+
+ private final Row reusableElementData;
+ private final Row reusableRunnerInput;
+
+ public RunnerInputHandler() {
+ this.reusableElementData = new Row(3);
+ this.reusableRunnerInput = new Row(3);
+ this.reusableRunnerInput.setField(2, reusableElementData);
+ }
+
+ public Row buildRunnerInputData(
+ boolean isLeft, long timestamp, long watermark, Object elementData) {
+ reusableElementData.setField(0, isLeft);
+ if (isLeft) {
+ // The input row is a tuple of key and value.
+ reusableElementData.setField(1, elementData);
+ // need to set null since it is a reuse row.
+ reusableElementData.setField(2, null);
+ } else {
+ // need to set null since it is a reuse row.
+ reusableElementData.setField(1, null);
+ // The input row is a tuple of key and value.
+ reusableElementData.setField(2, elementData);
+ }
+
+ reusableRunnerInput.setField(0, timestamp);
+ reusableRunnerInput.setField(1, watermark);
+ return reusableRunnerInput;
+ }
+
+ public static TypeInformation<Row> getRunnerInputTypeInfo(
+ TypeInformation<?> leftInputType, TypeInformation<?> rightInputType) {
+ // structure: [timestamp, watermark, [isLeft, leftInput, rightInput]]
+ return Types.ROW(
+ Types.LONG,
+ Types.LONG,
+ new RowTypeInfo(Types.BOOLEAN, leftInputType, rightInputType));
+ }
+ }
}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/OneInputPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/OneInputPythonFunctionOperator.java
deleted file mode 100644
index a315683e..0000000
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/OneInputPythonFunctionOperator.java
+++ /dev/null
@@ -1,214 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.flink.streaming.api.operators.python;
-
-import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.common.typeinfo.TypeInformation;
-import org.apache.flink.api.common.typeinfo.Types;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
-import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
-import org.apache.flink.fnexecution.v1.FlinkFnApi;
-import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
-import org.apache.flink.streaming.api.operators.TimestampedCollector;
-import org.apache.flink.streaming.api.operators.python.collector.RunnerOutputCollector;
-import org.apache.flink.streaming.api.utils.PythonTypeUtils;
-import org.apache.flink.table.functions.python.PythonEnv;
-import org.apache.flink.types.Row;
-import org.apache.flink.util.Preconditions;
-
-import java.util.Collections;
-import java.util.Map;
-
-import static org.apache.flink.streaming.api.utils.ProtoUtils.createRawTypeCoderInfoDescriptorProto;
-
-/**
- * {@link OneInputPythonFunctionOperator} is responsible for launching beam runner which will start
- * a python harness to execute user defined python function.
- *
- * <p>The operator will buffer the timestamp of input elements in a queue, and set into the produced
- * output element.
- */
-@Internal
-public abstract class OneInputPythonFunctionOperator<IN, OUT>
- extends AbstractOneInputPythonFunctionOperator<IN, OUT>
- implements ResultTypeQueryable<OUT> {
-
- private static final long serialVersionUID = 1L;
-
- /** The options used to configure the Python worker process. */
- private final Map<String, String> jobOptions;
-
- /** The serialized python function to be executed. */
- private final DataStreamPythonFunctionInfo pythonFunctionInfo;
-
- /** The TypeInformation of input data. */
- private final TypeInformation<IN> inputTypeInfo;
-
- /** The TypeInformation of output data. */
- private final TypeInformation<OUT> outputTypeInfo;
-
- /** The TypeInformation of runner input data. */
- private transient TypeInformation<Row> runnerInputTypeInfo;
-
- /** The TypeInformation of runner output data. */
- private transient TypeInformation<Row> runnerOutputTypeInfo;
-
- /** The TypeSerializer of python worker input data. */
- private transient TypeSerializer<Row> runnerInputTypeSerializer;
-
- /** The TypeSerializer of python worker output data. */
- private transient TypeSerializer<Row> runnerOutputTypeSerializer;
-
- /** Reusable InputStream used to holding the execution results to be deserialized. */
- private transient ByteArrayInputStreamWithPos bais;
-
- /** InputStream Wrapper. */
- private transient DataInputViewStreamWrapper baisWrapper;
-
- /** Reusable OutputStream used to holding the serialized input elements. */
- protected transient ByteArrayOutputStreamWithPos baos;
-
- /** OutputStream Wrapper. */
- protected transient DataOutputViewStreamWrapper baosWrapper;
-
- private transient RunnerInputHandler runnerInputHandler;
-
- private transient RunnerOutputCollector<OUT> runnerOutputCollector;
-
- public OneInputPythonFunctionOperator(
- Configuration config,
- DataStreamPythonFunctionInfo pythonFunctionInfo,
- TypeInformation<IN> inputTypeInfo,
- TypeInformation<OUT> outputTypeInfo) {
- super(config);
- this.jobOptions = config.toMap();
- this.pythonFunctionInfo = Preconditions.checkNotNull(pythonFunctionInfo);
- this.inputTypeInfo = Preconditions.checkNotNull(inputTypeInfo);
- this.outputTypeInfo = Preconditions.checkNotNull(outputTypeInfo);
- }
-
- @Override
- public void open() throws Exception {
- bais = new ByteArrayInputStreamWithPos();
- baisWrapper = new DataInputViewStreamWrapper(bais);
- baos = new ByteArrayOutputStreamWithPos();
- baosWrapper = new DataOutputViewStreamWrapper(baos);
-
- runnerInputTypeInfo = RunnerInputHandler.getRunnerInputTypeInfo(inputTypeInfo);
- runnerOutputTypeInfo = RunnerOutputCollector.getRunnerOutputTypeInfo(outputTypeInfo);
- runnerInputTypeSerializer =
- PythonTypeUtils.TypeInfoToSerializerConverter.typeInfoSerializerConverter(
- runnerInputTypeInfo);
- runnerOutputTypeSerializer =
- PythonTypeUtils.TypeInfoToSerializerConverter.typeInfoSerializerConverter(
- runnerOutputTypeInfo);
-
- runnerInputHandler = new RunnerInputHandler();
- runnerOutputCollector = new RunnerOutputCollector<>(new TimestampedCollector<>(output));
-
- super.open();
- }
-
- @Override
- public PythonEnv getPythonEnv() {
- return pythonFunctionInfo.getPythonFunction().getPythonEnv();
- }
-
- @Override
- public TypeInformation<OUT> getProducedType() {
- return this.outputTypeInfo;
- }
-
- @Override
- public void emitResult(Tuple2<byte[], Integer> resultTuple) throws Exception {
- byte[] rawResult = resultTuple.f0;
- int length = resultTuple.f1;
- bais.setBuffer(rawResult, 0, length);
- Row runnerOutput = runnerOutputTypeSerializer.deserialize(baisWrapper);
- runnerOutputCollector.collect(runnerOutput);
- }
-
- public void processElement(long timestamp, long watermark, Object element) throws Exception {
- Row row = runnerInputHandler.buildRunnerInputData(timestamp, watermark, element);
- runnerInputTypeSerializer.serialize(row, baosWrapper);
- pythonFunctionRunner.process(baos.toByteArray());
- baos.reset();
- elementCount++;
- checkInvokeFinishBundleByCount();
- emitResults();
- }
-
- public Map<String, String> getInternalParameters() {
- return Collections.emptyMap();
- }
-
- public FlinkFnApi.CoderInfoDescriptor createInputCoderInfoDescriptor() {
- return createRawTypeCoderInfoDescriptorProto(
- runnerInputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
- }
-
- public FlinkFnApi.CoderInfoDescriptor createOutputCoderInfoDescriptor() {
- return createRawTypeCoderInfoDescriptorProto(
- runnerOutputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
- }
-
- // ----------------------------------------------------------------------
- // Getters
- // ----------------------------------------------------------------------
-
- protected Map<String, String> getJobOptions() {
- return jobOptions;
- }
-
- protected DataStreamPythonFunctionInfo getPythonFunctionInfo() {
- return pythonFunctionInfo;
- }
-
- protected TypeInformation<IN> getInputTypeInfo() {
- return inputTypeInfo;
- }
-
- private static final class RunnerInputHandler {
-
- /** Reusable row for element data. */
- private final Row reusableRunnerInput;
-
- public RunnerInputHandler() {
- this.reusableRunnerInput = new Row(3);
- }
-
- public Row buildRunnerInputData(long timestamp, long watermark, Object elementData) {
- reusableRunnerInput.setField(0, timestamp);
- reusableRunnerInput.setField(1, watermark);
- reusableRunnerInput.setField(2, elementData);
- return reusableRunnerInput;
- }
-
- public static TypeInformation<Row> getRunnerInputTypeInfo(
- TypeInformation<?> elementDataTypeInfo) {
- // structure: [timestamp, watermark, data]
- return Types.ROW(Types.LONG, Types.LONG, elementDataTypeInfo);
- }
- }
-}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonCoProcessOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonCoProcessOperator.java
index c34d25d..fd83c91 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonCoProcessOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonCoProcessOperator.java
@@ -25,13 +25,13 @@ import org.apache.flink.python.PythonFunctionRunner;
import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
import org.apache.flink.streaming.api.operators.InternalTimerService;
import org.apache.flink.streaming.api.runners.python.beam.BeamDataStreamPythonFunctionRunner;
+import org.apache.flink.streaming.api.utils.ProtoUtils;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import java.util.Collections;
import static org.apache.flink.python.Constants.STATELESS_FUNCTION_URN;
-import static org.apache.flink.streaming.api.utils.ProtoUtils.getUserDefinedDataStreamFunctionProto;
import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchExecutionMode;
/**
@@ -43,7 +43,7 @@ import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchEx
*/
@Internal
public class PythonCoProcessOperator<IN1, IN2, OUT>
- extends TwoInputPythonFunctionOperator<IN1, IN2, OUT> {
+ extends AbstractTwoInputPythonFunctionOperator<IN1, IN2, OUT> {
private static final long serialVersionUID = 1L;
@@ -71,7 +71,7 @@ public class PythonCoProcessOperator<IN1, IN2, OUT>
getRuntimeContext().getTaskName(),
createPythonEnvironmentManager(),
STATELESS_FUNCTION_URN,
- getUserDefinedDataStreamFunctionProto(
+ ProtoUtils.createUserDefinedDataStreamFunctionProtos(
getPythonFunctionInfo(),
getRuntimeContext(),
Collections.emptyMap(),
@@ -114,4 +114,15 @@ public class PythonCoProcessOperator<IN1, IN2, OUT>
super.processWatermark(mark);
currentWatermark = mark.getTimestamp();
}
+
+ @Override
+ public <T> AbstractDataStreamPythonFunctionOperator<T> copy(
+ DataStreamPythonFunctionInfo pythonFunctionInfo, TypeInformation<T> outputTypeInfo) {
+ return new PythonCoProcessOperator<>(
+ getConfig().getConfig(),
+ pythonFunctionInfo,
+ getLeftInputType(),
+ getRightInputType(),
+ outputTypeInfo);
+ }
}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonKeyedCoProcessOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonKeyedCoProcessOperator.java
index f50bd56..cf8d9ca 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonKeyedCoProcessOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonKeyedCoProcessOperator.java
@@ -49,7 +49,8 @@ import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchEx
/** KeyedCoProcessOperator. */
@Internal
-public class PythonKeyedCoProcessOperator<OUT> extends TwoInputPythonFunctionOperator<Row, Row, OUT>
+public class PythonKeyedCoProcessOperator<OUT>
+ extends AbstractTwoInputPythonFunctionOperator<Row, Row, OUT>
implements Triggerable<Row, VoidNamespace> {
private static final long serialVersionUID = 1L;
@@ -108,7 +109,7 @@ public class PythonKeyedCoProcessOperator<OUT> extends TwoInputPythonFunctionOpe
getRuntimeContext().getTaskName(),
createPythonEnvironmentManager(),
STATEFUL_FUNCTION_URN,
- ProtoUtils.getUserDefinedDataStreamStatefulFunctionProto(
+ ProtoUtils.createUserDefinedDataStreamStatefulFunctionProtos(
getPythonFunctionInfo(),
getRuntimeContext(),
Collections.emptyMap(),
@@ -215,4 +216,15 @@ public class PythonKeyedCoProcessOperator<OUT> extends TwoInputPythonFunctionOpe
public Object getCurrentKey() {
return keyForTimerService;
}
+
+ @Override
+ public <T> AbstractDataStreamPythonFunctionOperator<T> copy(
+ DataStreamPythonFunctionInfo pythonFunctionInfo, TypeInformation<T> outputTypeInfo) {
+ return new PythonKeyedCoProcessOperator<>(
+ getConfig().getConfig(),
+ pythonFunctionInfo,
+ getLeftInputType(),
+ getRightInputType(),
+ outputTypeInfo);
+ }
}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonKeyedProcessOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonKeyedProcessOperator.java
index 73db676..1d58371 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonKeyedProcessOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonKeyedProcessOperator.java
@@ -53,7 +53,8 @@ import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchEx
* state request from the python stateful user defined function.
*/
@Internal
-public class PythonKeyedProcessOperator<OUT> extends OneInputPythonFunctionOperator<Row, OUT>
+public class PythonKeyedProcessOperator<OUT>
+ extends AbstractOneInputPythonFunctionOperator<Row, OUT>
implements Triggerable<Row, Object> {
private static final long serialVersionUID = 1L;
@@ -138,7 +139,7 @@ public class PythonKeyedProcessOperator<OUT> extends OneInputPythonFunctionOpera
getRuntimeContext().getTaskName(),
createPythonEnvironmentManager(),
STATEFUL_FUNCTION_URN,
- ProtoUtils.getUserDefinedDataStreamStatefulFunctionProto(
+ ProtoUtils.createUserDefinedDataStreamStatefulFunctionProtos(
getPythonFunctionInfo(),
getRuntimeContext(),
Collections.emptyMap(),
@@ -233,4 +234,15 @@ public class PythonKeyedProcessOperator<OUT> extends OneInputPythonFunctionOpera
public Object getCurrentKey() {
return keyForTimerService;
}
+
+ @Override
+ public <T> AbstractDataStreamPythonFunctionOperator<T> copy(
+ DataStreamPythonFunctionInfo pythonFunctionInfo, TypeInformation<T> outputTypeInfo) {
+ return new PythonKeyedProcessOperator<>(
+ getConfig().getConfig(),
+ pythonFunctionInfo,
+ (RowTypeInfo) getInputTypeInfo(),
+ outputTypeInfo,
+ namespaceSerializer);
+ }
}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonProcessOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonProcessOperator.java
index 9d38f9b..131927f 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonProcessOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/PythonProcessOperator.java
@@ -26,6 +26,7 @@ import org.apache.flink.python.PythonFunctionRunner;
import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
import org.apache.flink.streaming.api.operators.InternalTimerService;
import org.apache.flink.streaming.api.runners.python.beam.BeamDataStreamPythonFunctionRunner;
+import org.apache.flink.streaming.api.utils.ProtoUtils;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
@@ -35,7 +36,6 @@ import java.util.HashMap;
import java.util.Map;
import static org.apache.flink.python.Constants.STATELESS_FUNCTION_URN;
-import static org.apache.flink.streaming.api.utils.ProtoUtils.getUserDefinedDataStreamFunctionProto;
import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchExecutionMode;
/**
@@ -43,7 +43,8 @@ import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchEx
* harness to execute user defined python ProcessFunction.
*/
@Internal
-public class PythonProcessOperator<IN, OUT> extends OneInputPythonFunctionOperator<IN, OUT> {
+public class PythonProcessOperator<IN, OUT>
+ extends AbstractOneInputPythonFunctionOperator<IN, OUT> {
private static final long serialVersionUID = 1L;
@@ -74,7 +75,7 @@ public class PythonProcessOperator<IN, OUT> extends OneInputPythonFunctionOperat
getRuntimeContext().getTaskName(),
createPythonEnvironmentManager(),
STATELESS_FUNCTION_URN,
- getUserDefinedDataStreamFunctionProto(
+ ProtoUtils.createUserDefinedDataStreamFunctionProtos(
getPythonFunctionInfo(),
getRuntimeContext(),
getInternalParameters(),
@@ -126,4 +127,11 @@ public class PythonProcessOperator<IN, OUT> extends OneInputPythonFunctionOperat
public void setNumPartitions(int numPartitions) {
this.numPartitions = numPartitions;
}
+
+ @Override
+ public <T> AbstractDataStreamPythonFunctionOperator<T> copy(
+ DataStreamPythonFunctionInfo pythonFunctionInfo, TypeInformation<T> outputTypeInfo) {
+ return new PythonProcessOperator<>(
+ getConfig().getConfig(), pythonFunctionInfo, getInputTypeInfo(), outputTypeInfo);
+ }
}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/TwoInputPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/TwoInputPythonFunctionOperator.java
deleted file mode 100644
index f5f1a46..0000000
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/TwoInputPythonFunctionOperator.java
+++ /dev/null
@@ -1,219 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.flink.streaming.api.operators.python;
-
-import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.common.typeinfo.TypeInformation;
-import org.apache.flink.api.common.typeinfo.Types;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
-import org.apache.flink.api.java.typeutils.RowTypeInfo;
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
-import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
-import org.apache.flink.fnexecution.v1.FlinkFnApi;
-import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
-import org.apache.flink.streaming.api.operators.TimestampedCollector;
-import org.apache.flink.streaming.api.operators.python.collector.RunnerOutputCollector;
-import org.apache.flink.table.functions.python.PythonEnv;
-import org.apache.flink.types.Row;
-import org.apache.flink.util.Preconditions;
-
-import java.util.Map;
-
-import static org.apache.flink.streaming.api.utils.ProtoUtils.createRawTypeCoderInfoDescriptorProto;
-import static org.apache.flink.streaming.api.utils.PythonTypeUtils.TypeInfoToSerializerConverter.typeInfoSerializerConverter;
-
-/**
- * {@link TwoInputPythonFunctionOperator} is responsible for launching beam runner which will start
- * a python harness to execute two-input user defined python function.
- */
-@Internal
-public abstract class TwoInputPythonFunctionOperator<IN1, IN2, OUT>
- extends AbstractTwoInputPythonFunctionOperator<IN1, IN2, OUT>
- implements ResultTypeQueryable<OUT> {
-
- private static final long serialVersionUID = 1L;
-
- /** The options used to configure the Python worker process. */
- private final Map<String, String> jobOptions;
-
- /** The serialized python function to be executed. */
- private final DataStreamPythonFunctionInfo pythonFunctionInfo;
-
- /** The left input type. */
- private final TypeInformation<IN1> inputTypeInfo1;
-
- /** The right input type. */
- private final TypeInformation<IN2> inputTypeInfo2;
-
- /** The output type. */
- private final TypeInformation<OUT> outputTypeInfo;
-
- /** The TypeInformation of python worker input data. */
- private transient TypeInformation<Row> runnerInputTypeInfo;
-
- private transient TypeInformation<Row> runnerOutputTypeInfo;
-
- /** The TypeSerializer of python worker input data. */
- private transient TypeSerializer<Row> runnerInputTypeSerializer;
-
- /** The TypeSerializer of the runner output. */
- private transient TypeSerializer<Row> runnerOutputTypeSerializer;
-
- private transient ByteArrayInputStreamWithPos bais;
- private transient DataInputViewStreamWrapper baisWrapper;
- protected transient ByteArrayOutputStreamWithPos baos;
- protected transient DataOutputViewStreamWrapper baosWrapper;
-
- private transient RunnerInputHandler runnerInputHandler;
- private transient RunnerOutputCollector<OUT> runnerOutputCollector;
-
- public TwoInputPythonFunctionOperator(
- Configuration config,
- DataStreamPythonFunctionInfo pythonFunctionInfo,
- TypeInformation<IN1> inputTypeInfo1,
- TypeInformation<IN2> inputTypeInfo2,
- TypeInformation<OUT> outputTypeInfo) {
- super(config);
- this.jobOptions = config.toMap();
- this.pythonFunctionInfo = Preconditions.checkNotNull(pythonFunctionInfo);
- this.inputTypeInfo1 = Preconditions.checkNotNull(inputTypeInfo1);
- this.inputTypeInfo2 = Preconditions.checkNotNull(inputTypeInfo2);
- this.outputTypeInfo = Preconditions.checkNotNull(outputTypeInfo);
- }
-
- @Override
- public void open() throws Exception {
- bais = new ByteArrayInputStreamWithPos();
- baisWrapper = new DataInputViewStreamWrapper(bais);
- baos = new ByteArrayOutputStreamWithPos();
- baosWrapper = new DataOutputViewStreamWrapper(baos);
-
- runnerInputTypeInfo =
- RunnerInputHandler.getRunnerInputTypeInfo(inputTypeInfo1, inputTypeInfo2);
- runnerOutputTypeInfo = RunnerOutputCollector.getRunnerOutputTypeInfo(outputTypeInfo);
- runnerInputTypeSerializer = typeInfoSerializerConverter(runnerInputTypeInfo);
- runnerOutputTypeSerializer = typeInfoSerializerConverter(runnerOutputTypeInfo);
-
- runnerInputHandler = new RunnerInputHandler();
- runnerOutputCollector = new RunnerOutputCollector<>(new TimestampedCollector<>(output));
-
- super.open();
- }
-
- @Override
- public PythonEnv getPythonEnv() {
- return pythonFunctionInfo.getPythonFunction().getPythonEnv();
- }
-
- @Override
- public TypeInformation<OUT> getProducedType() {
- return outputTypeInfo;
- }
-
- @Override
- public void emitResult(Tuple2<byte[], Integer> resultTuple) throws Exception {
- byte[] rawResult = resultTuple.f0;
- int length = resultTuple.f1;
- bais.setBuffer(rawResult, 0, length);
- Row runnerOutput = runnerOutputTypeSerializer.deserialize(baisWrapper);
- runnerOutputCollector.collect(runnerOutput);
- }
-
- public void processElement(boolean isLeft, long timestamp, long watermark, Object element)
- throws Exception {
- Row row = runnerInputHandler.buildRunnerInputData(isLeft, timestamp, watermark, element);
- runnerInputTypeSerializer.serialize(row, baosWrapper);
- pythonFunctionRunner.process(baos.toByteArray());
- baos.reset();
- elementCount++;
- checkInvokeFinishBundleByCount();
- emitResults();
- }
-
- public FlinkFnApi.CoderInfoDescriptor createInputCoderInfoDescriptor() {
- return createRawTypeCoderInfoDescriptorProto(
- runnerInputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
- }
-
- public FlinkFnApi.CoderInfoDescriptor createOutputCoderInfoDescriptor() {
- return createRawTypeCoderInfoDescriptorProto(
- runnerOutputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
- }
-
- // ----------------------------------------------------------------------
- // Getters
- // ----------------------------------------------------------------------
-
- protected Map<String, String> getJobOptions() {
- return jobOptions;
- }
-
- protected DataStreamPythonFunctionInfo getPythonFunctionInfo() {
- return pythonFunctionInfo;
- }
-
- protected TypeInformation<IN1> getLeftInputType() {
- return inputTypeInfo1;
- }
-
- private static final class RunnerInputHandler {
-
- private final Row reusableElementData;
- private final Row reusableRunnerInput;
-
- public RunnerInputHandler() {
- this.reusableElementData = new Row(3);
- this.reusableRunnerInput = new Row(3);
- this.reusableRunnerInput.setField(2, reusableElementData);
- }
-
- public Row buildRunnerInputData(
- boolean isLeft, long timestamp, long watermark, Object elementData) {
- reusableElementData.setField(0, isLeft);
- if (isLeft) {
- // The input row is a tuple of key and value.
- reusableElementData.setField(1, elementData);
- // need to set null since it is a reuse row.
- reusableElementData.setField(2, null);
- } else {
- // need to set null since it is a reuse row.
- reusableElementData.setField(1, null);
- // The input row is a tuple of key and value.
- reusableElementData.setField(2, elementData);
- }
-
- reusableRunnerInput.setField(0, timestamp);
- reusableRunnerInput.setField(1, watermark);
- return reusableRunnerInput;
- }
-
- public static TypeInformation<Row> getRunnerInputTypeInfo(
- TypeInformation<?> leftInputType, TypeInformation<?> rightInputType) {
- // structure: [timestamp, watermark, [isLeft, leftInput, rightInput]]
- return Types.ROW(
- Types.LONG,
- Types.LONG,
- new RowTypeInfo(Types.BOOLEAN, leftInputType, rightInputType));
- }
- }
-}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java
index ed192b4..7347498 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java
@@ -44,9 +44,12 @@ import static org.apache.flink.python.Constants.INPUT_COLLECTION_ID;
import static org.apache.flink.python.Constants.MAIN_INPUT_NAME;
import static org.apache.flink.python.Constants.MAIN_OUTPUT_NAME;
import static org.apache.flink.python.Constants.OUTPUT_COLLECTION_ID;
+import static org.apache.flink.python.Constants.STATELESS_FUNCTION_URN;
import static org.apache.flink.python.Constants.TIMER_ID;
-import static org.apache.flink.python.Constants.TRANSFORM_ID;
+import static org.apache.flink.python.Constants.WINDOW_STRATEGY;
import static org.apache.flink.python.Constants.WRAPPER_TIMER_CODER_ID;
+import static org.apache.flink.streaming.api.utils.ProtoUtils.createCoderProto;
+import static org.apache.flink.streaming.api.utils.ProtoUtils.createReviseOutputDataStreamFunctionProto;
/**
* {@link BeamDataStreamPythonFunctionRunner} is responsible for starting a beam python harness to
@@ -55,21 +58,24 @@ import static org.apache.flink.python.Constants.WRAPPER_TIMER_CODER_ID;
@Internal
public class BeamDataStreamPythonFunctionRunner extends BeamPythonFunctionRunner {
- @Nullable private final FlinkFnApi.CoderInfoDescriptor timerCoderDescriptor;
+ private static final String TRANSFORM_ID_PREFIX = "transform-";
+ private static final String COLLECTION_PREFIX = "collection-";
+ private static final String CODER_PREFIX = "coder-";
- private final String functionUrn;
- private final FlinkFnApi.UserDefinedDataStreamFunction userDefinedDataStreamFunction;
+ @Nullable private final FlinkFnApi.CoderInfoDescriptor timerCoderDescriptor;
+ private final String headOperatorFunctionUrn;
+ private final List<FlinkFnApi.UserDefinedDataStreamFunction> userDefinedDataStreamFunctions;
public BeamDataStreamPythonFunctionRunner(
String taskName,
PythonEnvironmentManager environmentManager,
- String functionUrn,
- FlinkFnApi.UserDefinedDataStreamFunction userDefinedDataStreamFunction,
+ String headOperatorFunctionUrn,
+ List<FlinkFnApi.UserDefinedDataStreamFunction> userDefinedDataStreamFunctions,
Map<String, String> jobOptions,
@Nullable FlinkMetricContainer flinkMetricContainer,
- KeyedStateBackend stateBackend,
- TypeSerializer keySerializer,
- TypeSerializer namespaceSerializer,
+ KeyedStateBackend<?> stateBackend,
+ TypeSerializer<?> keySerializer,
+ TypeSerializer<?> namespaceSerializer,
@Nullable TimerRegistration timerRegistration,
MemoryManager memoryManager,
double managedMemoryFraction,
@@ -89,54 +95,95 @@ public class BeamDataStreamPythonFunctionRunner extends BeamPythonFunctionRunner
managedMemoryFraction,
inputCoderDescriptor,
outputCoderDescriptor);
- this.functionUrn = Preconditions.checkNotNull(functionUrn);
- this.userDefinedDataStreamFunction =
- Preconditions.checkNotNull(userDefinedDataStreamFunction);
+ this.headOperatorFunctionUrn = Preconditions.checkNotNull(headOperatorFunctionUrn);
+ Preconditions.checkArgument(
+ userDefinedDataStreamFunctions != null
+ && userDefinedDataStreamFunctions.size() >= 1);
+ this.userDefinedDataStreamFunctions = userDefinedDataStreamFunctions;
this.timerCoderDescriptor = timerCoderDescriptor;
}
@Override
- protected Map<String, RunnerApi.PTransform> getTransforms() {
- // Use ParDoPayload as a wrapper of the actual payload as timer is only supported in ParDo
- RunnerApi.ParDoPayload.Builder payloadBuilder =
- RunnerApi.ParDoPayload.newBuilder()
- .setDoFn(
- RunnerApi.FunctionSpec.newBuilder()
- .setUrn(functionUrn)
- .setPayload(
- org.apache.beam.vendor.grpc.v1p26p0.com.google
- .protobuf.ByteString.copyFrom(
- userDefinedDataStreamFunction
- .toByteArray()))
- .build());
+ protected void buildTransforms(RunnerApi.Components.Builder componentsBuilder) {
+ for (int i = 0; i < userDefinedDataStreamFunctions.size() + 1; i++) {
+ String functionUrn;
+ if (i == 0) {
+ functionUrn = headOperatorFunctionUrn;
+ } else {
+ functionUrn = STATELESS_FUNCTION_URN;
+ }
- if (timerCoderDescriptor != null) {
- payloadBuilder.putTimerFamilySpecs(
- TIMER_ID,
- RunnerApi.TimerFamilySpec.newBuilder()
- .setTimeDomain(
- RunnerApi.TimeDomain.Enum
- .EVENT_TIME) // always set it as event time, this field
- // is not used
- .setTimerFamilyCoderId(WRAPPER_TIMER_CODER_ID)
- .build());
- }
+ FlinkFnApi.UserDefinedDataStreamFunction functionProto;
+ if (i < userDefinedDataStreamFunctions.size()) {
+ functionProto = userDefinedDataStreamFunctions.get(i);
+ } else {
+ // the last function in the operation tree is used to prune the watermark column
+ functionProto = createReviseOutputDataStreamFunctionProto();
+ }
+
+ // Use ParDoPayload as a wrapper of the actual payload as timer is only supported in
+ // ParDo
+ final RunnerApi.ParDoPayload.Builder payloadBuilder =
+ RunnerApi.ParDoPayload.newBuilder()
+ .setDoFn(
+ RunnerApi.FunctionSpec.newBuilder()
+ .setUrn(functionUrn)
+ .setPayload(
+ org.apache.beam.vendor.grpc.v1p26p0.com.google
+ .protobuf.ByteString.copyFrom(
+ functionProto.toByteArray()))
+ .build());
+
+ // Timer is only available in the head operator
+ if (i == 0 && timerCoderDescriptor != null) {
+ payloadBuilder.putTimerFamilySpecs(
+ TIMER_ID,
+ RunnerApi.TimerFamilySpec.newBuilder()
+ // this field is not used, always set it as event time
+ .setTimeDomain(RunnerApi.TimeDomain.Enum.EVENT_TIME)
+ .setTimerFamilyCoderId(WRAPPER_TIMER_CODER_ID)
+ .build());
+ }
- return Collections.singletonMap(
- TRANSFORM_ID,
- RunnerApi.PTransform.newBuilder()
- .setUniqueName(TRANSFORM_ID)
- .setSpec(
- RunnerApi.FunctionSpec.newBuilder()
- .setUrn(
- BeamUrns.getUrn(
- RunnerApi.StandardPTransforms.Primitives
- .PAR_DO))
- .setPayload(payloadBuilder.build().toByteString())
+ final String transformName = TRANSFORM_ID_PREFIX + i;
+
+ final RunnerApi.PTransform.Builder transformBuilder =
+ RunnerApi.PTransform.newBuilder()
+ .setUniqueName(transformName)
+ .setSpec(
+ RunnerApi.FunctionSpec.newBuilder()
+ .setUrn(
+ BeamUrns.getUrn(
+ RunnerApi.StandardPTransforms.Primitives
+ .PAR_DO))
+ .setPayload(payloadBuilder.build().toByteString())
+ .build());
+
+ // prepare inputs
+ if (i == 0) {
+ transformBuilder.putInputs(MAIN_INPUT_NAME, INPUT_COLLECTION_ID);
+ } else {
+ transformBuilder.putInputs(MAIN_INPUT_NAME, COLLECTION_PREFIX + (i - 1));
+ }
+
+ // prepare outputs
+ if (i == userDefinedDataStreamFunctions.size()) {
+ transformBuilder.putOutputs(MAIN_OUTPUT_NAME, OUTPUT_COLLECTION_ID);
+ } else {
+ transformBuilder.putOutputs(MAIN_OUTPUT_NAME, COLLECTION_PREFIX + i);
+
+ componentsBuilder
+ .putPcollections(
+ COLLECTION_PREFIX + i,
+ RunnerApi.PCollection.newBuilder()
+ .setWindowingStrategyId(WINDOW_STRATEGY)
+ .setCoderId(CODER_PREFIX + i)
.build())
- .putInputs(MAIN_INPUT_NAME, INPUT_COLLECTION_ID)
- .putOutputs(MAIN_OUTPUT_NAME, OUTPUT_COLLECTION_ID)
- .build());
+ .putCoders(CODER_PREFIX + i, createCoderProto(inputCoderDescriptor));
+ }
+
+ componentsBuilder.putTransforms(transformName, transformBuilder.build());
+ }
}
@Override
@@ -144,7 +191,7 @@ public class BeamDataStreamPythonFunctionRunner extends BeamPythonFunctionRunner
if (timerCoderDescriptor != null) {
RunnerApi.ExecutableStagePayload.TimerId timerId =
RunnerApi.ExecutableStagePayload.TimerId.newBuilder()
- .setTransformId(TRANSFORM_ID)
+ .setTransformId(TRANSFORM_ID_PREFIX + 0)
.setLocalName(TIMER_ID)
.build();
return Collections.singletonList(TimerReference.fromTimerId(timerId, components));
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
index 2aa38b2..0e6fbba 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
@@ -85,13 +85,14 @@ import java.util.Map;
import java.util.Optional;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.function.BiConsumer;
+import java.util.stream.Collectors;
import static org.apache.beam.runners.core.construction.BeamUrns.getUrn;
import static org.apache.flink.python.Constants.INPUT_COLLECTION_ID;
import static org.apache.flink.python.Constants.OUTPUT_COLLECTION_ID;
import static org.apache.flink.python.Constants.TIMER_CODER_ID;
-import static org.apache.flink.python.Constants.TRANSFORM_ID;
import static org.apache.flink.python.Constants.WINDOW_CODER_ID;
+import static org.apache.flink.python.Constants.WINDOW_STRATEGY;
import static org.apache.flink.python.Constants.WRAPPER_TIMER_CODER_ID;
import static org.apache.flink.streaming.api.utils.ProtoUtils.createCoderProto;
@@ -103,8 +104,6 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
private static final String INPUT_CODER_ID = "input_coder";
private static final String OUTPUT_CODER_ID = "output_coder";
- private static final String WINDOW_STRATEGY = "windowing_strategy";
-
private static final String MANAGED_MEMORY_RESOURCE_ID = "python-process-managed-memory";
private static final String PYTHON_WORKER_MEMORY_LIMIT = "_PYTHON_WORKER_MEMORY_LIMIT";
@@ -444,8 +443,7 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
WRAPPER_TIMER_CODER_ID, wrapperTimerCoderProto);
});
- getTransforms().forEach(componentsBuilder::putTransforms);
-
+ buildTransforms(componentsBuilder);
RunnerApi.Components components = componentsBuilder.build();
PipelineNode.PCollectionNode input =
@@ -456,9 +454,9 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
List<UserStateReference> userStates = Collections.EMPTY_LIST;
List<TimerReference> timers = getTimers(components);
List<PipelineNode.PTransformNode> transforms =
- Collections.singletonList(
- PipelineNode.pTransform(
- TRANSFORM_ID, components.getTransformsOrThrow(TRANSFORM_ID)));
+ components.getTransformsMap().keySet().stream()
+ .map(id -> PipelineNode.pTransform(id, components.getTransformsOrThrow(id)))
+ .collect(Collectors.toList());
List<PipelineNode.PCollectionNode> outputs =
Collections.singletonList(
PipelineNode.pCollection(
@@ -512,7 +510,7 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
.build();
}
- protected abstract Map<String, RunnerApi.PTransform> getTransforms();
+ protected abstract void buildTransforms(RunnerApi.Components.Builder componentsBuilder);
protected abstract List<TimerReference> getTimers(RunnerApi.Components components);
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/utils/ProtoUtils.java b/flink-python/src/main/java/org/apache/flink/streaming/api/utils/ProtoUtils.java
index fbdab2b..9c60500 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/utils/ProtoUtils.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/utils/ProtoUtils.java
@@ -28,10 +28,13 @@ import org.apache.flink.table.functions.python.PythonFunctionInfo;
import org.apache.flink.table.functions.python.PythonFunctionKind;
import org.apache.flink.table.planner.typeutils.DataViewUtils;
import org.apache.flink.table.types.logical.RowType;
+import org.apache.flink.util.Preconditions;
import com.google.protobuf.ByteString;
import org.apache.beam.model.pipeline.v1.RunnerApi;
+import java.util.ArrayList;
+import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
@@ -122,7 +125,7 @@ public enum ProtoUtils {
return builder.build();
}
- public static FlinkFnApi.UserDefinedDataStreamFunction getUserDefinedDataStreamFunctionProto(
+ public static FlinkFnApi.UserDefinedDataStreamFunction createUserDefinedDataStreamFunctionProto(
DataStreamPythonFunctionInfo dataStreamPythonFunctionInfo,
RuntimeContext runtimeContext,
Map<String, String> internalParameters,
@@ -174,21 +177,60 @@ public enum ProtoUtils {
}
public static FlinkFnApi.UserDefinedDataStreamFunction
- getUserDefinedDataStreamStatefulFunctionProto(
+ createReviseOutputDataStreamFunctionProto() {
+ return FlinkFnApi.UserDefinedDataStreamFunction.newBuilder()
+ .setFunctionType(
+ FlinkFnApi.UserDefinedDataStreamFunction.FunctionType.REVISE_OUTPUT)
+ .build();
+ }
+
+ public static List<FlinkFnApi.UserDefinedDataStreamFunction>
+ createUserDefinedDataStreamFunctionProtos(
+ DataStreamPythonFunctionInfo dataStreamPythonFunctionInfo,
+ RuntimeContext runtimeContext,
+ Map<String, String> internalParameters,
+ boolean inBatchExecutionMode) {
+ List<FlinkFnApi.UserDefinedDataStreamFunction> results = new ArrayList<>();
+
+ Object[] inputs = dataStreamPythonFunctionInfo.getInputs();
+ if (inputs != null && inputs.length > 0) {
+ Preconditions.checkArgument(inputs.length == 1);
+ results.addAll(
+ createUserDefinedDataStreamFunctionProtos(
+ (DataStreamPythonFunctionInfo) inputs[0],
+ runtimeContext,
+ internalParameters,
+ inBatchExecutionMode));
+ }
+
+ results.add(
+ createUserDefinedDataStreamFunctionProto(
+ dataStreamPythonFunctionInfo,
+ runtimeContext,
+ internalParameters,
+ inBatchExecutionMode));
+ return results;
+ }
+
+ public static List<FlinkFnApi.UserDefinedDataStreamFunction>
+ createUserDefinedDataStreamStatefulFunctionProtos(
DataStreamPythonFunctionInfo dataStreamPythonFunctionInfo,
RuntimeContext runtimeContext,
Map<String, String> internalParameters,
- TypeInformation keyTypeInfo,
+ TypeInformation<?> keyTypeInfo,
boolean inBatchExecutionMode) {
- FlinkFnApi.UserDefinedDataStreamFunction userDefinedDataStreamFunction =
- getUserDefinedDataStreamFunctionProto(
+ List<FlinkFnApi.UserDefinedDataStreamFunction> results =
+ createUserDefinedDataStreamFunctionProtos(
dataStreamPythonFunctionInfo,
runtimeContext,
internalParameters,
inBatchExecutionMode);
+
+ // set the key typeinfo for the head operator
FlinkFnApi.TypeInfo builtKeyTypeInfo =
PythonTypeUtils.TypeInfoToProtoConverter.toTypeInfoProto(keyTypeInfo);
- return userDefinedDataStreamFunction.toBuilder().setKeyTypeInfo(builtKeyTypeInfo).build();
+ results.set(0, results.get(0).toBuilder().setKeyTypeInfo(builtKeyTypeInfo).build());
+ return results;
}
public static RunnerApi.Coder createCoderProto(
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractOneInputPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractOneInputPythonFunctionOperator.java
similarity index 91%
copy from flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractOneInputPythonFunctionOperator.java
copy to flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractOneInputPythonFunctionOperator.java
index 5c8a35a..503bc09 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractOneInputPythonFunctionOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractOneInputPythonFunctionOperator.java
@@ -16,12 +16,13 @@
* limitations under the License.
*/
-package org.apache.flink.streaming.api.operators.python;
+package org.apache.flink.table.runtime.operators.python;
import org.apache.flink.annotation.Internal;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.operators.BoundedOneInput;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.operators.python.AbstractPythonFunctionOperator;
/** Base class for all one input stream operators to execute Python functions. */
@Internal
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractStatelessFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractStatelessFunctionOperator.java
index 7476632..f83f560 100644
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractStatelessFunctionOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractStatelessFunctionOperator.java
@@ -27,7 +27,6 @@ import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.core.memory.ManagedMemoryUseCase;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.PythonFunctionRunner;
-import org.apache.flink.streaming.api.operators.python.AbstractOneInputPythonFunctionOperator;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.types.logical.RowType;
@@ -129,16 +128,13 @@ public abstract class AbstractStatelessFunctionOperator<IN, OUT, UDFIN>
@Override
public PythonFunctionRunner createPythonFunctionRunner() throws IOException {
- return new BeamTablePythonFunctionRunner(
+ return BeamTablePythonFunctionRunner.stateless(
getRuntimeContext().getTaskName(),
createPythonEnvironmentManager(),
getFunctionUrn(),
getUserDefinedFunctionsProto(),
jobOptions,
getFlinkMetricContainer(),
- null,
- null,
- null,
getContainingTask().getEnvironment().getMemoryManager(),
getOperatorConfig()
.getManagedMemoryFractionOperatorUseCaseOfSlot(
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java
index c845e61..19deb7d 100644
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java
@@ -30,7 +30,6 @@ import org.apache.flink.core.memory.ManagedMemoryUseCase;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.PythonFunctionRunner;
import org.apache.flink.python.PythonOptions;
-import org.apache.flink.streaming.api.operators.python.AbstractOneInputPythonFunctionOperator;
import org.apache.flink.streaming.api.utils.ProtoUtils;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.table.data.RowData;
@@ -39,6 +38,7 @@ import org.apache.flink.table.functions.python.PythonEnv;
import org.apache.flink.table.planner.plan.utils.KeySelectorUtil;
import org.apache.flink.table.planner.typeutils.DataViewUtils;
import org.apache.flink.table.runtime.keyselector.RowDataKeySelector;
+import org.apache.flink.table.runtime.operators.python.AbstractOneInputPythonFunctionOperator;
import org.apache.flink.table.runtime.operators.python.utils.StreamRecordRowDataWrappingCollector;
import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
@@ -177,7 +177,7 @@ public abstract class AbstractPythonStreamAggregateOperator
@Override
public PythonFunctionRunner createPythonFunctionRunner() throws Exception {
- return new BeamTablePythonFunctionRunner(
+ return BeamTablePythonFunctionRunner.stateful(
getRuntimeContext().getTaskName(),
createPythonEnvironmentManager(),
getFunctionUrn(),
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
index ecfeb5a..c0d86d6 100644
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
@@ -59,9 +59,9 @@ public class BeamTablePythonFunctionRunner extends BeamPythonFunctionRunner {
GeneratedMessageV3 userDefinedFunctionProto,
Map<String, String> jobOptions,
FlinkMetricContainer flinkMetricContainer,
- KeyedStateBackend keyedStateBackend,
- TypeSerializer keySerializer,
- TypeSerializer namespaceSerializer,
+ KeyedStateBackend<?> keyedStateBackend,
+ TypeSerializer<?> keySerializer,
+ TypeSerializer<?> namespaceSerializer,
MemoryManager memoryManager,
double managedMemoryFraction,
FlinkFnApi.CoderInfoDescriptor inputCoderDescriptor,
@@ -84,8 +84,8 @@ public class BeamTablePythonFunctionRunner extends BeamPythonFunctionRunner {
}
@Override
- protected Map<String, RunnerApi.PTransform> getTransforms() {
- return Collections.singletonMap(
+ protected void buildTransforms(RunnerApi.Components.Builder componentsBuilder) {
+ componentsBuilder.putTransforms(
TRANSFORM_ID,
RunnerApi.PTransform.newBuilder()
.setUniqueName(TRANSFORM_ID)
@@ -116,4 +116,61 @@ public class BeamTablePythonFunctionRunner extends BeamPythonFunctionRunner {
public void processTimer(byte[] timerData) throws Exception {
throw new UnsupportedOperationException();
}
+
+ public static BeamTablePythonFunctionRunner stateless(
+ String taskName,
+ PythonEnvironmentManager environmentManager,
+ String functionUrn,
+ GeneratedMessageV3 userDefinedFunctionProto,
+ Map<String, String> jobOptions,
+ FlinkMetricContainer flinkMetricContainer,
+ MemoryManager memoryManager,
+ double managedMemoryFraction,
+ FlinkFnApi.CoderInfoDescriptor inputCoderDescriptor,
+ FlinkFnApi.CoderInfoDescriptor outputCoderDescriptor) {
+ return new BeamTablePythonFunctionRunner(
+ taskName,
+ environmentManager,
+ functionUrn,
+ userDefinedFunctionProto,
+ jobOptions,
+ flinkMetricContainer,
+ null,
+ null,
+ null,
+ memoryManager,
+ managedMemoryFraction,
+ inputCoderDescriptor,
+ outputCoderDescriptor);
+ }
+
+ public static BeamTablePythonFunctionRunner stateful(
+ String taskName,
+ PythonEnvironmentManager environmentManager,
+ String functionUrn,
+ GeneratedMessageV3 userDefinedFunctionProto,
+ Map<String, String> jobOptions,
+ FlinkMetricContainer flinkMetricContainer,
+ KeyedStateBackend<?> keyedStateBackend,
+ TypeSerializer<?> keySerializer,
+ TypeSerializer<?> namespaceSerializer,
+ MemoryManager memoryManager,
+ double managedMemoryFraction,
+ FlinkFnApi.CoderInfoDescriptor inputCoderDescriptor,
+ FlinkFnApi.CoderInfoDescriptor outputCoderDescriptor) {
+ return new BeamTablePythonFunctionRunner(
+ taskName,
+ environmentManager,
+ functionUrn,
+ userDefinedFunctionProto,
+ jobOptions,
+ flinkMetricContainer,
+ keyedStateBackend,
+ keySerializer,
+ namespaceSerializer,
+ memoryManager,
+ managedMemoryFraction,
+ inputCoderDescriptor,
+ outputCoderDescriptor);
+ }
}
diff --git a/flink-python/src/test/java/org/apache/flink/python/chain/PythonOperatorChainingOptimizerTest.java b/flink-python/src/test/java/org/apache/flink/python/chain/PythonOperatorChainingOptimizerTest.java
new file mode 100644
index 0000000..55b7d90
--- /dev/null
+++ b/flink-python/src/test/java/org/apache/flink/python/chain/PythonOperatorChainingOptimizerTest.java
@@ -0,0 +1,544 @@
+/*
+ * 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.
+ */
+
+package org.apache.flink.python.chain;
+
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.dag.Transformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.api.scala.typeutils.Types;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.memory.ManagedMemoryUseCase;
+import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunction;
+import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionInfo;
+import org.apache.flink.streaming.api.operators.ChainingStrategy;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
+import org.apache.flink.streaming.api.operators.python.PythonKeyedCoProcessOperator;
+import org.apache.flink.streaming.api.operators.python.PythonKeyedProcessOperator;
+import org.apache.flink.streaming.api.operators.python.PythonProcessOperator;
+import org.apache.flink.streaming.api.transformations.OneInputTransformation;
+import org.apache.flink.streaming.api.transformations.SourceTransformation;
+import org.apache.flink.streaming.api.transformations.TwoInputTransformation;
+
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+
+/** Tests for {@link PythonOperatorChainingOptimizer}. */
+public class PythonOperatorChainingOptimizerTest {
+
+ @Test
+ public void testChainedTransformationPropertiesCorrectlySet() throws Exception {
+ PythonKeyedProcessOperator<?> keyedProcessOperator =
+ createKeyedProcessOperator(
+ "f1", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+ PythonProcessOperator<?, ?> processOperator =
+ createProcessOperator("f2", Types.STRING(), Types.STRING());
+
+ Transformation<?> sourceTransformation = mock(SourceTransformation.class);
+ OneInputTransformation<?, ?> keyedProcessTransformation =
+ new OneInputTransformation(
+ sourceTransformation,
+ "keyedProcess",
+ keyedProcessOperator,
+ keyedProcessOperator.getProducedType(),
+ 2);
+ keyedProcessTransformation.setUid("uid");
+ keyedProcessTransformation.setSlotSharingGroup("group");
+ keyedProcessTransformation.setCoLocationGroupKey("col");
+ keyedProcessTransformation.setMaxParallelism(64);
+ keyedProcessTransformation.declareManagedMemoryUseCaseAtOperatorScope(
+ ManagedMemoryUseCase.OPERATOR, 5);
+ keyedProcessTransformation.declareManagedMemoryUseCaseAtSlotScope(
+ ManagedMemoryUseCase.PYTHON);
+ keyedProcessTransformation.declareManagedMemoryUseCaseAtSlotScope(
+ ManagedMemoryUseCase.STATE_BACKEND);
+ keyedProcessTransformation.setBufferTimeout(1000L);
+ keyedProcessTransformation.setChainingStrategy(ChainingStrategy.HEAD);
+ Transformation<?> processTransformation =
+ new OneInputTransformation(
+ keyedProcessTransformation,
+ "process",
+ processOperator,
+ processOperator.getProducedType(),
+ 2);
+ processTransformation.setSlotSharingGroup("group");
+ processTransformation.declareManagedMemoryUseCaseAtOperatorScope(
+ ManagedMemoryUseCase.OPERATOR, 10);
+ processTransformation.declareManagedMemoryUseCaseAtSlotScope(ManagedMemoryUseCase.PYTHON);
+ processTransformation.setMaxParallelism(64);
+ processTransformation.setBufferTimeout(500L);
+
+ List<Transformation<?>> transformations = new ArrayList<>();
+ transformations.add(sourceTransformation);
+ transformations.add(keyedProcessTransformation);
+ transformations.add(processTransformation);
+
+ PythonOperatorChainingOptimizer optimizer =
+ new PythonOperatorChainingOptimizer(transformations, true);
+ List<Transformation<?>> optimized = optimizer.optimize();
+ assertEquals(2, optimized.size());
+
+ OneInputTransformation<?, ?> chainedTransformation =
+ (OneInputTransformation<?, ?>) optimized.get(1);
+ assertEquals(2, chainedTransformation.getParallelism());
+ assertEquals(sourceTransformation.getOutputType(), chainedTransformation.getInputType());
+ assertEquals(processOperator.getProducedType(), chainedTransformation.getOutputType());
+ assertEquals(keyedProcessTransformation.getUid(), chainedTransformation.getUid());
+ assertEquals("group", chainedTransformation.getSlotSharingGroup().get().getName());
+ assertEquals("col", chainedTransformation.getCoLocationGroupKey());
+ assertEquals(64, chainedTransformation.getMaxParallelism());
+ assertEquals(500L, chainedTransformation.getBufferTimeout());
+ assertEquals(
+ 15,
+ (int)
+ chainedTransformation
+ .getManagedMemoryOperatorScopeUseCaseWeights()
+ .getOrDefault(ManagedMemoryUseCase.OPERATOR, 0));
+ assertEquals(
+ ChainingStrategy.HEAD,
+ chainedTransformation.getOperatorFactory().getChainingStrategy());
+ assertTrue(
+ chainedTransformation
+ .getManagedMemorySlotScopeUseCases()
+ .contains(ManagedMemoryUseCase.PYTHON));
+ assertTrue(
+ chainedTransformation
+ .getManagedMemorySlotScopeUseCases()
+ .contains(ManagedMemoryUseCase.STATE_BACKEND));
+
+ OneInputStreamOperator<?, ?> chainedOperator = chainedTransformation.getOperator();
+ assertTrue(chainedOperator instanceof PythonKeyedProcessOperator);
+ validateChainedPythonFunctions(
+ ((PythonKeyedProcessOperator<?>) chainedOperator).getPythonFunctionInfo(),
+ "f2",
+ "f1");
+ }
+
+ @Test
+ public void testChainingMultipleOperators() throws Exception {
+ PythonKeyedProcessOperator<?> keyedProcessOperator =
+ createKeyedProcessOperator(
+ "f1", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+ PythonProcessOperator<?, ?> processOperator1 =
+ createProcessOperator("f2", Types.STRING(), Types.LONG());
+ PythonProcessOperator<?, ?> processOperator2 =
+ createProcessOperator("f3", Types.LONG(), Types.INT());
+
+ Transformation<?> sourceTransformation = mock(SourceTransformation.class);
+ OneInputTransformation<?, ?> keyedProcessTransformation =
+ new OneInputTransformation(
+ sourceTransformation,
+ "keyedProcess",
+ keyedProcessOperator,
+ keyedProcessOperator.getProducedType(),
+ 2);
+ Transformation<?> processTransformation1 =
+ new OneInputTransformation(
+ keyedProcessTransformation,
+ "process",
+ processOperator1,
+ processOperator1.getProducedType(),
+ 2);
+ Transformation<?> processTransformation2 =
+ new OneInputTransformation(
+ processTransformation1,
+ "process",
+ processOperator2,
+ processOperator2.getProducedType(),
+ 2);
+
+ List<Transformation<?>> transformations = new ArrayList<>();
+ transformations.add(sourceTransformation);
+ transformations.add(keyedProcessTransformation);
+ transformations.add(processTransformation1);
+ transformations.add(processTransformation2);
+
+ PythonOperatorChainingOptimizer optimizer =
+ new PythonOperatorChainingOptimizer(transformations, true);
+ List<Transformation<?>> optimized = optimizer.optimize();
+ assertEquals(2, optimized.size());
+
+ OneInputTransformation<?, ?> chainedTransformation =
+ (OneInputTransformation<?, ?>) optimized.get(1);
+ assertEquals(sourceTransformation.getOutputType(), chainedTransformation.getInputType());
+ assertEquals(processOperator2.getProducedType(), chainedTransformation.getOutputType());
+
+ OneInputStreamOperator<?, ?> chainedOperator = chainedTransformation.getOperator();
+ assertTrue(chainedOperator instanceof PythonKeyedProcessOperator);
+ validateChainedPythonFunctions(
+ ((PythonKeyedProcessOperator<?>) chainedOperator).getPythonFunctionInfo(),
+ "f3",
+ "f2",
+ "f1");
+ }
+
+ @Test
+ public void testChainingNonKeyedOperators() throws Exception {
+ PythonProcessOperator<?, ?> processOperator1 =
+ createProcessOperator(
+ "f1", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+ PythonProcessOperator<?, ?> processOperator2 =
+ createProcessOperator("f2", Types.STRING(), Types.INT());
+
+ Transformation<?> sourceTransformation = mock(SourceTransformation.class);
+ OneInputTransformation<?, ?> processTransformation1 =
+ new OneInputTransformation(
+ sourceTransformation,
+ "Process1",
+ processOperator1,
+ processOperator1.getProducedType(),
+ 2);
+ Transformation<?> processTransformation2 =
+ new OneInputTransformation(
+ processTransformation1,
+ "process2",
+ processOperator2,
+ processOperator2.getProducedType(),
+ 2);
+
+ List<Transformation<?>> transformations = new ArrayList<>();
+ transformations.add(sourceTransformation);
+ transformations.add(processTransformation1);
+ transformations.add(processTransformation2);
+
+ PythonOperatorChainingOptimizer optimizer =
+ new PythonOperatorChainingOptimizer(transformations, true);
+ List<Transformation<?>> optimized = optimizer.optimize();
+ assertEquals(2, optimized.size());
+
+ OneInputTransformation<?, ?> chainedTransformation =
+ (OneInputTransformation<?, ?>) optimized.get(1);
+ assertEquals(sourceTransformation.getOutputType(), chainedTransformation.getInputType());
+ assertEquals(processOperator2.getProducedType(), chainedTransformation.getOutputType());
+
+ OneInputStreamOperator<?, ?> chainedOperator = chainedTransformation.getOperator();
+ assertTrue(chainedOperator instanceof PythonProcessOperator);
+ validateChainedPythonFunctions(
+ ((PythonProcessOperator<?, ?>) chainedOperator).getPythonFunctionInfo(),
+ "f2",
+ "f1");
+ }
+
+ @Test
+ public void testContinuousKeyedOperators() throws Exception {
+ PythonKeyedProcessOperator<?> keyedProcessOperator1 =
+ createKeyedProcessOperator(
+ "f1",
+ new RowTypeInfo(Types.INT(), Types.INT()),
+ new RowTypeInfo(Types.INT(), Types.INT()));
+ PythonKeyedProcessOperator<?> keyedProcessOperator2 =
+ createKeyedProcessOperator(
+ "f2", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+
+ Transformation<?> sourceTransformation = mock(SourceTransformation.class);
+ OneInputTransformation<?, ?> processTransformation1 =
+ new OneInputTransformation(
+ sourceTransformation,
+ "KeyedProcess1",
+ keyedProcessOperator1,
+ keyedProcessOperator1.getProducedType(),
+ 2);
+ OneInputTransformation<?, ?> processTransformation2 =
+ new OneInputTransformation(
+ processTransformation1,
+ "KeyedProcess2",
+ keyedProcessOperator2,
+ keyedProcessOperator2.getProducedType(),
+ 2);
+
+ List<Transformation<?>> transformations = new ArrayList<>();
+ transformations.add(sourceTransformation);
+ transformations.add(processTransformation1);
+ transformations.add(processTransformation2);
+
+ PythonOperatorChainingOptimizer optimizer =
+ new PythonOperatorChainingOptimizer(transformations, true);
+ List<Transformation<?>> optimized = optimizer.optimize();
+ assertEquals(3, optimized.size());
+
+ assertEquals(processTransformation1, optimized.get(1));
+ assertEquals(processTransformation2, optimized.get(2));
+ }
+
+ @Test
+ public void testMultipleChainedOperators() throws Exception {
+ PythonKeyedProcessOperator<?> keyedProcessOperator1 =
+ createKeyedProcessOperator(
+ "f1", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+ PythonProcessOperator<?, ?> processOperator1 =
+ createProcessOperator(
+ "f2", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+ PythonProcessOperator<?, ?> processOperator2 =
+ createProcessOperator(
+ "f3", new RowTypeInfo(Types.INT(), Types.INT()), Types.LONG());
+
+ PythonKeyedProcessOperator<?> keyedProcessOperator2 =
+ createKeyedProcessOperator(
+ "f4", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+ PythonProcessOperator<?, ?> processOperator3 =
+ createProcessOperator(
+ "f5", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+
+ Transformation<?> sourceTransformation = mock(SourceTransformation.class);
+ OneInputTransformation<?, ?> keyedProcessTransformation1 =
+ new OneInputTransformation(
+ sourceTransformation,
+ "keyedProcess",
+ keyedProcessOperator1,
+ keyedProcessOperator1.getProducedType(),
+ 2);
+ Transformation<?> processTransformation1 =
+ new OneInputTransformation(
+ keyedProcessTransformation1,
+ "process",
+ processOperator1,
+ processOperator1.getProducedType(),
+ 2);
+ Transformation<?> processTransformation2 =
+ new OneInputTransformation(
+ processTransformation1,
+ "process",
+ processOperator2,
+ processOperator2.getProducedType(),
+ 2);
+
+ OneInputTransformation<?, ?> keyedProcessTransformation2 =
+ new OneInputTransformation(
+ processTransformation2,
+ "keyedProcess",
+ keyedProcessOperator2,
+ keyedProcessOperator2.getProducedType(),
+ 2);
+ Transformation<?> processTransformation3 =
+ new OneInputTransformation(
+ keyedProcessTransformation2,
+ "process",
+ processOperator3,
+ processOperator3.getProducedType(),
+ 2);
+
+ List<Transformation<?>> transformations = new ArrayList<>();
+ transformations.add(sourceTransformation);
+ transformations.add(keyedProcessTransformation1);
+ transformations.add(processTransformation1);
+ transformations.add(processTransformation2);
+ transformations.add(keyedProcessTransformation2);
+ transformations.add(processTransformation3);
+
+ PythonOperatorChainingOptimizer optimizer =
+ new PythonOperatorChainingOptimizer(transformations, true);
+ List<Transformation<?>> optimized = optimizer.optimize();
+ assertEquals(3, optimized.size());
+
+ OneInputTransformation<?, ?> chainedTransformation1 =
+ (OneInputTransformation<?, ?>) optimized.get(1);
+ assertEquals(sourceTransformation.getOutputType(), chainedTransformation1.getInputType());
+ assertEquals(processOperator2.getProducedType(), chainedTransformation1.getOutputType());
+
+ OneInputTransformation<?, ?> chainedTransformation2 =
+ (OneInputTransformation<?, ?>) optimized.get(2);
+ assertEquals(processOperator2.getProducedType(), chainedTransformation2.getInputType());
+ assertEquals(processOperator3.getProducedType(), chainedTransformation2.getOutputType());
+
+ OneInputStreamOperator<?, ?> chainedOperator1 = chainedTransformation1.getOperator();
+ assertTrue(chainedOperator1 instanceof PythonKeyedProcessOperator);
+ validateChainedPythonFunctions(
+ ((PythonKeyedProcessOperator<?>) chainedOperator1).getPythonFunctionInfo(),
+ "f3",
+ "f2",
+ "f1");
+
+ OneInputStreamOperator<?, ?> chainedOperator2 = chainedTransformation2.getOperator();
+ assertTrue(chainedOperator2 instanceof PythonKeyedProcessOperator);
+ validateChainedPythonFunctions(
+ ((PythonKeyedProcessOperator<?>) chainedOperator2).getPythonFunctionInfo(),
+ "f5",
+ "f4");
+ }
+
+ @Test
+ public void testChainingTwoInputOperators() throws Exception {
+ PythonKeyedCoProcessOperator<?> keyedCoProcessOperator1 =
+ createCoKeyedProcessOperator(
+ "f1",
+ new RowTypeInfo(Types.INT(), Types.STRING()),
+ new RowTypeInfo(Types.INT(), Types.INT()),
+ Types.STRING());
+ PythonProcessOperator<?, ?> processOperator1 =
+ createProcessOperator(
+ "f2", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+ PythonProcessOperator<?, ?> processOperator2 =
+ createProcessOperator(
+ "f3", new RowTypeInfo(Types.INT(), Types.INT()), Types.LONG());
+
+ PythonKeyedProcessOperator<?> keyedProcessOperator2 =
+ createKeyedProcessOperator(
+ "f4", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+ PythonProcessOperator<?, ?> processOperator3 =
+ createProcessOperator(
+ "f5", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
+
+ Transformation<?> sourceTransformation1 = mock(SourceTransformation.class);
+ Transformation<?> sourceTransformation2 = mock(SourceTransformation.class);
+ TwoInputTransformation<?, ?, ?> keyedCoProcessTransformation =
+ new TwoInputTransformation(
+ sourceTransformation1,
+ sourceTransformation2,
+ "keyedCoProcess",
+ keyedCoProcessOperator1,
+ keyedCoProcessOperator1.getProducedType(),
+ 2);
+ Transformation<?> processTransformation1 =
+ new OneInputTransformation(
+ keyedCoProcessTransformation,
+ "process",
+ processOperator1,
+ processOperator1.getProducedType(),
+ 2);
+ Transformation<?> processTransformation2 =
+ new OneInputTransformation(
+ processTransformation1,
+ "process",
+ processOperator2,
+ processOperator2.getProducedType(),
+ 2);
+
+ OneInputTransformation<?, ?> keyedProcessTransformation =
+ new OneInputTransformation(
+ processTransformation2,
+ "keyedProcess",
+ keyedProcessOperator2,
+ keyedProcessOperator2.getProducedType(),
+ 2);
+ Transformation<?> processTransformation3 =
+ new OneInputTransformation(
+ keyedProcessTransformation,
+ "process",
+ processOperator3,
+ processOperator3.getProducedType(),
+ 2);
+
+ List<Transformation<?>> transformations = new ArrayList<>();
+ transformations.add(sourceTransformation1);
+ transformations.add(sourceTransformation2);
+ transformations.add(keyedCoProcessTransformation);
+ transformations.add(processTransformation1);
+ transformations.add(processTransformation2);
+ transformations.add(keyedProcessTransformation);
+ transformations.add(processTransformation3);
+
+ PythonOperatorChainingOptimizer optimizer =
+ new PythonOperatorChainingOptimizer(transformations, true);
+ List<Transformation<?>> optimized = optimizer.optimize();
+ assertEquals(4, optimized.size());
+
+ TwoInputTransformation<?, ?, ?> chainedTransformation1 =
+ (TwoInputTransformation<?, ?, ?>) optimized.get(2);
+ assertEquals(sourceTransformation1.getOutputType(), chainedTransformation1.getInputType1());
+ assertEquals(sourceTransformation2.getOutputType(), chainedTransformation1.getInputType2());
+ assertEquals(processOperator2.getProducedType(), chainedTransformation1.getOutputType());
+
+ OneInputTransformation<?, ?> chainedTransformation2 =
+ (OneInputTransformation<?, ?>) optimized.get(3);
+ assertEquals(processOperator2.getProducedType(), chainedTransformation2.getInputType());
+ assertEquals(processOperator3.getProducedType(), chainedTransformation2.getOutputType());
+
+ TwoInputStreamOperator<?, ?, ?> chainedOperator1 = chainedTransformation1.getOperator();
+ assertTrue(chainedOperator1 instanceof PythonKeyedCoProcessOperator);
+ validateChainedPythonFunctions(
+ ((PythonKeyedCoProcessOperator<?>) chainedOperator1).getPythonFunctionInfo(),
+ "f3",
+ "f2",
+ "f1");
+
+ OneInputStreamOperator<?, ?> chainedOperator2 = chainedTransformation2.getOperator();
+ assertTrue(chainedOperator2 instanceof PythonKeyedProcessOperator);
+ validateChainedPythonFunctions(
+ ((PythonKeyedProcessOperator<?>) chainedOperator2).getPythonFunctionInfo(),
+ "f5",
+ "f4");
+ }
+
+ // ----------------------- Utility Methods -----------------------
+
+ private void validateChainedPythonFunctions(
+ DataStreamPythonFunctionInfo pythonFunctionInfo,
+ String... expectedChainedPythonFunctions) {
+ for (String expectedPythonFunction : expectedChainedPythonFunctions) {
+ assertArrayEquals(
+ expectedPythonFunction.getBytes(),
+ pythonFunctionInfo.getPythonFunction().getSerializedPythonFunction());
+ Object[] inputs = pythonFunctionInfo.getInputs();
+ if (inputs.length > 0) {
+ assertEquals(1, inputs.length);
+ pythonFunctionInfo = (DataStreamPythonFunctionInfo) inputs[0];
+ } else {
+ pythonFunctionInfo = null;
+ }
+ }
+
+ assertNull(pythonFunctionInfo);
+ }
+
+ private static <OUT> PythonKeyedProcessOperator<OUT> createKeyedProcessOperator(
+ String functionContent,
+ RowTypeInfo inputTypeInfo,
+ TypeInformation<OUT> outputTypeInfo) {
+ return new PythonKeyedProcessOperator<>(
+ new Configuration(),
+ new DataStreamPythonFunctionInfo(
+ new DataStreamPythonFunction(functionContent.getBytes(), null), -1),
+ inputTypeInfo,
+ outputTypeInfo);
+ }
+
+ private static <OUT> PythonKeyedCoProcessOperator<OUT> createCoKeyedProcessOperator(
+ String functionContent,
+ RowTypeInfo inputTypeInfo1,
+ RowTypeInfo inputTypeInfo2,
+ TypeInformation<OUT> outputTypeInfo) {
+ return new PythonKeyedCoProcessOperator(
+ new Configuration(),
+ new DataStreamPythonFunctionInfo(
+ new DataStreamPythonFunction(functionContent.getBytes(), null), -1),
+ inputTypeInfo1,
+ inputTypeInfo2,
+ outputTypeInfo);
+ }
+
+ private static <IN, OUT> PythonProcessOperator<IN, OUT> createProcessOperator(
+ String functionContent,
+ TypeInformation<IN> inputTypeInfo,
+ TypeInformation<OUT> outputTypeInfo) {
+ return new PythonProcessOperator<>(
+ new Configuration(),
+ new DataStreamPythonFunctionInfo(
+ new DataStreamPythonFunction(functionContent.getBytes(), null), -1),
+ inputTypeInfo,
+ outputTypeInfo);
+ }
+}
[flink] 01/02: [hotfix][python][tests] Refactor the test cases in
DataStreamTests to reduce the number of itcases
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 1b8bcf2dedf197bd82ff3e281cf450e62bdeb4a5
Author: Dian Fu <di...@apache.org>
AuthorDate: Mon Aug 9 16:29:43 2021 +0800
[hotfix][python][tests] Refactor the test cases in DataStreamTests to reduce the number of itcases
---
.../pyflink/datastream/tests/test_data_stream.py | 567 ++++++++-------------
1 file changed, 203 insertions(+), 364 deletions(-)
diff --git a/flink-python/pyflink/datastream/tests/test_data_stream.py b/flink-python/pyflink/datastream/tests/test_data_stream.py
index 24a230f..57dc80d 100644
--- a/flink-python/pyflink/datastream/tests/test_data_stream.py
+++ b/flink-python/pyflink/datastream/tests/test_data_stream.py
@@ -55,82 +55,135 @@ class DataStreamTests(object):
def tearDown(self) -> None:
self.test_sink.clear()
- def test_map(self):
- self.env.set_parallelism(1)
+ def assert_equals_sorted(self, expected, actual):
+ expected.sort()
+ actual.sort()
+ self.assertEqual(expected, actual)
+
+ def test_basic_operations(self):
+ ds = self.env.from_collection([('ab', decimal.Decimal(1)),
+ ('bdc', decimal.Decimal(2)),
+ ('cfgs', decimal.Decimal(3)),
+ ('deeefg', decimal.Decimal(4))],
+ type_info=Types.ROW([Types.STRING(), Types.BIG_DEC()]))
+
+ class MyMapFunction(MapFunction):
+ def map(self, value):
+ return Row(value[0], value[1] + 1, value[2])
+
+ class MyFlatMapFunction(FlatMapFunction):
+ def flat_map(self, value):
+ if value[1] % 2 == 0:
+ yield value
+
+ class MyFilterFunction(FilterFunction):
+ def filter(self, value):
+ return value[1] > 2
+
+ (ds.map(lambda i: (i[0], len(i[0]), i[1]),
+ output_type=Types.TUPLE([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
+ .flat_map(MyFlatMapFunction(),
+ output_type=Types.TUPLE([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
+ .filter(MyFilterFunction())
+ .map(MyMapFunction(),
+ output_type=Types.ROW([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
+ .add_sink(self.test_sink))
+ self.env.execute('test_basic_operations')
+ results = self.test_sink.get_results()
+ expected = ["+I[cfgs, 5, 3]",
+ "+I[deeefg, 7, 4]"]
+ self.assert_equals_sorted(expected, results)
+
+ def test_basic_operations_with_callable_function(self):
ds = self.env.from_collection([('ab', decimal.Decimal(1)),
('bdc', decimal.Decimal(2)),
('cfgs', decimal.Decimal(3)),
('deeefg', decimal.Decimal(4))],
type_info=Types.ROW([Types.STRING(), Types.BIG_DEC()]))
- ds.map(MyMapFunction()).add_sink(self.test_sink)
- self.env.execute('map_function_test')
+
+ def add_1(value):
+ return Row(value[0], value[1] + 1, value[2])
+
+ def reserve_even(value):
+ if value[1] % 2 == 0:
+ yield value
+
+ (ds.map(lambda i: (i[0], len(i[0]), i[1]))
+ .flat_map(reserve_even)
+ .filter(lambda i: i[1] > 2)
+ .map(add_1)
+ .add_sink(self.test_sink))
+ self.env.execute('test_basic_operations_with_callable_function')
results = self.test_sink.get_results(True)
- expected = ["<Row('ab', 2, Decimal('1'))>", "<Row('bdc', 3, Decimal('2'))>",
- "<Row('cfgs', 4, Decimal('3'))>", "<Row('deeefg', 6, Decimal('4'))>"]
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
+ expected = ["<Row('cfgs', 5, Decimal('3'))>",
+ "<Row('deeefg', 7, Decimal('4'))>"]
+ self.assert_equals_sorted(expected, results)
- def test_map_with_output_type(self):
- ds = self.env.from_collection([('ab', 1), ('bdc', 2), ('cfgs', 3), ('deeefg', 4)],
- type_info=Types.ROW([Types.STRING(), Types.INT()]))
+ def test_basic_co_operations(self):
+ python_file_dir = os.path.join(self.tempdir, "python_file_dir_" + str(uuid.uuid4()))
+ os.mkdir(python_file_dir)
+ python_file_path = os.path.join(python_file_dir, "test_stream_dependency_manage_lib.py")
+ with open(python_file_path, 'w') as f:
+ f.write("def add_two(a):\n return a + 2")
- ds.map(MyMapFunction(), output_type=Types.ROW([Types.STRING(), Types.INT(), Types.INT()]))\
- .add_sink(self.test_sink)
- self.env.execute('map_function_test')
- results = self.test_sink.get_results(False)
- expected = ['+I[ab, 2, 1]', '+I[bdc, 3, 2]', '+I[cfgs, 4, 3]', '+I[deeefg, 6, 4]']
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
+ class MyCoFlatMapFunction(CoFlatMapFunction):
- def test_map_with_callable_function(self):
- ds = self.env.from_collection([('ab', 1), ('bdc', 2), ('cfgs', 3), ('deeefg', 4)],
- type_info=Types.TUPLE([Types.STRING(), Types.INT()]))
+ def flat_map1(self, value):
+ yield value + 1
- def map_func(value):
- result = Row(value[0], len(value[0]), value[1])
- return result
+ def flat_map2(self, value):
+ yield value - 1
- ds.map(map_func, output_type=Types.ROW([Types.STRING(), Types.INT(), Types.INT()]))\
- .add_sink(self.test_sink)
- self.env.execute('map_function_test')
- results = self.test_sink.get_results(False)
- expected = ['+I[ab, 2, 1]', '+I[bdc, 3, 2]', '+I[cfgs, 4, 3]', '+I[deeefg, 6, 4]']
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
+ class MyCoMapFunction(CoMapFunction):
- def test_co_map(self):
- self.env.set_parallelism(1)
- ds1 = self.env.from_collection([(1, 1), (2, 2), (3, 3)],
- type_info=Types.ROW([Types.INT(), Types.INT()]))
- ds2 = self.env.from_collection([("a", "a"), ("b", "b"), ("c", "c")],
- type_info=Types.ROW([Types.STRING(), Types.STRING()]))
- ds1.connect(ds2).map(MyCoMapFunction()).add_sink(self.test_sink)
- self.env.execute('co_map_function_test')
+ def map1(self, value):
+ from test_stream_dependency_manage_lib import add_two
+ return add_two(value)
+
+ def map2(self, value):
+ return value + 1
+
+ self.env.add_python_file(python_file_path)
+ ds_1 = self.env.from_collection([1, 2, 3, 4, 5])
+ ds_2 = ds_1.map(lambda x: x * 2)
+
+ (ds_1.connect(ds_2).flat_map(MyCoFlatMapFunction())
+ .connect(ds_2).map(MyCoMapFunction())
+ .add_sink(self.test_sink))
+ self.env.execute("test_basic_co_operations")
results = self.test_sink.get_results(True)
- expected = ['2', '3', '4', 'a', 'b', 'c']
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
+ expected = ['4', '5', '6', '7', '8', '3', '5', '7', '9', '11', '3', '5', '7', '9', '11']
+ self.assert_equals_sorted(expected, results)
- def test_co_map_with_output_type(self):
- self.env.set_parallelism(1)
- ds1 = self.env.from_collection([(1, 1), (2, 2), (3, 3)],
- type_info=Types.ROW([Types.INT(), Types.INT()]))
- ds2 = self.env.from_collection([("a", "a"), ("b", "b"), ("c", "c")],
- type_info=Types.ROW([Types.STRING(), Types.STRING()]))
- ds1.connect(ds2).map(MyCoMapFunction(), output_type=Types.STRING()).add_sink(self.test_sink)
- self.env.execute('co_map_function_test')
- results = self.test_sink.get_results(False)
- expected = ['2', '3', '4', 'a', 'b', 'c']
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
+ def test_basic_co_operations_with_output_type(self):
+ class MyCoMapFunction(CoMapFunction):
- def test_keyed_co_process_function(self):
- self.env.set_parallelism(1)
+ def map1(self, value):
+ return value + 2
+
+ def map2(self, value):
+ return value + 1
+
+ class MyCoFlatMapFunction(CoFlatMapFunction):
+
+ def flat_map1(self, value):
+ yield value + 1
+
+ def flat_map2(self, value):
+ yield value - 1
+
+ ds_1 = self.env.from_collection([1, 2, 3, 4, 5])
+ ds_2 = ds_1.map(lambda x: x * 2)
+
+ (ds_1.connect(ds_2).flat_map(MyCoFlatMapFunction(), output_type=Types.INT())
+ .connect(ds_2).map(MyCoMapFunction(), output_type=Types.INT())
+ .add_sink(self.test_sink))
+ self.env.execute("test_basic_co_operations_with_output_type")
+ results = self.test_sink.get_results()
+ expected = ['4', '5', '6', '7', '8', '3', '5', '7', '9', '11', '3', '5', '7', '9', '11']
+ self.assert_equals_sorted(expected, results)
+
+ def test_keyed_co_process(self):
ds1 = self.env.from_collection([("a", 1), ("b", 2), ("c", 3)],
type_info=Types.ROW([Types.STRING(), Types.INT()]))
ds2 = self.env.from_collection([("b", 2), ("c", 3), ("d", 4)],
@@ -145,7 +198,7 @@ class DataStreamTests(object):
.key_by(lambda x: x[0], lambda x: x[0]) \
.process(MyKeyedCoProcessFunction()) \
.add_sink(self.test_sink)
- self.env.execute('co_map_function_test')
+ self.env.execute('test_keyed_co_process_function')
results = self.test_sink.get_results(True)
expected = ["<Row('a', 1)>",
"<Row('b', 1)>",
@@ -154,52 +207,9 @@ class DataStreamTests(object):
"<Row('c', 2)>",
"<Row('d', 1)>",
"<Row('on_timer', 3)>"]
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
-
- def test_connected_streams_with_dependency(self):
- python_file_dir = os.path.join(self.tempdir, "python_file_dir_" + str(uuid.uuid4()))
- os.mkdir(python_file_dir)
- python_file_path = os.path.join(python_file_dir, "test_stream_dependency_manage_lib.py")
- with open(python_file_path, 'w') as f:
- f.write("def add_two(a):\n return a + 2")
+ self.assert_equals_sorted(expected, results)
- class TestCoMapFunction(CoMapFunction):
-
- def map1(self, value):
- from test_stream_dependency_manage_lib import add_two
- return add_two(value)
-
- def map2(self, value):
- return value + 1
-
- self.env.add_python_file(python_file_path)
- ds = self.env.from_collection([1, 2, 3, 4, 5])
- ds_1 = ds.map(lambda x: x * 2)
- ds.connect(ds_1).map(TestCoMapFunction()).add_sink(self.test_sink)
- self.env.execute("test co-map add python file")
- result = self.test_sink.get_results(True)
- expected = ['11', '3', '3', '4', '5', '5', '6', '7', '7', '9']
- result.sort()
- expected.sort()
- self.assertEqual(expected, result)
-
- def test_co_map_function_with_data_types(self):
- self.env.set_parallelism(1)
- ds1 = self.env.from_collection([(1, 1), (2, 2), (3, 3)],
- type_info=Types.ROW([Types.INT(), Types.INT()]))
- ds2 = self.env.from_collection([("a", "a"), ("b", "b"), ("c", "c")],
- type_info=Types.ROW([Types.STRING(), Types.STRING()]))
- ds1.connect(ds2).map(MyCoMapFunction(), output_type=Types.STRING()).add_sink(self.test_sink)
- self.env.execute('co_map_function_test')
- results = self.test_sink.get_results(False)
- expected = ['2', '3', '4', 'a', 'b', 'c']
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
-
- def test_key_by_and_co_map(self):
+ def test_keyed_co_map(self):
ds1 = self.env.from_collection([('a', 0), ('b', 0), ('c', 1), ('d', 1), ('e', 2)],
type_info=Types.ROW([Types.STRING(), Types.INT()])) \
.key_by(MyKeySelector(), key_type=Types.INT())
@@ -238,78 +248,15 @@ class DataStreamTests(object):
.map(AssertKeyCoMapFunction())\
.add_sink(self.test_sink)
- self.env.execute('key_by_test')
+ self.env.execute()
results = self.test_sink.get_results(True)
expected = ["Row(f0='e', f1=2)", "Row(f0='a', f1=0)", "Row(f0='b', f1=0)",
"Row(f0='c', f1=1)", "Row(f0='d', f1=1)", "Row(f0='e', f1=2)",
"Row(f0='a', f1=0)", "Row(f0='b', f1=0)", "Row(f0='c', f1=1)",
"Row(f0='d', f1=1)"]
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
- def test_flat_map(self):
- ds = self.env.from_collection([('a', 0), ('ab', 1), ('bdc', 2), ('cfgs', 3), ('deeefg', 4)],
- type_info=Types.ROW([Types.STRING(), Types.INT()]))
- ds.flat_map(MyFlatMapFunction(), output_type=Types.ROW([Types.STRING(), Types.INT()]))\
- .add_sink(self.test_sink)
-
- self.env.execute('flat_map_test')
- results = self.test_sink.get_results(False)
- expected = ['+I[a, 0]', '+I[bdc, 2]', '+I[deeefg, 4]']
- results.sort()
- expected.sort()
-
- self.assertEqual(expected, results)
-
- def test_flat_map_with_callable_function(self):
- ds = self.env.from_collection([('a', 0), ('ab', 1), ('bdc', 2), ('cfgs', 3), ('deeefg', 4)],
- type_info=Types.ROW([Types.STRING(), Types.INT()]))
-
- def flat_map(value):
- if value[1] % 2 == 0:
- yield value
-
- ds.flat_map(flat_map, output_type=Types.ROW([Types.STRING(), Types.INT()]))\
- .add_sink(self.test_sink)
- self.env.execute('flat_map_test')
- results = self.test_sink.get_results(False)
- expected = ['+I[a, 0]', '+I[bdc, 2]', '+I[deeefg, 4]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
-
- def test_co_flat_map(self):
- self.env.set_parallelism(1)
- ds1 = self.env.from_collection([(1, 1), (2, 2), (3, 3)],
- type_info=Types.ROW([Types.INT(), Types.INT()]))
- ds2 = self.env.from_collection([("a", "a"), ("b", "b"), ("c", "c")],
- type_info=Types.ROW([Types.STRING(), Types.STRING()]))
- ds1.connect(ds2).flat_map(MyCoFlatMapFunction()).add_sink(self.test_sink)
- self.env.execute('co_flat_map_function_test')
- results = self.test_sink.get_results(True)
- expected = ['2', '2', '3', '3', '4', '4', 'b']
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
-
- def test_co_flat_map_with_output_type(self):
- self.env.set_parallelism(1)
- ds1 = self.env.from_collection([(1, 1), (2, 2), (3, 3)],
- type_info=Types.ROW([Types.INT(), Types.INT()]))
- ds2 = self.env.from_collection([("a", "a"), ("b", "b"), ("c", "c")],
- type_info=Types.ROW([Types.STRING(), Types.STRING()]))
- ds1.connect(ds2).flat_map(MyCoFlatMapFunction(), output_type=Types.STRING())\
- .add_sink(self.test_sink)
- self.env.execute('co_flat_map_function_test')
- results = self.test_sink.get_results(False)
- expected = ['2', '2', '3', '3', '4', '4', 'b']
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
-
- def test_key_by_and_co_flat_map(self):
- self.env.set_parallelism(1)
+ def test_keyed_co_flat_map(self):
ds1 = self.env.from_collection([(1, 1), (2, 2), (3, 3)],
type_info=Types.ROW([Types.INT(), Types.INT()]))
ds2 = self.env.from_collection([("a", "a"), ("b", "b"), ("c", "c"), ("a", "a")],
@@ -317,35 +264,10 @@ class DataStreamTests(object):
ds1.connect(ds2).key_by(lambda x: 1, lambda x: 1) \
.flat_map(MyRichCoFlatMapFunction(), output_type=Types.STRING()) \
.add_sink(self.test_sink)
- self.env.execute('co_flat_map_function_test')
+ self.env.execute('test_keyed_co_flat_map')
results = self.test_sink.get_results(False)
expected = ['2', '2', '3', '3', '4', '4', 'a', 'b', 'c']
- expected.sort()
- results.sort()
- self.assertEqual(expected, results)
-
- def test_filter(self):
- ds = self.env.from_collection([(1, 'Hi', 'Hello'), (2, 'Hello', 'Hi')])
- ds.filter(MyFilterFunction()).add_sink(self.test_sink)
- self.env.execute("test filter")
- results = self.test_sink.get_results(True)
- expected = ["(2, 'Hello', 'Hi')"]
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
-
- def test_filter_with_lambda_function(self):
- ds = self.env.from_collection([(1, 'Hi', 'Hello'), (2, 'Hello', 'Hi')],
- type_info=Types.ROW(
- [Types.INT(), Types.STRING(), Types.STRING()])
- )
- ds.filter(lambda x: x[0] % 2 == 0).add_sink(self.test_sink)
- self.env.execute("test filter")
- results = self.test_sink.get_results(False)
- expected = ['+I[2, Hello, Hi]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
def test_execute_and_collect(self):
test_data = ['pyflink', 'datastream', 'execute', 'collect']
@@ -389,14 +311,11 @@ class DataStreamTests(object):
ds = self.env.from_collection(test_data).map(lambda a: a)
with ds.execute_and_collect() as results:
actual = [result for result in results]
- actual.sort()
- expected.sort()
- self.assertEqual(expected, actual)
+ self.assert_equals_sorted(expected, actual)
- def test_key_by_map(self):
+ def test_keyed_map(self):
from pyflink.util.java_utils import get_j_env_configuration
from pyflink.common import Configuration
- self.env.set_parallelism(1)
config = Configuration(
j_configuration=get_j_env_configuration(self.env._j_stream_execution_environment))
config.set_integer("python.fn-execution.bundle.size", 1)
@@ -443,15 +362,13 @@ class DataStreamTests(object):
return state_value
keyed_stream.map(AssertKeyMapFunction()).add_sink(self.test_sink)
- self.env.execute('key_by_test')
+ self.env.execute('test_keyed_map')
results = self.test_sink.get_results(True)
expected = ["Row(f0='e', f1=2)", "Row(f0='a', f1=0)", "Row(f0='b', f1=1)",
"Row(f0='c', f1=0)", "Row(f0='d', f1=1)"]
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
- def test_key_by_flat_map(self):
+ def test_keyed_flat_map(self):
ds = self.env.from_collection([('a', 0), ('b', 0), ('c', 1), ('d', 1), ('e', 2)],
type_info=Types.ROW([Types.STRING(), Types.INT()]))
keyed_stream = ds.key_by(MyKeySelector(), key_type=Types.INT())
@@ -487,15 +404,13 @@ class DataStreamTests(object):
yield value
keyed_stream.flat_map(AssertKeyMapFunction()).add_sink(self.test_sink)
- self.env.execute('key_by_test')
+ self.env.execute('test_keyed_flat_map')
results = self.test_sink.get_results(True)
expected = ["Row(f0='e', f1=2)", "Row(f0='a', f1=0)", "Row(f0='b', f1=0)",
"Row(f0='c', f1=1)", "Row(f0='d', f1=1)"]
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
- def test_key_by_filter(self):
+ def test_keyed_filter(self):
ds = self.env.from_collection([('a', 0), ('b', 0), ('c', 1), ('d', 1), ('e', 2)],
type_info=Types.ROW([Types.STRING(), Types.INT()]))
keyed_stream = ds.key_by(MyKeySelector())
@@ -536,9 +451,7 @@ class DataStreamTests(object):
self.env.execute('key_by_test')
results = self.test_sink.get_results(False)
expected = ['+I[a, 0]', '+I[c, 1]', '+I[e, 2]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
def test_multi_key_by(self):
ds = self.env.from_collection([('a', 0), ('b', 0), ('c', 1), ('d', 1), ('e', 2)],
@@ -549,13 +462,10 @@ class DataStreamTests(object):
self.env.execute("test multi key by")
results = self.test_sink.get_results(False)
expected = ['+I[d, 1]', '+I[c, 1]', '+I[a, 0]', '+I[b, 0]', '+I[e, 2]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
def test_print_without_align_output(self):
# No need to align output typeinfo since we have specified the type info of the DataStream.
- self.env.set_parallelism(1)
ds = self.env.from_collection([('ab', 1), ('bdc', 2), ('cfgs', 3), ('deeefg', 4)],
type_info=Types.ROW([Types.STRING(), Types.INT()]))
ds.print()
@@ -582,9 +492,7 @@ class DataStreamTests(object):
self.env.execute("test primitive array type info")
results = self.test_sink.get_results()
expected = ['+I[1, [1.1, 1.2, 1.3]]', '+I[2, [2.1, 2.2, 2.3]]', '+I[3, [3.1, 3.2, 3.3]]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
def test_basic_array_type_info(self):
ds = self.env.from_collection([(1, [1.1, None, 1.30], [None, 'hi', 'flink']),
@@ -603,9 +511,7 @@ class DataStreamTests(object):
expected = ['+I[1, [1.1, null, 1.3], [null, hi, flink]]',
'+I[2, [null, 2.2, 2.3], [hello, null, flink]]',
'+I[3, [3.1, 3.2, null], [hello, hi, null]]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
def test_object_array_type_info(self):
ds = self.env.from_collection([(1, [1.1, None, 1.30], [None, 'hi', 'flink']),
@@ -624,9 +530,7 @@ class DataStreamTests(object):
expected = ['+I[1, [1.1, null, 1.3], [null, hi, flink]]',
'+I[2, [null, 2.2, 2.3], [hello, null, flink]]',
'+I[3, [3.1, 3.2, null], [hello, hi, null]]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
def test_sql_timestamp_type_info(self):
ds = self.env.from_collection([(datetime.date(2021, 1, 9),
@@ -643,7 +547,7 @@ class DataStreamTests(object):
self.env.execute("test sql timestamp type info")
results = self.test_sink.get_results()
expected = ['+I[2021-01-09, 12:00:00, 2021-01-09 12:00:00.011]']
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
def test_process_function(self):
self.env.set_parallelism(1)
@@ -668,21 +572,18 @@ class DataStreamTests(object):
data_stream.assign_timestamps_and_watermarks(watermark_strategy)\
.process(MyProcessFunction(), output_type=Types.STRING()).add_sink(self.test_sink)
self.env.execute('test process function')
- result = self.test_sink.get_results()
- expected_result = ["current timestamp: 1603708211000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=1, f1='1603708211000')",
- "current timestamp: 1603708224000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=2, f1='1603708224000')",
- "current timestamp: 1603708226000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=3, f1='1603708226000')",
- "current timestamp: 1603708289000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=4, f1='1603708289000')"]
- result.sort()
- expected_result.sort()
- self.assertEqual(expected_result, result)
+ results = self.test_sink.get_results()
+ expected = ["current timestamp: 1603708211000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=1, f1='1603708211000')",
+ "current timestamp: 1603708224000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=2, f1='1603708224000')",
+ "current timestamp: 1603708226000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=3, f1='1603708226000')",
+ "current timestamp: 1603708289000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=4, f1='1603708289000')"]
+ self.assert_equals_sorted(expected, results)
def test_keyed_process_function_with_state(self):
- self.env.set_parallelism(1)
self.env.get_config().set_auto_watermark_interval(2000)
self.env.set_stream_time_characteristic(TimeCharacteristic.EventTime)
data_stream = self.env.from_collection([(1, 'hi', '1603708211000'),
@@ -742,28 +643,26 @@ class DataStreamTests(object):
.process(MyProcessFunction(), output_type=Types.STRING()) \
.add_sink(self.test_sink)
self.env.execute('test time stamp assigner with keyed process function')
- result = self.test_sink.get_results()
- expected_result = ["current key: hi, current value state: None, current list state: [], "
- "current map state: {}, current value: Row(f0=1, f1='hi', "
- "f2='1603708211000')",
- "current key: hello, current value state: None, "
- "current list state: [], current map state: {}, current value: Row(f0=2,"
- " f1='hello', f2='1603708224000')",
- "current key: hi, current value state: 1, current list state: [1], "
- "current map state: {1: hi}, current value: Row(f0=3, f1='hi', "
- "f2='1603708226000')",
- "current key: hello, current value state: 2, current list state: [2], "
- "current map state: {2: hello}, current value: Row(f0=4, f1='hello', "
- "f2='1603708289000')",
- "current key: hi, current value state: 3, current list state: [1, 3], "
- "current map state: {1: hi, 3: hi}, current value: Row(f0=5, f1='hi', "
- "f2='1603708291000')",
- "current key: hello, current value state: 4, current list state: [2, 4],"
- " current map state: {2: hello, 4: hello}, current value: Row(f0=6, "
- "f1='hello', f2='1603708293000')"]
- result.sort()
- expected_result.sort()
- self.assertEqual(expected_result, result)
+ results = self.test_sink.get_results()
+ expected = ["current key: hi, current value state: None, current list state: [], "
+ "current map state: {}, current value: Row(f0=1, f1='hi', "
+ "f2='1603708211000')",
+ "current key: hello, current value state: None, "
+ "current list state: [], current map state: {}, current value: Row(f0=2,"
+ " f1='hello', f2='1603708224000')",
+ "current key: hi, current value state: 1, current list state: [1], "
+ "current map state: {1: hi}, current value: Row(f0=3, f1='hi', "
+ "f2='1603708226000')",
+ "current key: hello, current value state: 2, current list state: [2], "
+ "current map state: {2: hello}, current value: Row(f0=4, f1='hello', "
+ "f2='1603708289000')",
+ "current key: hi, current value state: 3, current list state: [1, 3], "
+ "current map state: {1: hi, 3: hi}, current value: Row(f0=5, f1='hi', "
+ "f2='1603708291000')",
+ "current key: hello, current value state: 4, current list state: [2, 4],"
+ " current map state: {2: hello, 4: hello}, current value: Row(f0=6, "
+ "f1='hello', f2='1603708293000')"]
+ self.assert_equals_sorted(expected, results)
def test_reducing_state(self):
self.env.set_parallelism(2)
@@ -833,11 +732,9 @@ class DataStreamTests(object):
.process(MyProcessFunction(), output_type=Types.TUPLE([Types.INT(), Types.STRING()])) \
.add_sink(self.test_sink)
self.env.execute('test_aggregating_state')
- result = self.test_sink.get_results()
- expected_result = ['(1,hi)', '(2,hello)', '(4,hi)', '(6,hello)', '(9,hi)', '(12,hello)']
- result.sort()
- expected_result.sort()
- self.assertEqual(expected_result, result)
+ results = self.test_sink.get_results()
+ expected = ['(1,hi)', '(2,hello)', '(4,hi)', '(6,hello)', '(9,hi)', '(12,hello)']
+ self.assert_equals_sorted(expected, results)
def test_count_window(self):
self.env.set_parallelism(2)
@@ -850,11 +747,9 @@ class DataStreamTests(object):
.add_sink(self.test_sink)
self.env.execute('test_count_window')
- result = self.test_sink.get_results()
- expected_result = ['(hello,12)', '(hi,9)']
- result.sort()
- expected_result.sort()
- self.assertEqual(expected_result, result)
+ results = self.test_sink.get_results()
+ expected = ['(hello,12)', '(hi,9)']
+ self.assert_equals_sorted(expected, results)
def test_time_window(self):
self.env.set_parallelism(1)
@@ -870,11 +765,9 @@ class DataStreamTests(object):
.add_sink(self.test_sink)
self.env.execute('test_time_window')
- result = self.test_sink.get_results()
- expected_result = ['(hi,1)', '(hi,1)', '(hi,2)', '(hi,3)']
- result.sort()
- expected_result.sort()
- self.assertEqual(expected_result, result)
+ results = self.test_sink.get_results()
+ expected = ['(hi,1)', '(hi,1)', '(hi,2)', '(hi,3)']
+ self.assert_equals_sorted(expected, results)
class StreamingModeDataStreamTests(DataStreamTests, PyFlinkStreamingTestCase):
@@ -1154,21 +1047,19 @@ class StreamingModeDataStreamTests(DataStreamTests, PyFlinkStreamingTestCase):
.key_by(lambda x: x[0], key_type=Types.INT()) \
.process(MyProcessFunction(), output_type=Types.STRING()).add_sink(self.test_sink)
self.env.execute('test time stamp assigner with keyed process function')
- result = self.test_sink.get_results()
+ results = self.test_sink.get_results()
# Because the watermark interval is too long, no watermark was sent before processing these
# data. So all current watermarks are Long.MIN_VALUE.
- expected_result = ["current key: 1, current timestamp: 1603708211000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=1, f1='1603708211000')",
- "current key: 2, current timestamp: 1603708224000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=2, f1='1603708224000')",
- "current key: 3, current timestamp: 1603708226000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=3, f1='1603708226000')",
- "current key: 4, current timestamp: 1603708289000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=4, f1='1603708289000')",
- "on timer: 3"]
- result.sort()
- expected_result.sort()
- self.assertEqual(expected_result, result)
+ expected = ["current key: 1, current timestamp: 1603708211000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=1, f1='1603708211000')",
+ "current key: 2, current timestamp: 1603708224000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=2, f1='1603708224000')",
+ "current key: 3, current timestamp: 1603708226000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=3, f1='1603708226000')",
+ "current key: 4, current timestamp: 1603708289000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=4, f1='1603708289000')",
+ "on timer: 3"]
+ self.assert_equals_sorted(expected, results)
def test_reduce(self):
ds = self.env.from_collection([(1, 'a'), (2, 'a'), (3, 'a'), (4, 'b')],
@@ -1177,11 +1068,9 @@ class StreamingModeDataStreamTests(DataStreamTests, PyFlinkStreamingTestCase):
.reduce(lambda a, b: Row(a[0] + b[0], b[1])) \
.add_sink(self.test_sink)
self.env.execute('reduce_function_test')
- result = self.test_sink.get_results()
+ results = self.test_sink.get_results()
expected = ["+I[1, a]", "+I[3, a]", "+I[6, a]", "+I[4, b]"]
- expected.sort()
- result.sort()
- self.assertEqual(expected, result)
+ self.assert_equals_sorted(expected, results)
def test_reduce_with_state(self):
ds = self.env.from_collection([('a', 0), ('b', 0), ('c', 1), ('d', 1), ('e', 1)],
@@ -1195,9 +1084,7 @@ class StreamingModeDataStreamTests(DataStreamTests, PyFlinkStreamingTestCase):
self.env.execute('key_by_test')
results = self.test_sink.get_results(False)
expected = ['+I[a, 0]', '+I[ab, 0]', '+I[c, 1]', '+I[cd, 1]', '+I[cde, 1]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
def test_function_with_error(self):
ds = self.env.from_collection([('a', 0), ('b', 0), ('c', 1), ('d', 1), ('e', 1)],
@@ -1230,9 +1117,7 @@ class StreamingModeDataStreamTests(DataStreamTests, PyFlinkStreamingTestCase):
self.env.execute("test_data_with_custom_class")
results = self.test_sink.get_results(True)
expected = ['c', 'c', 'b']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
+ self.assert_equals_sorted(expected, results)
class BatchModeDataStreamTests(DataStreamTests, PyFlinkBatchTestCase):
@@ -1271,18 +1156,16 @@ class BatchModeDataStreamTests(DataStreamTests, PyFlinkBatchTestCase):
.key_by(lambda x: x[0], key_type=Types.INT()) \
.process(MyProcessFunction(), output_type=Types.STRING()).add_sink(self.test_sink)
self.env.execute('test time stamp assigner with keyed process function')
- result = self.test_sink.get_results()
- expected_result = ["current key: 1, current timestamp: 1603708211000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=1, f1='1603708211000')",
- "current key: 2, current timestamp: 1603708224000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=2, f1='1603708224000')",
- "current key: 3, current timestamp: 1603708226000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=3, f1='1603708226000')",
- "current key: 4, current timestamp: 1603708289000, current watermark: "
- "-9223372036854775808, current_value: Row(f0=4, f1='1603708289000')"]
- result.sort()
- expected_result.sort()
- self.assertEqual(expected_result, result)
+ results = self.test_sink.get_results()
+ expected = ["current key: 1, current timestamp: 1603708211000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=1, f1='1603708211000')",
+ "current key: 2, current timestamp: 1603708224000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=2, f1='1603708224000')",
+ "current key: 3, current timestamp: 1603708226000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=3, f1='1603708226000')",
+ "current key: 4, current timestamp: 1603708289000, current watermark: "
+ "-9223372036854775808, current_value: Row(f0=4, f1='1603708289000')"]
+ self.assert_equals_sorted(expected, results)
def test_reduce(self):
ds = self.env.from_collection([(1, 'a'), (4, 'b'), (2, 'a'), (3, 'b'), (3, 'a')],
@@ -1291,11 +1174,9 @@ class BatchModeDataStreamTests(DataStreamTests, PyFlinkBatchTestCase):
.reduce(lambda a, b: Row(a[0] + b[0], b[1])) \
.add_sink(self.test_sink)
self.env.execute('reduce_function_test')
- result = self.test_sink.get_results()
+ results = self.test_sink.get_results()
expected = ["+I[6, a]", "+I[7, b]"]
- expected.sort()
- result.sort()
- self.assertEqual(expected, result)
+ self.assert_equals_sorted(expected, results)
def test_reduce_with_state(self):
ds = self.env.from_collection([('a', 0), ('c', 1), ('d', 1), ('b', 0), ('e', 1)],
@@ -1309,23 +1190,7 @@ class BatchModeDataStreamTests(DataStreamTests, PyFlinkBatchTestCase):
self.env.execute('key_by_test')
results = self.test_sink.get_results(False)
expected = ['+I[ab, 0]', '+I[cde, 1]']
- results.sort()
- expected.sort()
- self.assertEqual(expected, results)
-
-
-class MyMapFunction(MapFunction):
-
- def map(self, value):
- result = Row(value[0], len(value[0]), value[1])
- return result
-
-
-class MyFlatMapFunction(FlatMapFunction):
-
- def flat_map(self, value):
- if value[1] % 2 == 0:
- yield value
+ self.assert_equals_sorted(expected, results)
class MyKeySelector(KeySelector):
@@ -1333,32 +1198,6 @@ class MyKeySelector(KeySelector):
return value[1]
-class MyFilterFunction(FilterFunction):
-
- def filter(self, value):
- return value[0] % 2 == 0
-
-
-class MyCoMapFunction(CoMapFunction):
-
- def map1(self, value):
- return str(value[0] + 1)
-
- def map2(self, value):
- return value[0]
-
-
-class MyCoFlatMapFunction(CoFlatMapFunction):
-
- def flat_map1(self, value):
- yield str(value[0] + 1)
- yield str(value[0] + 1)
-
- def flat_map2(self, value):
- if value[0] == 'b':
- yield value[0]
-
-
class MyRichCoFlatMapFunction(CoFlatMapFunction):
def __init__(self):