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 f1e34e6 [FLINK-18866][python] Support filter() operation for Python
DataStream API. (#13098)
f1e34e6 is described below
commit f1e34e659aed9f7bbf7cd7c38e10e97ca88e4b73
Author: Shuiqiang Chen <acqua....@alibaba-inc.com>
AuthorDate: Tue Aug 11 19:44:51 2020 +0800
[FLINK-18866][python] Support filter() operation for Python DataStream API.
(#13098)
---
flink-python/pyflink/datastream/data_stream.py | 54 ++++++++++++++++++++--
flink-python/pyflink/datastream/functions.py | 36 +++++++++++++++
.../pyflink/datastream/tests/test_data_stream.py | 43 +++++++++++++++++
3 files changed, 128 insertions(+), 5 deletions(-)
diff --git a/flink-python/pyflink/datastream/data_stream.py
b/flink-python/pyflink/datastream/data_stream.py
index d6f44b0..4d1b0dc 100644
--- a/flink-python/pyflink/datastream/data_stream.py
+++ b/flink-python/pyflink/datastream/data_stream.py
@@ -22,8 +22,9 @@ from pyflink.common import typeinfo, ExecutionConfig
from pyflink.common.typeinfo import RowTypeInfo, PickledBytesTypeInfo, Types
from pyflink.common.typeinfo import TypeInformation
from pyflink.datastream.functions import _get_python_env,
FlatMapFunctionWrapper, FlatMapFunction, \
- MapFunction, MapFunctionWrapper, Function, FunctionWrapper, SinkFunction, \
- KeySelectorFunctionWrapper, KeySelector, ReduceFunction,
ReduceFunctionWrapper
+ MapFunction, MapFunctionWrapper, Function, FunctionWrapper, SinkFunction,
FilterFunction, \
+ FilterFunctionWrapper, KeySelectorFunctionWrapper, KeySelector,
ReduceFunction, \
+ ReduceFunctionWrapper
from pyflink.java_gateway import get_gateway
@@ -262,10 +263,41 @@ class DataStream(object):
output_type_info]))
._j_data_stream
.keyBy(PickledKeySelector(is_key_pickled_byte_array),
-
key_type_info.get_java_type_info()))
+
key_type_info.get_java_type_info()), self)
generated_key_stream._original_data_type_info = output_type_info
return generated_key_stream
+ def filter(self, func: Union[Callable, FilterFunction]) -> 'DataStream':
+ """
+ Applies a Filter transformation on a DataStream. The transformation
calls a FilterFunction
+ for each element of the DataStream and retains only those element for
which the function
+ returns true. Elements for which the function returns false are
filtered. The user can also
+ extend RichFilterFunction to gain access to other features provided by
the RichFunction
+ interface.
+
+ :param func: The FilterFunction that is called for each element of the
DataStream.
+ :return: The filtered DataStream.
+ """
+ class FilterFlatMap(FlatMapFunction):
+ def __init__(self, filter_func):
+ self._func = filter_func
+
+ def flat_map(self, value):
+ if self._func.filter(value):
+ yield value
+
+ if isinstance(func, Callable):
+ func = FilterFunctionWrapper(func)
+ elif not isinstance(func, FilterFunction):
+ raise TypeError("func must be a Callable or instance of
FilterFunction.")
+
+ j_input_type = self._j_data_stream.getTransformation().getOutputType()
+ type_info = typeinfo._from_java_type(j_input_type)
+ j_data_stream = self.flat_map(FilterFlatMap(func),
type_info=type_info)._j_data_stream
+ filtered_stream = DataStream(j_data_stream)
+ filtered_stream.name("Filter")
+ return filtered_stream
+
def _get_java_python_function_operator(self, func: Union[Function,
FunctionWrapper],
type_info: TypeInformation,
func_name: str,
func_type: int):
@@ -432,14 +464,16 @@ class KeyedStream(DataStream):
Reduce-style operations, such as reduce and sum work on elements that have
the same key.
"""
- def __init__(self, j_keyed_stream):
+ def __init__(self, j_keyed_stream, origin_stream: DataStream):
"""
Constructor of KeyedStream.
:param j_keyed_stream: A java KeyedStream object.
+ :param origin_stream: The DataStream before key by.
"""
super(KeyedStream, self).__init__(j_data_stream=j_keyed_stream)
self._original_data_type_info = None
+ self._origin_stream = origin_stream
def map(self, func: Union[Callable, MapFunction], type_info:
TypeInformation = None) \
-> 'DataStream':
@@ -483,7 +517,17 @@ class KeyedStream(DataStream):
j_python_data_stream_scalar_function_operator
))
- def _values(self):
+ def filter(self, func: Union[Callable, FilterFunction]) -> 'DataStream':
+ return self._values().filter(func)
+
+ def add_sink(self, sink_func: SinkFunction) -> 'DataStreamSink':
+ return self._values().add_sink(sink_func)
+
+ def key_by(self, key_selector: Union[Callable, KeySelector],
+ key_type_info: TypeInformation = None) -> 'KeyedStream':
+ return self._origin_stream.key_by(key_selector, key_type_info)
+
+ def _values(self) -> 'DataStream':
"""
Since python KeyedStream is in the format of Row(key_value,
original_data), it is used for
getting the original_data.
diff --git a/flink-python/pyflink/datastream/functions.py
b/flink-python/pyflink/datastream/functions.py
index 30f8c6f..c5049e8 100644
--- a/flink-python/pyflink/datastream/functions.py
+++ b/flink-python/pyflink/datastream/functions.py
@@ -131,6 +131,29 @@ class KeySelector(Function):
pass
+class FilterFunction(Function):
+ """
+ A filter function is a predicate applied individually to each record. The
predicate decides
+ whether to keep the element, or to discard it.
+ The basic syntax for using a FilterFunction is as follows:
+ :
+ >>> ds = ...
+ >>> result = ds.filter(MyFilterFunction())
+ Note that the system assumes that the function does not modify the
elements on which the
+ predicate is applied. Violating this assumption can lead to incorrect
results.
+ """
+
+ @abc.abstractmethod
+ def filter(self, value):
+ """
+ The filter function that evaluates the predicate.
+
+ :param value: The value to be filtered.
+ :return: True for values that should be retained, false for values to
be filtered out.
+ """
+ pass
+
+
class FunctionWrapper(object):
"""
A basic wrapper class for user defined function.
@@ -188,6 +211,19 @@ class FlatMapFunctionWrapper(FunctionWrapper):
return self._func(value)
+class FilterFunctionWrapper(FunctionWrapper):
+ """
+ A wrapper class for FilterFunction. It's used for wrapping up user
defined function in a
+ FilterFunction when user does not implement a FilterFunction but
directly pass a function
+ object or a lambda function to filter() function.
+ """
+ def __init__(self, func):
+ super(FilterFunctionWrapper, self).__init__(func)
+
+ def filter(self, value):
+ return self._func(value)
+
+
class ReduceFunctionWrapper(FunctionWrapper):
"""
A wrapper class for ReduceFunction. It's used for wrapping up user defined
function in a
diff --git a/flink-python/pyflink/datastream/tests/test_data_stream.py
b/flink-python/pyflink/datastream/tests/test_data_stream.py
index 6d4fa49..1abd80d 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.functions import FilterFunction
from pyflink.datastream.functions import KeySelector
from pyflink.datastream.functions import MapFunction, FlatMapFunction
from pyflink.datastream.tests.test_util import DataStreamTestSinkFunction
@@ -148,6 +149,29 @@ class DataStreamTests(PyFlinkTestCase):
expected.sort()
self.assertEqual(expected, results)
+ def test_filter_without_data_types(self):
+ ds = self.env.from_collection([(1, 'Hi', 'Hello'), (2, 'Hello', 'Hi')])
+ ds.filter(MyFilterFunction()).add_sink(self.test_sink)
+ self.env.execute("test filter")
+ results = self.test_sink.get_results(True)
+ expected = ["(2, 'Hello', 'Hi')"]
+ results.sort()
+ expected.sort()
+ self.assertEqual(expected, results)
+
+ def test_filter_with_data_types(self):
+ ds = self.env.from_collection([(1, 'Hi', 'Hello'), (2, 'Hello', 'Hi')],
+ type_info=Types.ROW(
+ [Types.INT(), Types.STRING(),
Types.STRING()])
+ )
+ ds.filter(lambda x: x[0] % 2 == 0).add_sink(self.test_sink)
+ self.env.execute("test filter")
+ results = self.test_sink.get_results(False)
+ expected = ['2,Hello,Hi']
+ results.sort()
+ expected.sort()
+ self.assertEqual(expected, results)
+
def test_add_sink(self):
ds = self.env.from_collection([('ab', 1), ('bdc', 2), ('cfgs', 3),
('deeefg', 4)],
type_info=Types.ROW([Types.STRING(),
Types.INT()]))
@@ -188,6 +212,19 @@ class DataStreamTests(PyFlinkTestCase):
expected.sort()
self.assertEqual(expected, results)
+ def test_multi_key_by(self):
+ ds = self.env.from_collection([('a', 0), ('b', 0), ('c', 1), ('d', 1),
('e', 2)],
+ type_info=Types.ROW([Types.STRING(),
Types.INT()]))
+ ds.key_by(MyKeySelector(), key_type_info=Types.INT()).key_by(lambda x:
x[0])\
+ .add_sink(self.test_sink)
+
+ self.env.execute("test multi key by")
+ results = self.test_sink.get_results(False)
+ expected = ['d,1', 'c,1', 'a,0', 'b,0', 'e,2']
+ results.sort()
+ expected.sort()
+ self.assertEqual(expected, results)
+
def tearDown(self) -> None:
self.test_sink.get_results()
@@ -209,3 +246,9 @@ class MyFlatMapFunction(FlatMapFunction):
class MyKeySelector(KeySelector):
def get_key(self, value):
return value[1]
+
+
+class MyFilterFunction(FilterFunction):
+
+ def filter(self, value):
+ return value[0] % 2 == 0
