This is an automated email from the ASF dual-hosted git repository. hequn pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push: new 917c3d2 [FLINK-18947][python] Support partition_custom() for Python DataStream API. (#13155) 917c3d2 is described below commit 917c3d2f294f17ea8031dce2220f142c67520463 Author: Shuiqiang Chen <acqua....@alibaba-inc.com> AuthorDate: Tue Aug 18 09:31:26 2020 +0800 [FLINK-18947][python] Support partition_custom() for Python DataStream API. (#13155) --- flink-python/pom.xml | 1 + flink-python/pyflink/datastream/data_stream.py | 89 +++++++++++++++++++++- flink-python/pyflink/datastream/functions.py | 44 ++++++++++- .../pyflink/datastream/tests/test_data_stream.py | 24 ++++++ .../python/PartitionCustomKeySelector.java | 36 +++++++++ ...treamPythonPartitionCustomFunctionOperator.java | 75 ++++++++++++++++++ .../env/beam/ProcessPythonEnvironmentManager.java | 4 + .../apache/flink/python/util/PythonConfigUtil.java | 26 ++++++- .../python/AbstractPythonFunctionOperatorBase.java | 3 +- .../util/PartitionCustomTestMapFunction.java | 48 ++++++++++++ 10 files changed, 343 insertions(+), 7 deletions(-) diff --git a/flink-python/pom.xml b/flink-python/pom.xml index 3246761..3b5f5aa 100644 --- a/flink-python/pom.xml +++ b/flink-python/pom.xml @@ -460,6 +460,7 @@ under the License. <includes> <include>**/DataStreamTestCollectSink.class</include> <include>**/MyCustomSourceFunction.class</include> + <include>**/PartitionCustomTestMapFunction.class</include> </includes> <outputDirectory> ${project.build.directory}/data-stream-test diff --git a/flink-python/pyflink/datastream/data_stream.py b/flink-python/pyflink/datastream/data_stream.py index e227a51..2c5f526 100644 --- a/flink-python/pyflink/datastream/data_stream.py +++ b/flink-python/pyflink/datastream/data_stream.py @@ -15,7 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################ - +import os from typing import Callable, Union from pyflink.common import typeinfo, ExecutionConfig @@ -24,7 +24,8 @@ from pyflink.common.typeinfo import TypeInformation from pyflink.datastream.functions import _get_python_env, FlatMapFunctionWrapper, FlatMapFunction, \ MapFunction, MapFunctionWrapper, Function, FunctionWrapper, SinkFunction, FilterFunction, \ FilterFunctionWrapper, KeySelectorFunctionWrapper, KeySelector, ReduceFunction, \ - ReduceFunctionWrapper, CoMapFunction, CoFlatMapFunction + ReduceFunctionWrapper, CoMapFunction, CoFlatMapFunction, Partitioner, \ + PartitionerFunctionWrapper from pyflink.java_gateway import get_gateway @@ -445,6 +446,77 @@ class DataStream(object): """ return DataStream(self._j_data_stream.broadcast()) + def partition_custom(self, partitioner: Union[Callable, Partitioner], + key_selector: Union[Callable, KeySelector]) -> 'DataStream': + """ + Partitions a DataStream on the key returned by the selector, using a custom partitioner. + This method takes the key selector to get the key to partition on, and a partitioner that + accepts the key type. + + Note that this method works only on single field keys, i.e. the selector cannet return + tuples of fields. + + :param partitioner: The partitioner to assign partitions to keys. + :param key_selector: The KeySelector with which the DataStream is partitioned. + :return: The partitioned DataStream. + """ + if callable(key_selector): + key_selector = KeySelectorFunctionWrapper(key_selector) + if not isinstance(key_selector, (KeySelector, KeySelectorFunctionWrapper)): + raise TypeError("Parameter key_selector should be a type of KeySelector.") + + if callable(partitioner): + partitioner = PartitionerFunctionWrapper(partitioner) + if not isinstance(partitioner, (Partitioner, PartitionerFunctionWrapper)): + raise TypeError("Parameter partitioner should be a type of Partitioner.") + + gateway = get_gateway() + data_stream_num_partitions_env_key = gateway.jvm\ + .org.apache.flink.datastream.runtime.operators.python\ + .DataStreamPythonPartitionCustomFunctionOperator.DATA_STREAM_NUM_PARTITIONS + + class PartitionCustomMapFunction(MapFunction): + """ + A wrapper class for partition_custom map function. It indicates that it is a partition + custom operation that we need to apply DataStreamPythonPartitionCustomFunctionOperator + to run the map function. + """ + + def __init__(self): + self.num_partitions = None + + def map(self, value): + return self.partition_custom_map(value) + + def partition_custom_map(self, value): + if self.num_partitions is None: + self.num_partitions = int(os.environ[data_stream_num_partitions_env_key]) + partition = partitioner.partition(key_selector.get_key(value), self.num_partitions) + return partition, value + + def __repr__(self) -> str: + return '_Flink_PartitionCustomMapFunction' + + original_type_info = self.get_type() + intermediate_map_stream = self.map(PartitionCustomMapFunction(), + type_info=Types.ROW([Types.INT(), original_type_info])) + intermediate_map_stream.name( + gateway.jvm.org.apache.flink.python.util.PythonConfigUtil + .STREAM_PARTITION_CUSTOM_MAP_OPERATOR_NAME) + + JPartitionCustomKeySelector = gateway.jvm\ + .org.apache.flink.datastream.runtime.functions.python.PartitionCustomKeySelector + JIdParitioner = gateway.jvm\ + .org.apache.flink.api.java.functions.IdPartitioner + intermediate_map_stream = DataStream(intermediate_map_stream._j_data_stream + .partitionCustom(JIdParitioner(), + JPartitionCustomKeySelector())) + + values_map_stream = intermediate_map_stream.map(lambda x: x[1], original_type_info) + values_map_stream.name(gateway.jvm.org.apache.flink.python.util.PythonConfigUtil + .KEYED_STREAM_VALUE_OPERATOR_NAME) + return DataStream(values_map_stream._j_data_stream) + def _get_java_python_function_operator(self, func: Union[Function, FunctionWrapper], type_info: TypeInformation, func_name: str, func_type: int): @@ -506,8 +578,13 @@ class DataStream(object): j_python_data_stream_function_info) return j_python_data_stream_function_operator, j_output_type_info else: - DataStreamPythonFunctionOperator = gateway.jvm.org.apache.flink.datastream.runtime \ - .operators.python.DataStreamPythonStatelessFunctionOperator + if str(func) == '_Flink_PartitionCustomMapFunction': + DataStreamPythonFunctionOperator = gateway.jvm.org.apache.flink.datastream.runtime \ + .operators.python.DataStreamPythonPartitionCustomFunctionOperator + else: + DataStreamPythonFunctionOperator = gateway.jvm.org.apache.flink.datastream.runtime \ + .operators.python.DataStreamPythonStatelessFunctionOperator + j_python_data_stream_function_operator = DataStreamPythonFunctionOperator( j_conf, j_input_types, @@ -729,6 +806,10 @@ class KeyedStream(DataStream): def broadcast(self) -> 'DataStream': raise Exception('Cannot override partitioning for KeyedStream.') + def partition_custom(self, partitioner: Union[Callable, Partitioner], + key_selector: Union[Callable, KeySelector]) -> 'DataStream': + raise Exception('Cannot override partitioning for KeyedStream.') + def print(self, sink_identifier=None): return self._values().print() diff --git a/flink-python/pyflink/datastream/functions.py b/flink-python/pyflink/datastream/functions.py index 76f5025..3f576d4 100644 --- a/flink-python/pyflink/datastream/functions.py +++ b/flink-python/pyflink/datastream/functions.py @@ -17,7 +17,7 @@ ################################################################################ import abc -from typing import Union +from typing import Union, Any from py4j.java_gateway import JavaObject @@ -240,6 +240,23 @@ class FilterFunction(Function): pass +class Partitioner(Function): + """ + Function to implement a custom partition assignment for keys. + """ + + @abc.abstractmethod + def partition(self, key: Any, num_partitions: int) -> int: + """ + Computes the partition for the given key. + + :param key: The key. + :param num_partitions: The number of partitions to partition into. + :return: The partition index. + """ + pass + + class FunctionWrapper(object): """ A basic wrapper class for user defined function. @@ -360,6 +377,31 @@ class KeySelectorFunctionWrapper(FunctionWrapper): return self._func(value) +class PartitionerFunctionWrapper(FunctionWrapper): + """ + A wrapper class for Partitioner. It's used for wrapping up user defined function in a + Partitioner when user does not implement a Partitioner but directly pass a function + object or a lambda function to partition_custom() function. + """ + def __init__(self, func): + """ + The constructor of PartitionerFunctionWrapper. + + :param func: user defined function object. + """ + super(PartitionerFunctionWrapper, self).__init__(func) + + def partition(self, key: Any, num_partitions: int) -> int: + """ + A delegated partition function to invoke user defined function. + + :param key: The key. + :param num_partitions: The number of partitions to partition into. + :return: The partition index. + """ + return self._func(key, num_partitions) + + def _get_python_env(): """ An util function to get a python user defined function execution environment. diff --git a/flink-python/pyflink/datastream/tests/test_data_stream.py b/flink-python/pyflink/datastream/tests/test_data_stream.py index 2e352a7..655e847 100644 --- a/flink-python/pyflink/datastream/tests/test_data_stream.py +++ b/flink-python/pyflink/datastream/tests/test_data_stream.py @@ -19,6 +19,7 @@ import decimal from pyflink.common.typeinfo import Types from pyflink.datastream import StreamExecutionEnvironment +from pyflink.datastream.data_stream import DataStream from pyflink.datastream.functions import FilterFunction from pyflink.datastream.functions import KeySelector from pyflink.datastream.functions import MapFunction, FlatMapFunction @@ -404,6 +405,29 @@ class DataStreamTests(PyFlinkTestCase): pre_ship_strategy = shuffle_node['predecessors'][0]['ship_strategy'] self.assertEqual(pre_ship_strategy, 'SHUFFLE') + def test_partition_custom(self): + ds = self.env.from_collection([('a', 0), ('b', 0), ('c', 1), ('d', 1), ('e', 2), + ('f', 7), ('g', 7), ('h', 8), ('i', 8), ('j', 9)], + type_info=Types.ROW([Types.STRING(), Types.INT()])) + + expected_num_partitions = 5 + + def my_partitioner(key, num_partitions): + assert expected_num_partitions, num_partitions + return key % num_partitions + + partitioned_stream = ds.map(lambda x: x, type_info=Types.ROW([Types.STRING(), + Types.INT()]))\ + .set_parallelism(4).partition_custom(my_partitioner, lambda x: x[1]) + + JPartitionCustomTestMapFunction = get_gateway().jvm\ + .org.apache.flink.python.util.PartitionCustomTestMapFunction + test_map_stream = DataStream(partitioned_stream + ._j_data_stream.map(JPartitionCustomTestMapFunction())) + test_map_stream.set_parallelism(expected_num_partitions).add_sink(self.test_sink) + + self.env.execute('test_partition_custom') + def test_keyed_stream_partitioning(self): ds = self.env.from_collection([('ab', 1), ('bdc', 2), ('cfgs', 3), ('deeefg', 4)]) keyed_stream = ds.key_by(lambda x: x[1]) diff --git a/flink-python/src/main/java/org/apache/flink/datastream/runtime/functions/python/PartitionCustomKeySelector.java b/flink-python/src/main/java/org/apache/flink/datastream/runtime/functions/python/PartitionCustomKeySelector.java new file mode 100644 index 0000000..04391b3 --- /dev/null +++ b/flink-python/src/main/java/org/apache/flink/datastream/runtime/functions/python/PartitionCustomKeySelector.java @@ -0,0 +1,36 @@ +/* + * 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.datastream.runtime.functions.python; + +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.types.Row; + +/** + * The {@link PartitionCustomKeySelector} will return the first field of the input row value. The input value is + * generated by the + * {@link org.apache.flink.datastream.runtime.operators.python.DataStreamPythonPartitionCustomFunctionOperator} after + * executed user defined partitioner and keySelector function. The value of the first field will be the desired + * partition index. + */ +public class PartitionCustomKeySelector implements KeySelector<Row, Integer> { + @Override + public Integer getKey(Row value) throws Exception { + return (Integer) value.getField(0); + } +} diff --git a/flink-python/src/main/java/org/apache/flink/datastream/runtime/operators/python/DataStreamPythonPartitionCustomFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/datastream/runtime/operators/python/DataStreamPythonPartitionCustomFunctionOperator.java new file mode 100644 index 0000000..2a39935 --- /dev/null +++ b/flink-python/src/main/java/org/apache/flink/datastream/runtime/operators/python/DataStreamPythonPartitionCustomFunctionOperator.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.datastream.runtime.operators.python; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.CoreOptions; +import org.apache.flink.datastream.runtime.functions.python.DataStreamPythonFunctionInfo; +import org.apache.flink.datastream.runtime.runners.python.beam.BeamDataStreamPythonStatelessFunctionRunner; +import org.apache.flink.python.PythonFunctionRunner; +import org.apache.flink.python.env.PythonEnvironmentManager; +import org.apache.flink.python.env.beam.ProcessPythonEnvironmentManager; + +/** + * The {@link DataStreamPythonPartitionCustomFunctionOperator} enables us to set the number of partitions for current + * operator dynamically when generating the {@link org.apache.flink.streaming.api.graph.StreamGraph} before executing + * the job. The number of partitions will be set in environment variables for python Worker, so that we can obtain the + * number of partitions when executing user defined partitioner function. + */ +public class DataStreamPythonPartitionCustomFunctionOperator<IN, OUT> extends + DataStreamPythonStatelessFunctionOperator<IN, OUT> { + + public static final String DATA_STREAM_NUM_PARTITIONS = "DATA_STREAM_NUM_PARTITIONS"; + + private int numPartitions = CoreOptions.DEFAULT_PARALLELISM.defaultValue(); + + public DataStreamPythonPartitionCustomFunctionOperator( + Configuration config, + TypeInformation<IN> inputTypeInfo, + TypeInformation<OUT> outputTypeInfo, + DataStreamPythonFunctionInfo pythonFunctionInfo) { + super(config, inputTypeInfo, outputTypeInfo, pythonFunctionInfo); + } + + @Override + public PythonFunctionRunner createPythonFunctionRunner() throws Exception { + PythonEnvironmentManager pythonEnvironmentManager = createPythonEnvironmentManager(); + if (pythonEnvironmentManager instanceof ProcessPythonEnvironmentManager) { + ProcessPythonEnvironmentManager envManager = (ProcessPythonEnvironmentManager) pythonEnvironmentManager; + envManager.setEnvironmentVariable(DATA_STREAM_NUM_PARTITIONS, + String.valueOf(this.numPartitions)); + } + return new BeamDataStreamPythonStatelessFunctionRunner( + getRuntimeContext().getTaskName(), + pythonEnvironmentManager, + inputTypeInfo, + outputTypeInfo, + DATA_STREAM_STATELESS_PYTHON_FUNCTION_URN, + getUserDefinedDataStreamFunctionsProto(), + DATA_STREAM_MAP_FUNCTION_CODER_URN, + jobOptions, + getFlinkMetricContainer() + ); + } + + public void setNumPartitions(int numPartitions) { + this.numPartitions = numPartitions; + } +} diff --git a/flink-python/src/main/java/org/apache/flink/python/env/beam/ProcessPythonEnvironmentManager.java b/flink-python/src/main/java/org/apache/flink/python/env/beam/ProcessPythonEnvironmentManager.java index be11fad..6d7ce46 100644 --- a/flink-python/src/main/java/org/apache/flink/python/env/beam/ProcessPythonEnvironmentManager.java +++ b/flink-python/src/main/java/org/apache/flink/python/env/beam/ProcessPythonEnvironmentManager.java @@ -240,6 +240,10 @@ public final class ProcessPythonEnvironmentManager implements PythonEnvironmentM return env; } + public void setEnvironmentVariable(String key, String value) { + this.systemEnv.put(key, value); + } + private void constructFilesDirectory(Map<String, String> env) throws IOException { // link or copy python files to filesDirectory and add them to PYTHONPATH List<String> pythonFilePaths = new ArrayList<>(); 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 6605959..f6d13c5 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 @@ -18,6 +18,7 @@ package org.apache.flink.python.util; import org.apache.flink.configuration.Configuration; +import org.apache.flink.datastream.runtime.operators.python.DataStreamPythonPartitionCustomFunctionOperator; import org.apache.flink.datastream.runtime.operators.python.DataStreamPythonStatelessFunctionOperator; import org.apache.flink.python.PythonConfig; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; @@ -41,6 +42,7 @@ public class PythonConfigUtil { public static final String KEYED_STREAM_VALUE_OPERATOR_NAME = "_keyed_stream_values_operator"; public static final String STREAM_KEY_BY_MAP_OPERATOR_NAME = "_stream_key_by_map_operator"; + public static final String STREAM_PARTITION_CUSTOM_MAP_OPERATOR_NAME = "_partition_custom_map_operator"; /** * A static method to get the {@link StreamExecutionEnvironment} configuration merged with python dependency @@ -94,7 +96,8 @@ public class PythonConfigUtil { downStreamEdge.setPartitioner(new ForwardPartitioner()); } - if (streamNode.getOperatorName().equals(STREAM_KEY_BY_MAP_OPERATOR_NAME)) { + if (streamNode.getOperatorName().equals(STREAM_KEY_BY_MAP_OPERATOR_NAME) || + streamNode.getOperatorName().equals(STREAM_PARTITION_CUSTOM_MAP_OPERATOR_NAME)) { StreamEdge upStreamEdge = streamNode.getInEdges().get(0); StreamNode upStreamNode = streamGraph.getStreamNode(upStreamEdge.getSourceId()); chainStreamNode(upStreamEdge, streamNode, upStreamNode); @@ -139,9 +142,30 @@ public class PythonConfigUtil { } } } + + setStreamPartitionCustomOperatorNumPartitions(streamNodes, streamGraph); + return streamGraph; } + private static void setStreamPartitionCustomOperatorNumPartitions( + Collection<StreamNode> streamNodes, StreamGraph streamGraph){ + for (StreamNode streamNode : streamNodes) { + StreamOperatorFactory streamOperatorFactory = streamNode.getOperatorFactory(); + if (streamOperatorFactory instanceof SimpleOperatorFactory) { + StreamOperator streamOperator = ((SimpleOperatorFactory) streamOperatorFactory).getOperator(); + if (streamOperator instanceof DataStreamPythonPartitionCustomFunctionOperator) { + DataStreamPythonPartitionCustomFunctionOperator paritionCustomFunctionOperator = + (DataStreamPythonPartitionCustomFunctionOperator) streamOperator; + + // Update the numPartitions of PartitionCustomOperator after aligned all operators. + paritionCustomFunctionOperator.setNumPartitions( + streamGraph.getStreamNode(streamNode.getOutEdges().get(0).getTargetId()).getParallelism()); + } + } + } + } + /** * 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/operators/python/AbstractPythonFunctionOperatorBase.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperatorBase.java index 280bf0a..5a989ba 100644 --- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperatorBase.java +++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperatorBase.java @@ -38,6 +38,7 @@ import org.apache.flink.table.functions.python.PythonEnv; import org.apache.flink.util.Preconditions; import java.io.IOException; +import java.util.HashMap; import java.util.concurrent.ScheduledFuture; /** @@ -336,7 +337,7 @@ public abstract class AbstractPythonFunctionOperatorBase<OUT> return new ProcessPythonEnvironmentManager( dependencyInfo, getContainingTask().getEnvironment().getTaskManagerInfo().getTmpDirectories(), - System.getenv()); + new HashMap<>(System.getenv())); } else { throw new UnsupportedOperationException(String.format( "Execution type '%s' is not supported.", pythonEnv.getExecType())); diff --git a/flink-python/src/test/java/org/apache/flink/python/util/PartitionCustomTestMapFunction.java b/flink-python/src/test/java/org/apache/flink/python/util/PartitionCustomTestMapFunction.java new file mode 100644 index 0000000..5484194 --- /dev/null +++ b/flink-python/src/test/java/org/apache/flink/python/util/PartitionCustomTestMapFunction.java @@ -0,0 +1,48 @@ +/* + * 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.util; + +import org.apache.flink.api.common.functions.RichMapFunction; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.types.Row; + +/** + * {@link PartitionCustomTestMapFunction} is a dedicated MapFunction to make sure the specific field data is equal to + * current sub-task index. + */ +public class PartitionCustomTestMapFunction extends RichMapFunction<Row, Row> { + + private int currentTaskIndex; + + @Override + public void open(Configuration parameters) { + this.currentTaskIndex = getRuntimeContext().getIndexOfThisSubtask(); + } + + @Override + public Row map(Row value) throws Exception { + int expectedPartitionIndex = (Integer) (value.getField(1)) % getRuntimeContext() + .getNumberOfParallelSubtasks(); + if (expectedPartitionIndex != currentTaskIndex) { + throw new RuntimeException(String.format("the data: Row<%s> was sent to the wrong partition[%d], " + + "expected partition is [%d].", value.toString(), currentTaskIndex, expectedPartitionIndex)); + } + return value; + } +}