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/16 02:05:25 UTC
[flink] branch master updated (d6bccd1 -> 8d77f85)
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 d6bccd1 [FLINK-23738][tests] Disable ChangelogBackend in tests
new da83911 [FLINK-23616][python] Support to chain the Python DataStream operators in execute_and_collect
new 8d77f85 [FLINK-23616][python] Support chaining optimization when mixing use of Python Table API and Python DataStream API
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:
flink-python/pyflink/datastream/data_stream.py | 18 +-
.../datastream/stream_execution_environment.py | 14 +-
.../pyflink/datastream/tests/test_connectors.py | 34 +--
flink-python/pyflink/table/table_environment.py | 11 +-
.../table/tests/test_environment_settings.py | 5 +-
.../table/tests/test_table_environment_api.py | 3 +-
flink-python/pyflink/testing/test_case_utils.py | 18 --
flink-python/pyflink/util/java_utils.py | 22 ++
.../chain/PythonOperatorChainingOptimizer.java | 303 ++++++++++++++-------
.../apache/flink/python/util/PythonConfigUtil.java | 83 +++---
.../python/DataStreamPythonFunctionInfo.java | 11 +
.../AbstractDataStreamPythonFunctionOperator.java | 36 +++
.../AbstractOneInputPythonFunctionOperator.java | 5 -
.../operators/python/PythonCoProcessOperator.java | 6 +-
.../python/PythonKeyedCoProcessOperator.java | 6 +-
.../python/PythonKeyedProcessOperator.java | 6 +-
.../operators/python/PythonProcessOperator.java | 25 +-
.../python/ChainingOptimizingExecutor.java | 75 +++++
.../chain/PythonOperatorChainingOptimizerTest.java | 153 +++++++++--
.../flink/python/util/PythonConfigUtilTest.java | 5 +-
.../table/functions/python/PythonFunctionInfo.java | 7 +-
21 files changed, 594 insertions(+), 252 deletions(-)
create mode 100644 flink-python/src/main/java/org/apache/flink/table/executor/python/ChainingOptimizingExecutor.java
[flink] 02/02: [FLINK-23616][python] Support chaining optimization
when mixing use of Python Table API and Python DataStream API
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 8d77f85ff2495e162c874f8b006b5d28ca78f4ca
Author: Dian Fu <di...@apache.org>
AuthorDate: Fri Aug 13 16:35:33 2021 +0800
[FLINK-23616][python] Support chaining optimization when mixing use of Python Table API and Python DataStream API
This closes #16810.
---
.../pyflink/datastream/tests/test_connectors.py | 34 +++++-----
flink-python/pyflink/table/table_environment.py | 11 +++-
.../table/tests/test_environment_settings.py | 5 +-
.../table/tests/test_table_environment_api.py | 3 +-
flink-python/pyflink/testing/test_case_utils.py | 18 ------
.../python/ChainingOptimizingExecutor.java | 75 ++++++++++++++++++++++
.../chain/PythonOperatorChainingOptimizerTest.java | 54 ++++++++++++++++
7 files changed, 160 insertions(+), 40 deletions(-)
diff --git a/flink-python/pyflink/datastream/tests/test_connectors.py b/flink-python/pyflink/datastream/tests/test_connectors.py
index a089bd8..622f21f 100644
--- a/flink-python/pyflink/datastream/tests/test_connectors.py
+++ b/flink-python/pyflink/datastream/tests/test_connectors.py
@@ -28,8 +28,8 @@ from pyflink.datastream.connectors import FlinkKafkaConsumer, FlinkKafkaProducer
from pyflink.datastream.tests.test_util import DataStreamTestSinkFunction
from pyflink.java_gateway import get_gateway
from pyflink.testing.test_case_utils import PyFlinkTestCase, _load_specific_flink_module_jars, \
- get_private_field, invoke_java_object_method
-from pyflink.util.java_utils import load_java_class
+ invoke_java_object_method
+from pyflink.util.java_utils import load_java_class, get_field_value
class FlinkKafkaTest(PyFlinkTestCase):
@@ -61,22 +61,22 @@ class FlinkKafkaTest(PyFlinkTestCase):
flink_kafka_consumer.set_start_from_earliest()
flink_kafka_consumer.set_commit_offsets_on_checkpoints(True)
- j_properties = get_private_field(flink_kafka_consumer.get_java_function(), 'properties')
+ j_properties = get_field_value(flink_kafka_consumer.get_java_function(), 'properties')
self.assertEqual('localhost:9092', j_properties.getProperty('bootstrap.servers'))
self.assertEqual('test_group', j_properties.getProperty('group.id'))
- self.assertTrue(get_private_field(flink_kafka_consumer.get_java_function(),
- 'enableCommitOnCheckpoints'))
- j_start_up_mode = get_private_field(flink_kafka_consumer.get_java_function(), 'startupMode')
+ self.assertTrue(get_field_value(flink_kafka_consumer.get_java_function(),
+ 'enableCommitOnCheckpoints'))
+ j_start_up_mode = get_field_value(flink_kafka_consumer.get_java_function(), 'startupMode')
- j_deserializer = get_private_field(flink_kafka_consumer.get_java_function(), 'deserializer')
+ j_deserializer = get_field_value(flink_kafka_consumer.get_java_function(), 'deserializer')
j_deserialize_type_info = invoke_java_object_method(j_deserializer, "getProducedType")
deserialize_type_info = typeinfo._from_java_type(j_deserialize_type_info)
self.assertTrue(deserialize_type_info == type_info)
self.assertTrue(j_start_up_mode.equals(get_gateway().jvm
.org.apache.flink.streaming.connectors
.kafka.config.StartupMode.EARLIEST))
- j_topic_desc = get_private_field(flink_kafka_consumer.get_java_function(),
- 'topicsDescriptor')
+ j_topic_desc = get_field_value(flink_kafka_consumer.get_java_function(),
+ 'topicsDescriptor')
j_topics = invoke_java_object_method(j_topic_desc, 'getFixedTopics')
self.assertEqual(['test_source_topic'], list(j_topics))
@@ -86,12 +86,12 @@ class FlinkKafkaTest(PyFlinkTestCase):
flink_kafka_producer = flink_kafka_producer_clz(sink_topic, serialization_schema, props)
flink_kafka_producer.set_write_timestamp_to_kafka(False)
- j_producer_config = get_private_field(flink_kafka_producer.get_java_function(),
- 'producerConfig')
+ j_producer_config = get_field_value(flink_kafka_producer.get_java_function(),
+ 'producerConfig')
self.assertEqual('localhost:9092', j_producer_config.getProperty('bootstrap.servers'))
self.assertEqual('test_group', j_producer_config.getProperty('group.id'))
- self.assertFalse(get_private_field(flink_kafka_producer.get_java_function(),
- 'writeTimestampToKafka'))
+ self.assertFalse(get_field_value(flink_kafka_producer.get_java_function(),
+ 'writeTimestampToKafka'))
def tearDown(self):
# Change the ClassLoader back to the cached ContextClassLoader after the test case finished.
@@ -123,11 +123,11 @@ class FlinkJdbcSinkTest(PyFlinkTestCase):
ds.add_sink(jdbc_sink).name('jdbc sink')
plan = eval(self.env.get_execution_plan())
self.assertEqual('Sink: jdbc sink', plan['nodes'][1]['type'])
- j_output_format = get_private_field(jdbc_sink.get_java_function(), 'outputFormat')
+ j_output_format = get_field_value(jdbc_sink.get_java_function(), 'outputFormat')
connection_options = JdbcConnectionOptions(
- get_private_field(get_private_field(j_output_format, 'connectionProvider'),
- 'jdbcOptions'))
+ get_field_value(get_field_value(j_output_format, 'connectionProvider'),
+ 'jdbcOptions'))
self.assertEqual(jdbc_connection_options.get_db_url(), connection_options.get_db_url())
self.assertEqual(jdbc_connection_options.get_driver_name(),
connection_options.get_driver_name())
@@ -135,7 +135,7 @@ class FlinkJdbcSinkTest(PyFlinkTestCase):
self.assertEqual(jdbc_connection_options.get_user_name(),
connection_options.get_user_name())
- exec_options = JdbcExecutionOptions(get_private_field(j_output_format, 'executionOptions'))
+ exec_options = JdbcExecutionOptions(get_field_value(j_output_format, 'executionOptions'))
self.assertEqual(jdbc_execution_options.get_batch_interval_ms(),
exec_options.get_batch_interval_ms())
self.assertEqual(jdbc_execution_options.get_batch_size(),
diff --git a/flink-python/pyflink/table/table_environment.py b/flink-python/pyflink/table/table_environment.py
index 88afa5c..bd6c7b4 100644
--- a/flink-python/pyflink/table/table_environment.py
+++ b/flink-python/pyflink/table/table_environment.py
@@ -48,7 +48,7 @@ from pyflink.table.udf import UserDefinedFunctionWrapper, AggregateFunction, uda
from pyflink.table.utils import to_expression_jarray
from pyflink.util import java_utils
from pyflink.util.java_utils import get_j_env_configuration, is_local_deployment, load_java_class, \
- to_j_explain_detail_arr, to_jarray
+ to_j_explain_detail_arr, to_jarray, get_field
__all__ = [
'StreamTableEnvironment',
@@ -95,6 +95,7 @@ class TableEnvironment(object):
# specified by sys.executable if users have not specified it explicitly via configuration
# python.executable.
self._set_python_executable_for_local_executor()
+ self._config_chaining_optimization()
@staticmethod
def create(environment_settings: EnvironmentSettings) -> 'TableEnvironment':
@@ -1757,12 +1758,18 @@ class TableEnvironment(object):
name=str(function.__class__.__name__))
return function
+ def _config_chaining_optimization(self):
+ JChainingOptimizingExecutor = get_gateway().jvm.org.apache.flink.table.executor.python.\
+ ChainingOptimizingExecutor
+ exec_env_field = get_field(self._j_tenv.getClass(), "execEnv")
+ exec_env_field.set(self._j_tenv,
+ JChainingOptimizingExecutor(exec_env_field.get(self._j_tenv)))
+
class StreamTableEnvironment(TableEnvironment):
def __init__(self, j_tenv):
super(StreamTableEnvironment, self).__init__(j_tenv)
- self._j_tenv = j_tenv
@staticmethod
def create(stream_execution_environment: StreamExecutionEnvironment = None, # type: ignore
diff --git a/flink-python/pyflink/table/tests/test_environment_settings.py b/flink-python/pyflink/table/tests/test_environment_settings.py
index 525755f..5f74bb7 100644
--- a/flink-python/pyflink/table/tests/test_environment_settings.py
+++ b/flink-python/pyflink/table/tests/test_environment_settings.py
@@ -19,7 +19,8 @@ from pyflink.java_gateway import get_gateway
from pyflink.common import Configuration
from pyflink.table import EnvironmentSettings
-from pyflink.testing.test_case_utils import PyFlinkTestCase, get_private_field
+from pyflink.testing.test_case_utils import PyFlinkTestCase
+from pyflink.util.java_utils import get_field_value
class EnvironmentSettingsTests(PyFlinkTestCase):
@@ -120,7 +121,7 @@ class EnvironmentSettingsTests(PyFlinkTestCase):
CLASS_NAME = gateway.jvm.EnvironmentSettings.CLASS_NAME
builder = EnvironmentSettings.new_instance()
- BLINK_PLANNER_FACTORY = get_private_field(builder._j_builder, "BLINK_PLANNER_FACTORY")
+ BLINK_PLANNER_FACTORY = get_field_value(builder._j_builder, "BLINK_PLANNER_FACTORY")
self.assertEqual(
settings._j_environment_settings.toPlannerProperties()[CLASS_NAME],
diff --git a/flink-python/pyflink/table/tests/test_table_environment_api.py b/flink-python/pyflink/table/tests/test_table_environment_api.py
index 523b31e..4c5dfe6 100644
--- a/flink-python/pyflink/table/tests/test_table_environment_api.py
+++ b/flink-python/pyflink/table/tests/test_table_environment_api.py
@@ -295,7 +295,8 @@ class DataStreamConversionTestCases(PyFlinkTestCase):
expected = ['+I[1, Hi, Hello]', '+I[2, Hello, Hi]']
self.assert_equals(result, expected)
- ds = ds.map(lambda x: x, Types.ROW([Types.INT(), Types.STRING(), Types.STRING()]))
+ ds = ds.map(lambda x: x, Types.ROW([Types.INT(), Types.STRING(), Types.STRING()])) \
+ .map(lambda x: x, Types.ROW([Types.INT(), Types.STRING(), Types.STRING()]))
table = t_env.from_data_stream(ds, col('a'), col('b'), col('c'))
t_env.register_table_sink("ExprSink",
source_sink_utils.TestAppendSink(field_names, field_types))
diff --git a/flink-python/pyflink/testing/test_case_utils.py b/flink-python/pyflink/testing/test_case_utils.py
index 6777f60..ddb5a6b 100644
--- a/flink-python/pyflink/testing/test_case_utils.py
+++ b/flink-python/pyflink/testing/test_case_utils.py
@@ -26,7 +26,6 @@ import unittest
from abc import abstractmethod
from py4j.java_gateway import JavaObject
-from py4j.protocol import Py4JJavaError
from pyflink.common import JobExecutionResult
from pyflink.datastream.execution_mode import RuntimeExecutionMode
@@ -46,23 +45,6 @@ logging.basicConfig(stream=sys.stdout, level=log_level,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
-def get_private_field(java_obj, field_name):
- try:
- field = java_obj.getClass().getDeclaredField(field_name)
- field.setAccessible(True)
- return field.get(java_obj)
- except Py4JJavaError:
- cls = java_obj.getClass()
- while cls.getSuperclass() is not None:
- cls = cls.getSuperclass()
- try:
- field = cls.getDeclaredField(field_name)
- field.setAccessible(True)
- return field.get(java_obj)
- except Py4JJavaError:
- pass
-
-
def exec_insert_table(table, table_path) -> JobExecutionResult:
return table.execute_insert(table_path).get_job_client().get_job_execution_result().result()
diff --git a/flink-python/src/main/java/org/apache/flink/table/executor/python/ChainingOptimizingExecutor.java b/flink-python/src/main/java/org/apache/flink/table/executor/python/ChainingOptimizingExecutor.java
new file mode 100644
index 0000000..be26b92
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/table/executor/python/ChainingOptimizingExecutor.java
@@ -0,0 +1,75 @@
+/*
+ * 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.table.executor.python;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.JobExecutionResult;
+import org.apache.flink.api.dag.Pipeline;
+import org.apache.flink.api.dag.Transformation;
+import org.apache.flink.configuration.ReadableConfig;
+import org.apache.flink.core.execution.JobClient;
+import org.apache.flink.python.PythonOptions;
+import org.apache.flink.python.chain.PythonOperatorChainingOptimizer;
+import org.apache.flink.python.util.PythonConfigUtil;
+import org.apache.flink.table.delegation.Executor;
+import org.apache.flink.util.Preconditions;
+
+import java.util.List;
+
+/** {@link Executor} which will perform chaining optimization before generating the StreamGraph. */
+@Internal
+public class ChainingOptimizingExecutor implements Executor {
+
+ private final Executor executor;
+
+ public ChainingOptimizingExecutor(Executor executor) {
+ this.executor = Preconditions.checkNotNull(executor);
+ }
+
+ @Override
+ public ReadableConfig getConfiguration() {
+ return executor.getConfiguration();
+ }
+
+ @Override
+ public Pipeline createPipeline(
+ List<Transformation<?>> transformations,
+ ReadableConfig configuration,
+ String defaultJobName) {
+ List<Transformation<?>> chainedTransformations = transformations;
+ if (configuration
+ .getOptional(PythonOptions.PYTHON_OPERATOR_CHAINING_ENABLED)
+ .orElse(getConfiguration().get(PythonOptions.PYTHON_OPERATOR_CHAINING_ENABLED))) {
+ chainedTransformations = PythonOperatorChainingOptimizer.optimize(transformations);
+ }
+
+ PythonConfigUtil.setPartitionCustomOperatorNumPartitions(chainedTransformations);
+ return executor.createPipeline(chainedTransformations, configuration, defaultJobName);
+ }
+
+ @Override
+ public JobExecutionResult execute(Pipeline pipeline) throws Exception {
+ return executor.execute(pipeline);
+ }
+
+ @Override
+ public JobClient executeAsync(Pipeline pipeline) throws Exception {
+ return executor.executeAsync(pipeline);
+ }
+}
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
index a693316..2e24bc3 100644
--- 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
@@ -534,6 +534,60 @@ public class PythonOperatorChainingOptimizerTest {
"f1");
}
+ @Test
+ public void testSingleTransformation() {
+ 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(processTransformation2);
+
+ List<Transformation<?>> optimized =
+ PythonOperatorChainingOptimizer.optimize(transformations);
+ assertEquals(2, optimized.size());
+
+ OneInputTransformation<?, ?> chainedTransformation =
+ (OneInputTransformation<?, ?>) optimized.get(0);
+ 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");
+ }
+
// ----------------------- Utility Methods -----------------------
private void validateChainedPythonFunctions(
[flink] 01/02: [FLINK-23616][python] Support to chain the Python
DataStream operators in execute_and_collect
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 da83911087bb0099cb355445dfedd7bd2b321244
Author: Dian Fu <di...@apache.org>
AuthorDate: Thu Aug 12 13:17:02 2021 +0800
[FLINK-23616][python] Support to chain the Python DataStream operators in execute_and_collect
This closes #16804.
---
flink-python/pyflink/datastream/data_stream.py | 18 +-
.../datastream/stream_execution_environment.py | 14 +-
flink-python/pyflink/util/java_utils.py | 22 ++
.../chain/PythonOperatorChainingOptimizer.java | 303 ++++++++++++++-------
.../apache/flink/python/util/PythonConfigUtil.java | 83 +++---
.../python/DataStreamPythonFunctionInfo.java | 11 +
.../AbstractDataStreamPythonFunctionOperator.java | 36 +++
.../AbstractOneInputPythonFunctionOperator.java | 5 -
.../operators/python/PythonCoProcessOperator.java | 6 +-
.../python/PythonKeyedCoProcessOperator.java | 6 +-
.../python/PythonKeyedProcessOperator.java | 6 +-
.../operators/python/PythonProcessOperator.java | 25 +-
.../chain/PythonOperatorChainingOptimizerTest.java | 99 +++++--
.../flink/python/util/PythonConfigUtilTest.java | 5 +-
.../table/functions/python/PythonFunctionInfo.java | 7 +-
15 files changed, 434 insertions(+), 212 deletions(-)
diff --git a/flink-python/pyflink/datastream/data_stream.py b/flink-python/pyflink/datastream/data_stream.py
index d1156f2..bd2e3ec 100644
--- a/flink-python/pyflink/datastream/data_stream.py
+++ b/flink-python/pyflink/datastream/data_stream.py
@@ -604,7 +604,7 @@ class DataStream(object):
This method takes the key selector to get the key to partition on, and a partitioner that
accepts the key type.
- Note that this method works only on single field keys, i.e. the selector cannet return
+ Note that this method works only on single field keys, i.e. the selector cannot return
tuples of fields.
:param partitioner: The partitioner to assign partitions to keys.
@@ -675,6 +675,8 @@ class DataStream(object):
stream_with_partition_info = self.process(
CustomPartitioner(partitioner, key_selector),
output_type=Types.ROW([Types.INT(), original_type_info]))
+ stream_with_partition_info._j_data_stream.getTransformation().getOperatorFactory() \
+ .getOperator().setContainsPartitionCustom(True)
stream_with_partition_info.name(
gateway.jvm.org.apache.flink.python.util.PythonConfigUtil
@@ -731,6 +733,7 @@ class DataStream(object):
"""
JPythonConfigUtil = get_gateway().jvm.org.apache.flink.python.util.PythonConfigUtil
JPythonConfigUtil.configPythonOperator(self._j_data_stream.getExecutionEnvironment())
+ self._apply_chaining_optimization()
if job_execution_name is None and limit is None:
return CloseableIterator(self._j_data_stream.executeAndCollect(), self.get_type())
elif job_execution_name is not None and limit is None:
@@ -760,6 +763,19 @@ class DataStream(object):
j_data_stream_sink = self._align_output_type()._j_data_stream.print()
return DataStreamSink(j_data_stream_sink)
+ def _apply_chaining_optimization(self):
+ """
+ Chain the Python operators if possible.
+ """
+ gateway = get_gateway()
+ JPythonOperatorChainingOptimizer = gateway.jvm.org.apache.flink.python.chain. \
+ PythonOperatorChainingOptimizer
+ j_transformation = JPythonOperatorChainingOptimizer.apply(
+ self._j_data_stream.getExecutionEnvironment(),
+ self._j_data_stream.getTransformation())
+ self._j_data_stream = gateway.jvm.org.apache.flink.streaming.api.datastream.DataStream(
+ self._j_data_stream.getExecutionEnvironment(), j_transformation)
+
def _align_output_type(self) -> 'DataStream':
"""
Transform the pickled python object into String if the output type is PickledByteArrayInfo.
diff --git a/flink-python/pyflink/datastream/stream_execution_environment.py b/flink-python/pyflink/datastream/stream_execution_environment.py
index 7f3cda7..a7ba4c1 100644
--- a/flink-python/pyflink/datastream/stream_execution_environment.py
+++ b/flink-python/pyflink/datastream/stream_execution_environment.py
@@ -39,7 +39,8 @@ from pyflink.datastream.state_backend import _from_j_state_backend, StateBackend
from pyflink.datastream.time_characteristic import TimeCharacteristic
from pyflink.java_gateway import get_gateway
from pyflink.serializers import PickleSerializer
-from pyflink.util.java_utils import load_java_class, add_jars_to_context_class_loader, invoke_method
+from pyflink.util.java_utils import load_java_class, add_jars_to_context_class_loader, \
+ invoke_method, get_field_value
__all__ = ['StreamExecutionEnvironment']
@@ -867,15 +868,18 @@ class StreamExecutionEnvironment(object):
def _generate_stream_graph(self, clear_transformations: bool = False, job_name: str = None) \
-> JavaObject:
gateway = get_gateway()
+ JPythonConfigUtil = gateway.jvm.org.apache.flink.python.util.PythonConfigUtil
+ JPythonConfigUtil.configPythonOperator(self._j_stream_execution_environment)
+
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)
+ JPythonConfigUtil.setPartitionCustomOperatorNumPartitions(
+ get_field_value(self._j_stream_execution_environment, "transformations"))
+
+ j_stream_graph = self._j_stream_execution_environment.getStreamGraph(clear_transformations)
if job_name is not None:
j_stream_graph.setJobName(job_name)
-
return j_stream_graph
def is_unaligned_checkpoints_enabled(self):
diff --git a/flink-python/pyflink/util/java_utils.py b/flink-python/pyflink/util/java_utils.py
index 8a3045d..1ba6a3e 100644
--- a/flink-python/pyflink/util/java_utils.py
+++ b/flink-python/pyflink/util/java_utils.py
@@ -19,6 +19,7 @@
from datetime import timedelta
from py4j.java_gateway import JavaClass, get_java_class, JavaObject
+from py4j.protocol import Py4JJavaError
from pyflink.java_gateway import get_gateway
@@ -91,6 +92,27 @@ def get_j_env_configuration(j_env):
return field.get(j_env)
+def get_field_value(java_obj, field_name):
+ field = get_field(java_obj.getClass(), field_name)
+ return field.get(java_obj)
+
+
+def get_field(cls, field_name):
+ try:
+ field = cls.getDeclaredField(field_name)
+ field.setAccessible(True)
+ return field
+ except Py4JJavaError:
+ while cls.getSuperclass() is not None:
+ cls = cls.getSuperclass()
+ try:
+ field = cls.getDeclaredField(field_name)
+ field.setAccessible(True)
+ return field
+ except Py4JJavaError:
+ pass
+
+
def invoke_method(obj, object_type, method_name, args=None, arg_types=None):
env_clazz = load_java_class(object_type)
method = env_clazz.getDeclaredMethod(
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
index ebfdc85..4283c7c 100644
--- 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
@@ -18,8 +18,8 @@
package org.apache.flink.python.chain;
-import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.dag.Transformation;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.core.memory.ManagedMemoryUseCase;
import org.apache.flink.python.PythonOptions;
import org.apache.flink.python.util.PythonConfigUtil;
@@ -51,21 +51,22 @@ 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.Lists;
+import org.apache.flink.shaded.guava30.com.google.common.collect.Queues;
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.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
+import java.util.Queue;
import java.util.Set;
-import static org.apache.flink.util.Preconditions.checkNotNull;
-
/**
* A util class which attempts to chain all available Python functions.
*
@@ -91,47 +92,135 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
*/
public class PythonOperatorChainingOptimizer {
- private final List<Transformation<?>> transformations;
-
- private final boolean isChainingEnabled;
+ /**
+ * Perform chaining optimization. It will iterate the transformations defined in the given
+ * StreamExecutionEnvironment and update them with the chained transformations.
+ */
+ @SuppressWarnings("unchecked")
+ public static void apply(StreamExecutionEnvironment env) throws Exception {
+ if (env.getConfiguration().get(PythonOptions.PYTHON_OPERATOR_CHAINING_ENABLED)) {
+ final Field transformationsField =
+ StreamExecutionEnvironment.class.getDeclaredField("transformations");
+ transformationsField.setAccessible(true);
+ final List<Transformation<?>> transformations =
+ (List<Transformation<?>>) transformationsField.get(env);
+ transformationsField.set(env, optimize(transformations));
+ }
+ }
- // key-value pairs where the value is the output transformations of the key transformation
- private final transient Map<Transformation<?>, Set<Transformation<?>>> outputMap =
- new HashMap<>();
+ /**
+ * Perform chaining optimization. It will iterate the transformations defined in the given
+ * StreamExecutionEnvironment and update them with the chained transformations. Besides, it will
+ * return the transformation after chaining optimization for the given transformation.
+ */
+ @SuppressWarnings("unchecked")
+ public static Transformation<?> apply(
+ StreamExecutionEnvironment env, Transformation<?> transformation) throws Exception {
+ if (env.getConfiguration().get(PythonOptions.PYTHON_OPERATOR_CHAINING_ENABLED)) {
+ final Field transformationsField =
+ StreamExecutionEnvironment.class.getDeclaredField("transformations");
+ transformationsField.setAccessible(true);
+ final List<Transformation<?>> transformations =
+ (List<Transformation<?>>) transformationsField.get(env);
+ final Tuple2<List<Transformation<?>>, Transformation<?>> resultTuple =
+ optimize(transformations, transformation);
+ transformationsField.set(env, resultTuple.f0);
+ return resultTuple.f1;
+ } else {
+ return transformation;
+ }
+ }
- public PythonOperatorChainingOptimizer(
- List<Transformation<?>> transformations, boolean isChainingEnabled) {
- this.transformations = checkNotNull(transformations);
- this.isChainingEnabled = isChainingEnabled;
- buildOutputMap();
+ /**
+ * Perform chaining optimization. It will return the chained transformations for the given
+ * transformation list.
+ */
+ public static List<Transformation<?>> optimize(List<Transformation<?>> transformations) {
+ final Map<Transformation<?>, Set<Transformation<?>>> outputMap =
+ buildOutputMap(transformations);
+
+ final LinkedHashSet<Transformation<?>> chainedTransformations = new LinkedHashSet<>();
+ final Set<Transformation<?>> alreadyTransformed = Sets.newIdentityHashSet();
+ final Queue<Transformation<?>> toTransformQueue = Queues.newArrayDeque(transformations);
+ while (!toTransformQueue.isEmpty()) {
+ final Transformation<?> transformation = toTransformQueue.poll();
+ if (!alreadyTransformed.contains(transformation)) {
+ alreadyTransformed.add(transformation);
+
+ final ChainInfo chainInfo = chainWithInputIfPossible(transformation, outputMap);
+ chainedTransformations.add(chainInfo.newTransformation);
+ chainedTransformations.removeAll(chainInfo.oldTransformations);
+ alreadyTransformed.addAll(chainInfo.oldTransformations);
+
+ // Add the chained transformation and its inputs to the to-optimize list
+ toTransformQueue.add(chainInfo.newTransformation);
+ toTransformQueue.addAll(chainInfo.newTransformation.getInputs());
+ }
+ }
+ return new ArrayList<>(chainedTransformations);
}
- private void buildOutputMap() {
- for (Transformation<?> transformation : transformations) {
- for (Transformation<?> input : transformation.getInputs()) {
- Set<Transformation<?>> outputs =
- outputMap.computeIfAbsent(input, i -> Sets.newHashSet());
- outputs.add(transformation);
+ /**
+ * Perform chaining optimization. It will returns the chained transformations and the
+ * transformation after chaining optimization for the given transformation.
+ */
+ public static Tuple2<List<Transformation<?>>, Transformation<?>> optimize(
+ List<Transformation<?>> transformations, Transformation<?> targetTransformation) {
+ final Map<Transformation<?>, Set<Transformation<?>>> outputMap =
+ buildOutputMap(transformations);
+
+ final LinkedHashSet<Transformation<?>> chainedTransformations = new LinkedHashSet<>();
+ final Set<Transformation<?>> alreadyTransformed = Sets.newIdentityHashSet();
+ final Queue<Transformation<?>> toTransformQueue = Queues.newArrayDeque();
+ toTransformQueue.add(targetTransformation);
+ while (!toTransformQueue.isEmpty()) {
+ final Transformation<?> toTransform = toTransformQueue.poll();
+ if (!alreadyTransformed.contains(toTransform)) {
+ alreadyTransformed.add(toTransform);
+
+ final ChainInfo chainInfo = chainWithInputIfPossible(toTransform, outputMap);
+ chainedTransformations.add(chainInfo.newTransformation);
+ chainedTransformations.removeAll(chainInfo.oldTransformations);
+ alreadyTransformed.addAll(chainInfo.oldTransformations);
+
+ // Add the chained transformation and its inputs to the to-optimize list
+ toTransformQueue.add(chainInfo.newTransformation);
+ toTransformQueue.addAll(chainInfo.newTransformation.getInputs());
+
+ if (toTransform == targetTransformation) {
+ targetTransformation = chainInfo.newTransformation;
+ }
}
}
+ return Tuple2.of(new ArrayList<>(chainedTransformations), targetTransformation);
}
- @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);
+ /**
+ * Construct the key-value pairs where the value is the output transformations of the key
+ * transformation.
+ */
+ private static Map<Transformation<?>, Set<Transformation<?>>> buildOutputMap(
+ List<Transformation<?>> transformations) {
+ final Map<Transformation<?>, Set<Transformation<?>>> outputMap = new HashMap<>();
+ final Queue<Transformation<?>> toTransformQueue = Queues.newArrayDeque(transformations);
+ final Set<Transformation<?>> alreadyTransformed = Sets.newIdentityHashSet();
+ while (!toTransformQueue.isEmpty()) {
+ Transformation<?> transformation = toTransformQueue.poll();
+ if (!alreadyTransformed.contains(transformation)) {
+ alreadyTransformed.add(transformation);
+ for (Transformation<?> input : transformation.getInputs()) {
+ Set<Transformation<?>> outputs =
+ outputMap.computeIfAbsent(input, i -> Sets.newHashSet());
+ outputs.add(transformation);
+ }
+ toTransformQueue.addAll(transformation.getInputs());
}
}
- return optimizedTransformations;
+ return outputMap;
}
- private ChainInfo chainWithInputIfPossible(Transformation<?> transform) throws Exception {
+ private static ChainInfo chainWithInputIfPossible(
+ Transformation<?> transform, Map<Transformation<?>, Set<Transformation<?>>> outputMap) {
ChainInfo chainInfo = null;
if (transform instanceof OneInputTransformation
&& PythonConfigUtil.isPythonDataStreamOperator(transform)) {
@@ -166,27 +255,36 @@ public class PythonOperatorChainingOptimizer {
return chainInfo;
}
- private Transformation<?> createChainedTransformation(
+ private static Transformation<?> createChainedTransformation(
Transformation<?> upTransform, Transformation<?> downTransform) {
- AbstractDataStreamPythonFunctionOperator<?> upOperator =
+ final AbstractDataStreamPythonFunctionOperator<?> upOperator =
(AbstractDataStreamPythonFunctionOperator<?>)
((SimpleOperatorFactory<?>) getOperatorFactory(upTransform)).getOperator();
- PythonProcessOperator<?, ?> downOperator =
+ final PythonProcessOperator<?, ?> downOperator =
(PythonProcessOperator<?, ?>)
((SimpleOperatorFactory<?>) getOperatorFactory(downTransform))
.getOperator();
- DataStreamPythonFunctionInfo upPythonFunctionInfo = upOperator.getPythonFunctionInfo();
- DataStreamPythonFunctionInfo downPythonFunctionInfo = downOperator.getPythonFunctionInfo();
+ final DataStreamPythonFunctionInfo upPythonFunctionInfo =
+ upOperator.getPythonFunctionInfo().copy();
+ final DataStreamPythonFunctionInfo downPythonFunctionInfo =
+ downOperator.getPythonFunctionInfo().copy();
+
+ DataStreamPythonFunctionInfo headPythonFunctionInfoOfDownOperator = downPythonFunctionInfo;
+ while (headPythonFunctionInfoOfDownOperator.getInputs().length != 0) {
+ headPythonFunctionInfoOfDownOperator =
+ (DataStreamPythonFunctionInfo)
+ headPythonFunctionInfoOfDownOperator.getInputs()[0];
+ }
+ headPythonFunctionInfoOfDownOperator.setInputs(
+ new DataStreamPythonFunctionInfo[] {upPythonFunctionInfo});
- DataStreamPythonFunctionInfo chainedPythonFunctionInfo =
- new DataStreamPythonFunctionInfo(
- downPythonFunctionInfo.getPythonFunction(),
- upPythonFunctionInfo,
- downPythonFunctionInfo.getFunctionType());
+ final AbstractDataStreamPythonFunctionOperator<?> chainedOperator =
+ upOperator.copy(downPythonFunctionInfo, downOperator.getProducedType());
- AbstractDataStreamPythonFunctionOperator<?> chainedOperator =
- upOperator.copy(chainedPythonFunctionInfo, downOperator.getProducedType());
+ // set partition custom flag
+ chainedOperator.setContainsPartitionCustom(
+ downOperator.containsPartitionCustom() || upOperator.containsPartitionCustom());
PhysicalTransformation<?> chainedTransformation;
if (upOperator instanceof AbstractOneInputPythonFunctionOperator) {
@@ -269,27 +367,12 @@ public class PythonOperatorChainingOptimizer {
return chainedTransformation;
}
- private boolean isChainable(Transformation<?> upTransform, Transformation<?> downTransform) {
+ private static 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);
- PythonOperatorChainingOptimizer chainingOptimizer =
- new PythonOperatorChainingOptimizer(
- transformations,
- env.getConfiguration().get(PythonOptions.PYTHON_OPERATOR_CHAINING_ENABLED));
- transformationsField.set(env, chainingOptimizer.optimize());
+ && areOperatorsChainable(upTransform, downTransform);
}
private static boolean areOperatorsChainable(
@@ -298,10 +381,10 @@ public class PythonOperatorChainingOptimizer {
return false;
}
- AbstractDataStreamPythonFunctionOperator<?> upOperator =
+ final AbstractDataStreamPythonFunctionOperator<?> upOperator =
(AbstractDataStreamPythonFunctionOperator<?>)
((SimpleOperatorFactory<?>) getOperatorFactory(upTransform)).getOperator();
- AbstractDataStreamPythonFunctionOperator<?> downOperator =
+ final AbstractDataStreamPythonFunctionOperator<?> downOperator =
(AbstractDataStreamPythonFunctionOperator<?>)
((SimpleOperatorFactory<?>) getOperatorFactory(downTransform))
.getOperator();
@@ -356,6 +439,8 @@ public class PythonOperatorChainingOptimizer {
return isChainable;
}
+ // ----------------------- Utility Methods -----------------------
+
private static StreamOperatorFactory<?> getOperatorFactory(Transformation<?> transform) {
if (transform instanceof OneInputTransformation) {
return ((OneInputTransformation<?, ?>) transform).getOperatorFactory();
@@ -371,48 +456,54 @@ public class PythonOperatorChainingOptimizer {
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");
+ Transformation<?> newInput) {
+ try {
+ if (transformation instanceof OneInputTransformation
+ || transformation instanceof FeedbackTransformation
+ || transformation instanceof SideOutputTransformation
+ || transformation instanceof ReduceTransformation
+ || transformation instanceof SinkTransformation
+ || transformation instanceof LegacySinkTransformation
+ || transformation instanceof TimestampsAndWatermarksTransformation
+ || transformation instanceof PartitionTransformation) {
+ final Field inputField = transformation.getClass().getDeclaredField("input");
+ inputField.setAccessible(true);
+ inputField.set(transformation, newInput);
+ } else if (transformation instanceof TwoInputTransformation) {
+ final 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) {
+ final 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) {
+ final 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 {
- inputField = transformation.getClass().getDeclaredField("broadcastInput");
+ throw new RuntimeException("Unsupported transformation: " + transformation);
}
- inputField.setAccessible(true);
- inputField.set(transformation, newInput);
- } else {
- throw new RuntimeException("Unsupported transformation: " + transformation);
+ } catch (NoSuchFieldException | IllegalAccessException e) {
+ // This should never happen
+ throw new RuntimeException(e);
}
}
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 0834469..b0922be 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
@@ -25,13 +25,11 @@ import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.memory.ManagedMemoryUseCase;
import org.apache.flink.python.PythonConfig;
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.PythonProcessOperator;
import org.apache.flink.streaming.api.transformations.AbstractMultipleInputTransformation;
import org.apache.flink.streaming.api.transformations.OneInputTransformation;
import org.apache.flink.streaming.api.transformations.PartitionTransformation;
@@ -40,9 +38,15 @@ import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner;
import org.apache.flink.table.api.TableConfig;
import org.apache.flink.table.api.TableException;
+import org.apache.flink.shaded.guava30.com.google.common.collect.Queues;
+import org.apache.flink.shaded.guava30.com.google.common.collect.Sets;
+
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.util.List;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.Set;
/**
* A Util class to get the {@link StreamExecutionEnvironment} configuration and merged configuration
@@ -90,21 +94,6 @@ public class PythonConfigUtil {
return (Configuration) configurationField.get(env);
}
- /**
- * Generate a {@link StreamGraph} for transformations maintained by current {@link
- * StreamExecutionEnvironment}, and reset the merged env configurations with dependencies to
- * 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)
- throws IllegalAccessException, InvocationTargetException, NoSuchFieldException {
- configPythonOperator(env);
-
- return env.getStreamGraph(clearTransformations);
- }
-
@SuppressWarnings("unchecked")
public static void configPythonOperator(StreamExecutionEnvironment env)
throws IllegalAccessException, InvocationTargetException, NoSuchFieldException {
@@ -132,8 +121,6 @@ public class PythonConfigUtil {
}
}
}
-
- setPartitionCustomOperatorNumPartitions(transformations);
}
public static Configuration getMergedConfig(
@@ -175,6 +162,9 @@ public class PythonConfigUtil {
private static void alignTransformation(Transformation<?> transformation)
throws NoSuchFieldException, IllegalAccessException {
String transformName = transformation.getName();
+ if (transformation.getInputs().isEmpty()) {
+ return;
+ }
Transformation<?> inputTransformation = transformation.getInputs().get(0);
String inputTransformName = inputTransformation.getName();
if (inputTransformName.equals(KEYED_STREAM_VALUE_OPERATOR_NAME)) {
@@ -288,22 +278,53 @@ public class PythonConfigUtil {
return new PythonConfig(mergedConfig);
}
- private static void setPartitionCustomOperatorNumPartitions(
+ public static void setPartitionCustomOperatorNumPartitions(
List<Transformation<?>> transformations) {
// Update the numPartitions of PartitionCustomOperator after aligned all operators.
- for (Transformation<?> transformation : transformations) {
- Transformation<?> firstInputTransformation = transformation.getInputs().get(0);
- if (firstInputTransformation instanceof PartitionTransformation) {
- firstInputTransformation = firstInputTransformation.getInputs().get(0);
- }
- AbstractPythonFunctionOperator<?> pythonFunctionOperator =
- getPythonOperator(firstInputTransformation);
- if (pythonFunctionOperator instanceof PythonProcessOperator) {
- PythonProcessOperator<?, ?> partitionCustomFunctionOperator =
- (PythonProcessOperator<?, ?>) pythonFunctionOperator;
+ final Set<Transformation<?>> alreadyTransformed = Sets.newIdentityHashSet();
+ final Queue<Transformation<?>> toTransformQueue = Queues.newArrayDeque(transformations);
+ while (!toTransformQueue.isEmpty()) {
+ final Transformation<?> transformation = toTransformQueue.poll();
+ if (!alreadyTransformed.contains(transformation)
+ && !(transformation instanceof PartitionTransformation)) {
+ alreadyTransformed.add(transformation);
+
+ getNonPartitionTransformationInput(transformation)
+ .ifPresent(
+ input -> {
+ AbstractPythonFunctionOperator<?> pythonFunctionOperator =
+ getPythonOperator(input);
+ if (pythonFunctionOperator
+ instanceof AbstractDataStreamPythonFunctionOperator) {
+ AbstractDataStreamPythonFunctionOperator<?>
+ pythonDataStreamFunctionOperator =
+ (AbstractDataStreamPythonFunctionOperator<
+ ?>)
+ pythonFunctionOperator;
+ if (pythonDataStreamFunctionOperator
+ .containsPartitionCustom()) {
+ pythonDataStreamFunctionOperator.setNumPartitions(
+ transformation.getParallelism());
+ }
+ }
+ });
- partitionCustomFunctionOperator.setNumPartitions(transformation.getParallelism());
+ toTransformQueue.addAll(transformation.getInputs());
}
}
}
+
+ private static Optional<Transformation<?>> getNonPartitionTransformationInput(
+ Transformation<?> transformation) {
+ if (transformation.getInputs().size() != 1) {
+ return Optional.empty();
+ }
+
+ final Transformation<?> inputTransformation = transformation.getInputs().get(0);
+ if (inputTransformation instanceof PartitionTransformation) {
+ return getNonPartitionTransformationInput(inputTransformation);
+ } else {
+ return Optional.of(inputTransformation);
+ }
+ }
}
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 9aa1a05..5cf4b2d 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
@@ -44,4 +44,15 @@ public class DataStreamPythonFunctionInfo extends PythonFunctionInfo {
public int getFunctionType() {
return this.functionType;
}
+
+ public DataStreamPythonFunctionInfo copy() {
+ if (getInputs().length == 0) {
+ return new DataStreamPythonFunctionInfo(getPythonFunction(), this.functionType);
+ } else {
+ return new DataStreamPythonFunctionInfo(
+ getPythonFunction(),
+ ((DataStreamPythonFunctionInfo) getInputs()[0]).copy(),
+ this.functionType);
+ }
+ }
}
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
index 25619b5..b8feff9 100644
--- 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
@@ -26,6 +26,11 @@ import org.apache.flink.streaming.api.functions.python.DataStreamPythonFunctionI
import org.apache.flink.table.functions.python.PythonEnv;
import org.apache.flink.util.Preconditions;
+import javax.annotation.Nullable;
+
+import java.util.HashMap;
+import java.util.Map;
+
/** Base class for all Python DataStream operators. */
@Internal
public abstract class AbstractDataStreamPythonFunctionOperator<OUT>
@@ -33,6 +38,17 @@ public abstract class AbstractDataStreamPythonFunctionOperator<OUT>
private static final long serialVersionUID = 1L;
+ private static final String NUM_PARTITIONS = "NUM_PARTITIONS";
+
+ /** The number of partitions for the partition custom function. */
+ @Nullable private Integer numPartitions = null;
+
+ /**
+ * Whether it contains partition custom function. If true, the variable numPartitions should be
+ * set and the value should be set to the parallelism of the downstream operator.
+ */
+ private boolean containsPartitionCustom;
+
/** The serialized python function to be executed. */
private final DataStreamPythonFunctionInfo pythonFunctionInfo;
@@ -61,6 +77,26 @@ public abstract class AbstractDataStreamPythonFunctionOperator<OUT>
public abstract <T> AbstractDataStreamPythonFunctionOperator<T> copy(
DataStreamPythonFunctionInfo pythonFunctionInfo, TypeInformation<T> outputTypeInfo);
+ public Map<String, String> getInternalParameters() {
+ Map<String, String> internalParameters = new HashMap<>();
+ if (numPartitions != null) {
+ internalParameters.put(NUM_PARTITIONS, String.valueOf(numPartitions));
+ }
+ return internalParameters;
+ }
+
+ public void setNumPartitions(int numPartitions) {
+ this.numPartitions = numPartitions;
+ }
+
+ public void setContainsPartitionCustom(boolean containsPartitionCustom) {
+ this.containsPartitionCustom = containsPartitionCustom;
+ }
+
+ public boolean containsPartitionCustom() {
+ return this.containsPartitionCustom;
+ }
+
// ----------------------------------------------------------------------
// Getters
// ----------------------------------------------------------------------
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 f27983b..eaf07d4 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
@@ -38,7 +38,6 @@ import org.apache.flink.streaming.api.utils.PythonTypeUtils;
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;
@@ -147,10 +146,6 @@ public abstract class AbstractOneInputPythonFunctionOperator<IN, OUT>
emitResults();
}
- public Map<String, String> getInternalParameters() {
- return Collections.emptyMap();
- }
-
public FlinkFnApi.CoderInfoDescriptor createInputCoderInfoDescriptor() {
return createRawTypeCoderInfoDescriptorProto(
runnerInputTypeInfo, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE, false);
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 fd83c91..5508f6a 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
@@ -29,8 +29,6 @@ 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.PythonOperatorUtils.inBatchExecutionMode;
@@ -45,7 +43,7 @@ import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchEx
public class PythonCoProcessOperator<IN1, IN2, OUT>
extends AbstractTwoInputPythonFunctionOperator<IN1, IN2, OUT> {
- private static final long serialVersionUID = 1L;
+ private static final long serialVersionUID = 2L;
/** We listen to this ourselves because we don't have an {@link InternalTimerService}. */
private transient long currentWatermark;
@@ -74,7 +72,7 @@ public class PythonCoProcessOperator<IN1, IN2, OUT>
ProtoUtils.createUserDefinedDataStreamFunctionProtos(
getPythonFunctionInfo(),
getRuntimeContext(),
- Collections.emptyMap(),
+ getInternalParameters(),
inBatchExecutionMode(getKeyedStateBackend())),
getJobOptions(),
getFlinkMetricContainer(),
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 cf8d9ca..9038a28 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
@@ -40,8 +40,6 @@ import org.apache.flink.streaming.api.utils.PythonTypeUtils;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.types.Row;
-import java.util.Collections;
-
import static org.apache.flink.python.Constants.STATEFUL_FUNCTION_URN;
import static org.apache.flink.streaming.api.operators.python.timer.TimerUtils.createTimerDataCoderInfoDescriptorProto;
import static org.apache.flink.streaming.api.operators.python.timer.TimerUtils.createTimerDataTypeInfo;
@@ -53,7 +51,7 @@ public class PythonKeyedCoProcessOperator<OUT>
extends AbstractTwoInputPythonFunctionOperator<Row, Row, OUT>
implements Triggerable<Row, VoidNamespace> {
- private static final long serialVersionUID = 1L;
+ private static final long serialVersionUID = 2L;
/** TimerService for current operator to register or fire timer. */
private transient InternalTimerService<VoidNamespace> internalTimerService;
@@ -112,7 +110,7 @@ public class PythonKeyedCoProcessOperator<OUT>
ProtoUtils.createUserDefinedDataStreamStatefulFunctionProtos(
getPythonFunctionInfo(),
getRuntimeContext(),
- Collections.emptyMap(),
+ getInternalParameters(),
keyTypeInfo,
inBatchExecutionMode(getKeyedStateBackend())),
getJobOptions(),
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 1d58371..e3df754 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
@@ -40,8 +40,6 @@ import org.apache.flink.streaming.api.utils.PythonTypeUtils;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.types.Row;
-import java.util.Collections;
-
import static org.apache.flink.python.Constants.STATEFUL_FUNCTION_URN;
import static org.apache.flink.streaming.api.operators.python.timer.TimerUtils.createTimerDataCoderInfoDescriptorProto;
import static org.apache.flink.streaming.api.operators.python.timer.TimerUtils.createTimerDataTypeInfo;
@@ -57,7 +55,7 @@ public class PythonKeyedProcessOperator<OUT>
extends AbstractOneInputPythonFunctionOperator<Row, OUT>
implements Triggerable<Row, Object> {
- private static final long serialVersionUID = 1L;
+ private static final long serialVersionUID = 2L;
/** The TypeSerializer of the namespace. */
private final TypeSerializer namespaceSerializer;
@@ -142,7 +140,7 @@ public class PythonKeyedProcessOperator<OUT>
ProtoUtils.createUserDefinedDataStreamStatefulFunctionProtos(
getPythonFunctionInfo(),
getRuntimeContext(),
- Collections.emptyMap(),
+ getInternalParameters(),
keyTypeInfo,
inBatchExecutionMode(getKeyedStateBackend())),
getJobOptions(),
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 131927f..451418c 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
@@ -30,11 +30,6 @@ 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 javax.annotation.Nullable;
-
-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.PythonOperatorUtils.inBatchExecutionMode;
@@ -46,11 +41,7 @@ import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchEx
public class PythonProcessOperator<IN, OUT>
extends AbstractOneInputPythonFunctionOperator<IN, OUT> {
- private static final long serialVersionUID = 1L;
-
- private static final String NUM_PARTITIONS = "NUM_PARTITIONS";
-
- @Nullable private Integer numPartitions = null;
+ private static final long serialVersionUID = 2L;
/** We listen to this ourselves because we don't have an {@link InternalTimerService}. */
private transient long currentWatermark;
@@ -115,20 +106,6 @@ public class PythonProcessOperator<IN, OUT>
}
@Override
- public Map<String, String> getInternalParameters() {
- Map<String, String> internalParameters = super.getInternalParameters();
- if (numPartitions != null) {
- internalParameters = new HashMap<>(internalParameters);
- internalParameters.put(NUM_PARTITIONS, String.valueOf(numPartitions));
- }
- return internalParameters;
- }
-
- public void setNumPartitions(int numPartitions) {
- this.numPartitions = numPartitions;
- }
-
- @Override
public <T> AbstractDataStreamPythonFunctionOperator<T> copy(
DataStreamPythonFunctionInfo pythonFunctionInfo, TypeInformation<T> outputTypeInfo) {
return new PythonProcessOperator<>(
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
index 55b7d90..a693316 100644
--- 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
@@ -51,7 +51,7 @@ import static org.mockito.Mockito.mock;
public class PythonOperatorChainingOptimizerTest {
@Test
- public void testChainedTransformationPropertiesCorrectlySet() throws Exception {
+ public void testChainedTransformationPropertiesCorrectlySet() {
PythonKeyedProcessOperator<?> keyedProcessOperator =
createKeyedProcessOperator(
"f1", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
@@ -97,9 +97,8 @@ public class PythonOperatorChainingOptimizerTest {
transformations.add(keyedProcessTransformation);
transformations.add(processTransformation);
- PythonOperatorChainingOptimizer optimizer =
- new PythonOperatorChainingOptimizer(transformations, true);
- List<Transformation<?>> optimized = optimizer.optimize();
+ List<Transformation<?>> optimized =
+ PythonOperatorChainingOptimizer.optimize(transformations);
assertEquals(2, optimized.size());
OneInputTransformation<?, ?> chainedTransformation =
@@ -139,7 +138,7 @@ public class PythonOperatorChainingOptimizerTest {
}
@Test
- public void testChainingMultipleOperators() throws Exception {
+ public void testChainingMultipleOperators() {
PythonKeyedProcessOperator<?> keyedProcessOperator =
createKeyedProcessOperator(
"f1", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
@@ -177,9 +176,8 @@ public class PythonOperatorChainingOptimizerTest {
transformations.add(processTransformation1);
transformations.add(processTransformation2);
- PythonOperatorChainingOptimizer optimizer =
- new PythonOperatorChainingOptimizer(transformations, true);
- List<Transformation<?>> optimized = optimizer.optimize();
+ List<Transformation<?>> optimized =
+ PythonOperatorChainingOptimizer.optimize(transformations);
assertEquals(2, optimized.size());
OneInputTransformation<?, ?> chainedTransformation =
@@ -197,7 +195,7 @@ public class PythonOperatorChainingOptimizerTest {
}
@Test
- public void testChainingNonKeyedOperators() throws Exception {
+ public void testChainingNonKeyedOperators() {
PythonProcessOperator<?, ?> processOperator1 =
createProcessOperator(
"f1", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
@@ -225,9 +223,8 @@ public class PythonOperatorChainingOptimizerTest {
transformations.add(processTransformation1);
transformations.add(processTransformation2);
- PythonOperatorChainingOptimizer optimizer =
- new PythonOperatorChainingOptimizer(transformations, true);
- List<Transformation<?>> optimized = optimizer.optimize();
+ List<Transformation<?>> optimized =
+ PythonOperatorChainingOptimizer.optimize(transformations);
assertEquals(2, optimized.size());
OneInputTransformation<?, ?> chainedTransformation =
@@ -244,7 +241,7 @@ public class PythonOperatorChainingOptimizerTest {
}
@Test
- public void testContinuousKeyedOperators() throws Exception {
+ public void testContinuousKeyedOperators() {
PythonKeyedProcessOperator<?> keyedProcessOperator1 =
createKeyedProcessOperator(
"f1",
@@ -275,9 +272,8 @@ public class PythonOperatorChainingOptimizerTest {
transformations.add(processTransformation1);
transformations.add(processTransformation2);
- PythonOperatorChainingOptimizer optimizer =
- new PythonOperatorChainingOptimizer(transformations, true);
- List<Transformation<?>> optimized = optimizer.optimize();
+ List<Transformation<?>> optimized =
+ PythonOperatorChainingOptimizer.optimize(transformations);
assertEquals(3, optimized.size());
assertEquals(processTransformation1, optimized.get(1));
@@ -285,7 +281,7 @@ public class PythonOperatorChainingOptimizerTest {
}
@Test
- public void testMultipleChainedOperators() throws Exception {
+ public void testMultipleChainedOperators() {
PythonKeyedProcessOperator<?> keyedProcessOperator1 =
createKeyedProcessOperator(
"f1", new RowTypeInfo(Types.INT(), Types.INT()), Types.STRING());
@@ -349,9 +345,8 @@ public class PythonOperatorChainingOptimizerTest {
transformations.add(keyedProcessTransformation2);
transformations.add(processTransformation3);
- PythonOperatorChainingOptimizer optimizer =
- new PythonOperatorChainingOptimizer(transformations, true);
- List<Transformation<?>> optimized = optimizer.optimize();
+ List<Transformation<?>> optimized =
+ PythonOperatorChainingOptimizer.optimize(transformations);
assertEquals(3, optimized.size());
OneInputTransformation<?, ?> chainedTransformation1 =
@@ -381,7 +376,7 @@ public class PythonOperatorChainingOptimizerTest {
}
@Test
- public void testChainingTwoInputOperators() throws Exception {
+ public void testChainingTwoInputOperators() {
PythonKeyedCoProcessOperator<?> keyedCoProcessOperator1 =
createCoKeyedProcessOperator(
"f1",
@@ -451,9 +446,8 @@ public class PythonOperatorChainingOptimizerTest {
transformations.add(keyedProcessTransformation);
transformations.add(processTransformation3);
- PythonOperatorChainingOptimizer optimizer =
- new PythonOperatorChainingOptimizer(transformations, true);
- List<Transformation<?>> optimized = optimizer.optimize();
+ List<Transformation<?>> optimized =
+ PythonOperatorChainingOptimizer.optimize(transformations);
assertEquals(4, optimized.size());
TwoInputTransformation<?, ?, ?> chainedTransformation1 =
@@ -483,6 +477,63 @@ public class PythonOperatorChainingOptimizerTest {
"f4");
}
+ @Test
+ public void testChainingUnorderedTransformations() {
+ 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(processTransformation2);
+ transformations.add(processTransformation1);
+ transformations.add(keyedProcessTransformation);
+
+ List<Transformation<?>> optimized =
+ PythonOperatorChainingOptimizer.optimize(transformations);
+ 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");
+ }
+
// ----------------------- Utility Methods -----------------------
private void validateChainedPythonFunctions(
diff --git a/flink-python/src/test/java/org/apache/flink/python/util/PythonConfigUtilTest.java b/flink-python/src/test/java/org/apache/flink/python/util/PythonConfigUtilTest.java
index 013e37e..c6e997c 100644
--- a/flink-python/src/test/java/org/apache/flink/python/util/PythonConfigUtilTest.java
+++ b/flink-python/src/test/java/org/apache/flink/python/util/PythonConfigUtilTest.java
@@ -45,15 +45,14 @@ public class PythonConfigUtilTest {
}
@Test
- public void testJobName()
- throws IllegalAccessException, InvocationTargetException, NoSuchFieldException {
+ public void testJobName() {
String jobName = "MyTestJob";
Configuration config = new Configuration();
config.set(PipelineOptions.NAME, jobName);
StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(config);
env.fromCollection(Collections.singletonList("test")).addSink(new DiscardingSink<>());
- StreamGraph streamGraph = PythonConfigUtil.generateStreamGraphWithDependencies(env, true);
+ StreamGraph streamGraph = env.getStreamGraph(true);
assertEquals(jobName, streamGraph.getJobName());
}
}
diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunctionInfo.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunctionInfo.java
index deab8c0..c82bea6 100644
--- a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunctionInfo.java
+++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunctionInfo.java
@@ -39,7 +39,7 @@ public class PythonFunctionInfo implements Serializable {
* The input arguments, it could be an input offset of the input row or the execution result of
* another python function described as PythonFunctionInfo.
*/
- private final Object[] inputs;
+ private Object[] inputs;
public PythonFunctionInfo(PythonFunction pythonFunction, Object[] inputs) {
this.pythonFunction = Preconditions.checkNotNull(pythonFunction);
@@ -53,4 +53,9 @@ public class PythonFunctionInfo implements Serializable {
public Object[] getInputs() {
return this.inputs;
}
+
+ public void setInputs(Object[] inputs) {
+ Preconditions.checkNotNull(inputs);
+ this.inputs = inputs;
+ }
}