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/13 01:46:21 UTC
[flink] branch master updated: [FLINK-23309][python] Optimize the
finish bundle logic to make the Python UDF results could be processed in a
pipeline manner during finishBundle
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
The following commit(s) were added to refs/heads/master by this push:
new eab562b [FLINK-23309][python] Optimize the finish bundle logic to make the Python UDF results could be processed in a pipeline manner during finishBundle
eab562b is described below
commit eab562bd3c58089e8bc945a83826def3438dc25f
Author: huangxingbo <hx...@gmail.com>
AuthorDate: Mon Jul 12 16:55:54 2021 +0800
[FLINK-23309][python] Optimize the finish bundle logic to make the Python UDF results could be processed in a pipeline manner during finishBundle
This closes #16467.
---
.../fn_execution/beam/beam_coder_impl_fast.pxd | 2 +
.../fn_execution/beam/beam_coder_impl_fast.pyx | 7 ++-
.../fn_execution/beam/beam_coder_impl_slow.py | 27 +++++++--
.../pyflink/fn_execution/beam/beam_coders.py | 13 ++--
.../fn_execution/beam/beam_operations_fast.pxd | 5 +-
.../fn_execution/beam/beam_operations_fast.pyx | 10 +++-
.../fn_execution/beam/beam_operations_slow.py | 10 +++-
.../pyflink/fn_execution/beam/beam_stream_fast.pxd | 13 +++-
.../pyflink/fn_execution/beam/beam_stream_fast.pyx | 41 ++++++++++---
.../pyflink/fn_execution/beam/beam_stream_slow.py | 30 +++++++++-
.../pyflink/fn_execution/coder_impl_fast.pxd | 13 ++++
.../pyflink/fn_execution/coder_impl_fast.pyx | 70 ++++++++++++++++++++++
.../pyflink/fn_execution/coder_impl_slow.py | 10 ++--
flink-python/pyflink/fn_execution/coders.py | 25 ++------
flink-python/pyflink/fn_execution/stream_fast.pxd | 3 +
flink-python/pyflink/fn_execution/stream_fast.pyx | 21 +++++++
.../apache/flink/python/PythonFunctionRunner.java | 10 ++++
.../python/AbstractPythonFunctionOperator.java | 46 +++++++++++++-
.../python/beam/BeamPythonFunctionRunner.java | 13 ++++
19 files changed, 308 insertions(+), 61 deletions(-)
diff --git a/flink-python/pyflink/fn_execution/beam/beam_coder_impl_fast.pxd b/flink-python/pyflink/fn_execution/beam/beam_coder_impl_fast.pxd
index 0afa27f..0cf78de 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_coder_impl_fast.pxd
+++ b/flink-python/pyflink/fn_execution/beam/beam_coder_impl_fast.pxd
@@ -24,6 +24,7 @@ from apache_beam.coders.coder_impl cimport OutputStream as BOutputStream
from apache_beam.coders.coder_impl cimport StreamCoderImpl
from pyflink.fn_execution.coder_impl_fast cimport LengthPrefixBaseCoderImpl, FieldCoderImpl
+from pyflink.fn_execution.beam.beam_stream_fast cimport BeamTimeBasedOutputStream
from pyflink.fn_execution.stream_fast cimport OutputStream
cdef class PassThroughLengthPrefixCoderImpl(StreamCoderImpl):
@@ -37,3 +38,4 @@ cdef class FlinkFieldCoderBeamWrapper(StreamCoderImpl):
cdef class FlinkLengthPrefixCoderBeamWrapper(StreamCoderImpl):
cdef readonly LengthPrefixBaseCoderImpl _value_coder
+ cdef BeamTimeBasedOutputStream _output_stream
diff --git a/flink-python/pyflink/fn_execution/beam/beam_coder_impl_fast.pyx b/flink-python/pyflink/fn_execution/beam/beam_coder_impl_fast.pyx
index d8ab52f..294eff9 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_coder_impl_fast.pyx
+++ b/flink-python/pyflink/fn_execution/beam/beam_coder_impl_fast.pyx
@@ -26,7 +26,7 @@ from apache_beam.coders.coder_impl cimport OutputStream as BOutputStream
from apache_beam.coders.coder_impl cimport StreamCoderImpl
from pyflink.fn_execution.beam.beam_stream_fast cimport BeamInputStream
-from pyflink.fn_execution.beam.beam_stream_fast import BeamOutputStream
+from pyflink.fn_execution.beam.beam_stream_fast cimport BeamTimeBasedOutputStream
from pyflink.fn_execution.stream_fast cimport InputStream
cdef class PassThroughLengthPrefixCoderImpl(StreamCoderImpl):
@@ -94,10 +94,11 @@ cdef class FlinkLengthPrefixCoderBeamWrapper(StreamCoderImpl):
"""
def __cinit__(self, value_coder):
self._value_coder = value_coder
+ self._output_stream = BeamTimeBasedOutputStream()
cpdef encode_to_stream(self, value, BOutputStream out_stream, bint nested):
- output_stream = BeamOutputStream(out_stream)
- self._value_coder.encode_to_stream(value, output_stream)
+ self._output_stream.reset_output_stream(out_stream)
+ self._value_coder.encode_to_stream(value, self._output_stream)
cpdef decode_from_stream(self, BInputStream in_stream, bint nested):
cdef BeamInputStream input_stream = BeamInputStream(in_stream, in_stream.size())
diff --git a/flink-python/pyflink/fn_execution/beam/beam_coder_impl_slow.py b/flink-python/pyflink/fn_execution/beam/beam_coder_impl_slow.py
index 7cfa4c1..e87fff1 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_coder_impl_slow.py
+++ b/flink-python/pyflink/fn_execution/beam/beam_coder_impl_slow.py
@@ -20,7 +20,7 @@ from typing import Any
from apache_beam.coders.coder_impl import StreamCoderImpl, create_InputStream, create_OutputStream
from pyflink.fn_execution.stream_slow import OutputStream
-from pyflink.fn_execution.beam.beam_stream_slow import BeamInputStream
+from pyflink.fn_execution.beam.beam_stream_slow import BeamInputStream, BeamTimeBasedOutputStream
class PassThroughLengthPrefixCoderImpl(StreamCoderImpl):
@@ -40,9 +40,9 @@ class PassThroughLengthPrefixCoderImpl(StreamCoderImpl):
return 'PassThroughLengthPrefixCoderImpl[%s]' % self._value_coder
-class FlinkCoderBeamWrapper(StreamCoderImpl):
+class FlinkFieldCoderBeamWrapper(StreamCoderImpl):
"""
- Bridge between Beam coder and Flink coder.
+ Bridge between Beam coder and Flink coder for the low-level FieldCoder.
"""
def __init__(self, value_coder):
self._value_coder = value_coder
@@ -58,4 +58,23 @@ class FlinkCoderBeamWrapper(StreamCoderImpl):
return self._value_coder.decode_from_stream(data_input_stream)
def __repr__(self):
- return 'FlinkCoderBeamWrapper[%s]' % self._value_coder
+ return 'FlinkFieldCoderBeamWrapper[%s]' % self._value_coder
+
+
+class FlinkLengthPrefixCoderBeamWrapper(FlinkFieldCoderBeamWrapper):
+ """
+ Bridge between Beam coder and Flink coder for the top-level LengthPrefixCoder.
+ """
+ def __init__(self, value_coder):
+ super(FlinkLengthPrefixCoderBeamWrapper, self).__init__(value_coder)
+ self._output_stream = BeamTimeBasedOutputStream()
+
+ def encode_to_stream(self, value, out_stream: create_OutputStream, nested):
+ self._output_stream.reset_output_stream(out_stream)
+
+ self._value_coder.encode_to_stream(value, self._data_output_stream)
+ self._output_stream.write(self._data_output_stream.get())
+ self._data_output_stream.clear()
+
+ def __repr__(self):
+ return 'FlinkLengthPrefixCoderBeamWrapper[%s]' % self._value_coder
diff --git a/flink-python/pyflink/fn_execution/beam/beam_coders.py b/flink-python/pyflink/fn_execution/beam/beam_coders.py
index a8afb27..982ef26 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_coders.py
+++ b/flink-python/pyflink/fn_execution/beam/beam_coders.py
@@ -20,8 +20,7 @@ from apache_beam.coders.coders import FastCoder, LengthPrefixCoder
from apache_beam.portability import common_urns
from apache_beam.typehints import typehints
-from pyflink.fn_execution.coders import (
- ArrowCoder, OverWindowArrowCoder, LengthPrefixBaseCoder)
+from pyflink.fn_execution.coders import LengthPrefixBaseCoder
from pyflink.fn_execution.flink_fn_execution_pb2 import CoderInfoDescriptor
try:
@@ -30,8 +29,8 @@ try:
from pyflink.fn_execution.beam.beam_coder_impl_fast import FlinkLengthPrefixCoderBeamWrapper
except ImportError:
from pyflink.fn_execution.beam import beam_coder_impl_slow as beam_coder_impl
- FlinkFieldCoderBeamWrapper = beam_coder_impl.FlinkCoderBeamWrapper
- FlinkLengthPrefixCoderBeamWrapper = beam_coder_impl.FlinkCoderBeamWrapper
+ from pyflink.fn_execution.beam.beam_coder_impl_slow import FlinkFieldCoderBeamWrapper
+ from pyflink.fn_execution.beam.beam_coder_impl_slow import FlinkLengthPrefixCoderBeamWrapper
FLINK_CODER_URN = "flink:coder:v1"
@@ -66,11 +65,7 @@ class FlinkCoder(FastCoder):
def get_impl(self):
if isinstance(self._internal_coder, LengthPrefixBaseCoder):
- if isinstance(self._internal_coder._field_coder, (ArrowCoder, OverWindowArrowCoder)):
- from pyflink.fn_execution.beam.beam_coder_impl_slow import FlinkCoderBeamWrapper
- return FlinkCoderBeamWrapper(self._create_impl())
- else:
- return FlinkLengthPrefixCoderBeamWrapper(self._create_impl())
+ return FlinkLengthPrefixCoderBeamWrapper(self._create_impl())
else:
return FlinkFieldCoderBeamWrapper(self._create_impl())
diff --git a/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pxd b/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pxd
index dd5fd71..512f0ff 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pxd
+++ b/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pxd
@@ -17,10 +17,10 @@
################################################################################
# cython: language_level=3
-from apache_beam.coders.coder_impl cimport StreamCoderImpl
from apache_beam.runners.worker.operations cimport Operation
from apache_beam.utils.windowed_value cimport WindowedValue
+from pyflink.fn_execution.beam.beam_coder_impl_fast cimport FlinkLengthPrefixCoderBeamWrapper
from pyflink.fn_execution.coder_impl_fast cimport InputStreamWrapper
cdef class InputProcessor:
@@ -37,9 +37,10 @@ cdef class IntermediateInputProcessor(InputProcessor):
cdef class OutputProcessor:
cdef Operation _consumer
cpdef process_outputs(self, WindowedValue windowed_value, results)
+ cpdef close(self)
cdef class NetworkOutputProcessor(OutputProcessor):
- cdef StreamCoderImpl _value_coder_impl
+ cdef FlinkLengthPrefixCoderBeamWrapper _value_coder_impl
cdef class IntermediateOutputProcessor(OutputProcessor):
pass
diff --git a/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx b/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx
index 06fbce3..f7b535e 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx
+++ b/flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx
@@ -20,9 +20,11 @@
# cython: profile=True
# cython: boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True
from libc.stdint cimport *
+
from apache_beam.coders.coder_impl cimport OutputStream as BOutputStream
from apache_beam.utils cimport windowed_value
from apache_beam.utils.windowed_value cimport WindowedValue
+
from pyflink.fn_execution.coder_impl_fast cimport InputStreamWrapper
from apache_beam.runners.worker.bundle_processor import DataOutputOperation
@@ -75,19 +77,24 @@ cdef class OutputProcessor:
cpdef process_outputs(self, WindowedValue windowed_value, results):
pass
+ cpdef close(self):
+ pass
cdef class NetworkOutputProcessor(OutputProcessor):
def __init__(self, consumer):
assert isinstance(consumer, DataOutputOperation)
self._consumer = consumer
- self._value_coder_impl = consumer.windowed_coder.wrapped_value_coder.get_impl()
+ self._value_coder_impl = consumer.windowed_coder.wrapped_value_coder.get_impl()._value_coder
cpdef process_outputs(self, WindowedValue windowed_value, results):
output_stream = self._consumer.output_stream
self._value_coder_impl.encode_to_stream(results, output_stream, True)
output_stream.maybe_flush()
+ cpdef close(self):
+ if self._value_coder_impl._output_stream:
+ self._value_coder_impl._output_stream.close()
cdef class IntermediateOutputProcessor(OutputProcessor):
@@ -144,6 +151,7 @@ cdef class FunctionOperation(Operation):
cpdef teardown(self):
with self.scoped_finish_state:
self.operation.close()
+ self._output_processor.close()
cpdef process(self, WindowedValue o):
cdef InputStreamWrapper input_stream_wrapper
diff --git a/flink-python/pyflink/fn_execution/beam/beam_operations_slow.py b/flink-python/pyflink/fn_execution/beam/beam_operations_slow.py
index 580e3f8..c81c516 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_operations_slow.py
+++ b/flink-python/pyflink/fn_execution/beam/beam_operations_slow.py
@@ -34,19 +34,26 @@ class OutputProcessor(abc.ABC):
def process_outputs(self, windowed_value: WindowedValue, results: Iterable[Any]):
pass
+ def close(self):
+ pass
+
class NetworkOutputProcessor(OutputProcessor):
def __init__(self, consumer):
assert isinstance(consumer, DataOutputOperation)
self._consumer = consumer
- self._value_coder_impl = consumer.windowed_coder.wrapped_value_coder.get_impl()
+ self._value_coder_impl = consumer.windowed_coder.wrapped_value_coder.get_impl()._value_coder
def process_outputs(self, windowed_value: WindowedValue, results: Iterable[Any]):
output_stream = self._consumer.output_stream
self._value_coder_impl.encode_to_stream(results, output_stream, True)
output_stream.maybe_flush()
+ def close(self):
+ if self._value_coder_impl._output_stream:
+ self._value_coder_impl._output_stream.close()
+
class IntermediateOutputProcessor(OutputProcessor):
@@ -104,6 +111,7 @@ class FunctionOperation(Operation):
def teardown(self):
with self.scoped_finish_state:
self.operation.close()
+ self._output_processor.close()
def progress_metrics(self):
metrics = super(FunctionOperation, self).progress_metrics()
diff --git a/flink-python/pyflink/fn_execution/beam/beam_stream_fast.pxd b/flink-python/pyflink/fn_execution/beam/beam_stream_fast.pxd
index 499698f..3afffb4 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_stream_fast.pxd
+++ b/flink-python/pyflink/fn_execution/beam/beam_stream_fast.pxd
@@ -29,10 +29,17 @@ cdef class BeamInputStream(LengthPrefixInputStream):
cdef BInputStream _input_stream
cdef void _parse_input_stream(self, BInputStream input_stream)
-cdef class BeamOutputStream(LengthPrefixOutputStream):
+cdef class BeamSizeBasedOutputStream(LengthPrefixOutputStream):
cdef char*_output_data
cdef size_t _output_pos
cdef size_t _output_buffer_size
cdef BOutputStream _output_stream
- cdef void _maybe_flush(self)
- cdef void _parse_output_stream(self, BOutputStream output_stream)
+
+ cdef void reset_output_stream(self, BOutputStream output_stream)
+ cdef bint _maybe_flush(self)
+
+cdef class BeamTimeBasedOutputStream(BeamSizeBasedOutputStream):
+ cdef bint _flush_event
+ cdef object _periodic_flusher
+
+ cpdef void notify_flush(self)
diff --git a/flink-python/pyflink/fn_execution/beam/beam_stream_fast.pyx b/flink-python/pyflink/fn_execution/beam/beam_stream_fast.pyx
index f5981f3..c165f5e 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_stream_fast.pyx
+++ b/flink-python/pyflink/fn_execution/beam/beam_stream_fast.pyx
@@ -23,6 +23,8 @@
from libc.stdlib cimport realloc
from libc.string cimport memcpy
+from apache_beam.runners.worker.data_plane import PeriodicThread
+
cdef class BeamInputStream(LengthPrefixInputStream):
def __cinit__(self, input_stream, size):
self._input_buffer_size = size
@@ -54,14 +56,11 @@ cdef class BeamInputStream(LengthPrefixInputStream):
self._input_data = input_stream.allc
self._input_pos = input_stream.pos
-cdef class BeamOutputStream(LengthPrefixOutputStream):
- def __cinit__(self, output_stream):
- self._output_stream = output_stream
- self._parse_output_stream(output_stream)
-
+cdef class BeamSizeBasedOutputStream(LengthPrefixOutputStream):
cdef void write(self, char*data, size_t length):
cdef char bits
cdef size_t size = length
+ cdef size_t i
# the length of the variable prefix length will be less than 9 bytes
if self._output_buffer_size < self._output_pos + length + 9:
self._output_buffer_size += length + 9
@@ -91,12 +90,36 @@ cdef class BeamOutputStream(LengthPrefixOutputStream):
self._output_stream.flush()
self._output_pos = 0
- cdef void _parse_output_stream(self, BOutputStream output_stream):
+ cdef void reset_output_stream(self, BOutputStream output_stream):
+ self._output_stream = output_stream
self._output_data = output_stream.data
self._output_pos = output_stream.pos
self._output_buffer_size = output_stream.buffer_size
- cdef void _maybe_flush(self):
+ cdef bint _maybe_flush(self):
if self._output_pos > 10_000_000:
- self._output_stream.flush()
- self._output_pos = 0
+ self.flush()
+ return True
+ return False
+
+cdef class BeamTimeBasedOutputStream(BeamSizeBasedOutputStream):
+ def __init__(self, *args, **kwargs):
+ self._flush_event = False
+ self._periodic_flusher = PeriodicThread(1, self.notify_flush)
+ self._periodic_flusher.daemon = True
+ self._periodic_flusher.start()
+
+ cpdef void notify_flush(self):
+ self._flush_event = True
+
+ cpdef void close(self):
+ if self._periodic_flusher:
+ self._periodic_flusher.cancel()
+ self._periodic_flusher = None
+
+ cdef bint _maybe_flush(self):
+ if self._flush_event:
+ self.flush()
+ self._flush_event = False
+ elif BeamSizeBasedOutputStream._maybe_flush(self):
+ self._flush_event = False
diff --git a/flink-python/pyflink/fn_execution/beam/beam_stream_slow.py b/flink-python/pyflink/fn_execution/beam/beam_stream_slow.py
index 8dff781..5688beb 100644
--- a/flink-python/pyflink/fn_execution/beam/beam_stream_slow.py
+++ b/flink-python/pyflink/fn_execution/beam/beam_stream_slow.py
@@ -15,7 +15,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
-from apache_beam.coders.coder_impl import create_InputStream
+from apache_beam.coders.coder_impl import create_InputStream, create_OutputStream
+from apache_beam.runners.worker.data_plane import PeriodicThread
from pyflink.fn_execution.stream_slow import InputStream
@@ -33,3 +34,30 @@ class BeamInputStream(InputStream):
def size(self):
return self._input_stream.size()
+
+
+class BeamTimeBasedOutputStream(create_OutputStream):
+ def __init__(self):
+ super(BeamTimeBasedOutputStream).__init__()
+ self._flush_event = False
+ self._periodic_flusher = PeriodicThread(1, self.notify_flush)
+ self._periodic_flusher.daemon = True
+ self._periodic_flusher.start()
+ self._output_stream = None
+
+ def write(self, b: bytes):
+ self._output_stream.write(b)
+ if self._flush_event:
+ self._output_stream.flush()
+ self._flush_event = False
+
+ def reset_output_stream(self, output_stream: create_OutputStream):
+ self._output_stream = output_stream
+
+ def notify_flush(self):
+ self._flush_event = True
+
+ def close(self):
+ if self._periodic_flusher:
+ self._periodic_flusher.cancel()
+ self._periodic_flusher = None
diff --git a/flink-python/pyflink/fn_execution/coder_impl_fast.pxd b/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
index de6c373..5c3d6f6 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
+++ b/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
@@ -89,6 +89,19 @@ cdef class RowCoderImpl(FieldCoderImpl):
cdef list _field_names
cdef MaskUtils _mask_utils
+cdef class ArrowCoderImpl(FieldCoderImpl):
+ cdef object _schema
+ cdef list _field_types
+ cdef object _timezone
+ cdef object _resettable_io
+ cdef object _batch_reader
+
+ cdef list decode_one_batch_from_stream(self, InputStream in_stream, size_t size)
+
+cdef class OverWindowArrowCoderImpl(FieldCoderImpl):
+ cdef ArrowCoderImpl _arrow_coder
+ cdef IntCoderImpl _int_coder
+
cdef class TinyIntCoderImpl(FieldCoderImpl):
pass
diff --git a/flink-python/pyflink/fn_execution/coder_impl_fast.pyx b/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
index 7bf6f64..f66985d 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
+++ b/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
@@ -28,10 +28,13 @@ import pickle
from typing import List, Union
from cloudpickle import cloudpickle
+import pyarrow as pa
from pyflink.common import Row, RowKind
from pyflink.common.time import Instant
from pyflink.datastream.window import CountWindow, TimeWindow
+from pyflink.fn_execution.ResettableIO import ResettableIO
+from pyflink.table.utils import pandas_to_arrow, arrow_to_pandas
ROW_KIND_BIT_SIZE = 2
@@ -421,6 +424,73 @@ cdef class RowCoderImpl(FieldCoderImpl):
return 'RowCoderImpl[%s, %s]' % \
(', '.join(str(c) for c in self._field_coders), self._field_names)
+cdef class ArrowCoderImpl(FieldCoderImpl):
+ """
+ A coder for arrow format data.
+ """
+
+ def __init__(self, schema, row_type, timezone):
+ self._schema = schema
+ self._field_types = row_type.field_types()
+ self._timezone = timezone
+ self._resettable_io = ResettableIO()
+ self._batch_reader = self._load_from_stream(self._resettable_io)
+
+ cpdef encode_to_stream(self, cols, OutputStream out_stream):
+ self._resettable_io.set_output_stream(out_stream)
+ batch_writer = pa.RecordBatchStreamWriter(self._resettable_io, self._schema)
+ batch_writer.write_batch(
+ pandas_to_arrow(self._schema, self._timezone, self._field_types, cols))
+
+ cpdef decode_from_stream(self, InputStream in_stream, size_t size):
+ return self.decode_one_batch_from_stream(in_stream, size)
+
+ cdef list decode_one_batch_from_stream(self, InputStream in_stream, size_t size):
+ self._resettable_io.set_input_bytes(in_stream.read(size))
+ # there is only one arrow batch in the underlying input stream
+ return arrow_to_pandas(self._timezone, self._field_types, [next(self._batch_reader)])
+
+ def _load_from_stream(self, stream):
+ while stream.readable():
+ reader = pa.ipc.open_stream(stream)
+ yield reader.read_next_batch()
+
+ def __repr__(self):
+ return 'ArrowCoderImpl[%s]' % self._schema
+
+cdef class OverWindowArrowCoderImpl(FieldCoderImpl):
+ """
+ A coder for over window with arrow format data.
+ The data structure: [window data][arrow format data].
+ """
+ def __init__(self, arrow_coder_impl: ArrowCoderImpl):
+ self._arrow_coder = arrow_coder_impl
+ self._int_coder = IntCoderImpl()
+
+ cpdef encode_to_stream(self, cols, OutputStream out_stream):
+ self._arrow_coder.encode_to_stream(cols, out_stream)
+
+ cpdef decode_from_stream(self, InputStream in_stream, size_t size):
+ cdef int32_t window_num, window_size
+ cdef list window_boundaries_and_arrow_data
+ window_num = self._int_coder.decode_from_stream(in_stream, 0)
+ size -= 4
+ window_boundaries_and_arrow_data = []
+ for _ in range(window_num):
+ window_size = self._int_coder.decode_from_stream(in_stream, 0)
+ size -= 4
+ window_boundaries_and_arrow_data.append(
+ [self._int_coder.decode_from_stream(in_stream, 0)
+ for _ in range(window_size)])
+ size -= 4 * window_size
+ window_boundaries_and_arrow_data.append(
+ self._arrow_coder.decode_one_batch_from_stream(in_stream, size))
+ return window_boundaries_and_arrow_data
+
+ def __repr__(self):
+ return 'OverWindowArrowCoderImpl[%s]' % self._arrow_coder
+
+
cdef class TinyIntCoderImpl(FieldCoderImpl):
"""
A coder for tiny int value (from -128 to 127).
diff --git a/flink-python/pyflink/fn_execution/coder_impl_slow.py b/flink-python/pyflink/fn_execution/coder_impl_slow.py
index b9789b3..77aff11 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_slow.py
+++ b/flink-python/pyflink/fn_execution/coder_impl_slow.py
@@ -289,17 +289,17 @@ class ArrowCoderImpl(FieldCoderImpl):
def decode_from_stream(self, in_stream: InputStream, length=0):
return self.decode_one_batch_from_stream(in_stream, length)
+ def decode_one_batch_from_stream(self, in_stream: InputStream, size: int) -> List:
+ self._resettable_io.set_input_bytes(in_stream.read(size))
+ # there is only one arrow batch in the underlying input stream
+ return arrow_to_pandas(self._timezone, self._field_types, [next(self._batch_reader)])
+
@staticmethod
def _load_from_stream(stream):
while stream.readable():
reader = pa.ipc.open_stream(stream)
yield reader.read_next_batch()
- def decode_one_batch_from_stream(self, in_stream: InputStream, size: int) -> List:
- self._resettable_io.set_input_bytes(in_stream.read(size))
- # there is only one arrow batch in the underlying input stream
- return arrow_to_pandas(self._timezone, self._field_types, [next(self._batch_reader)])
-
def __repr__(self):
return 'ArrowCoderImpl[%s]' % self._schema
diff --git a/flink-python/pyflink/fn_execution/coders.py b/flink-python/pyflink/fn_execution/coders.py
index 7ba4aed..8ef5607 100644
--- a/flink-python/pyflink/fn_execution/coders.py
+++ b/flink-python/pyflink/fn_execution/coders.py
@@ -160,14 +160,8 @@ class IterableCoder(LengthPrefixBaseCoder):
self._separated_with_end_message = separated_with_end_message
def get_impl(self):
- if isinstance(self._field_coder, (ArrowCoder, OverWindowArrowCoder)):
- # ArrowCoder and OverWindowArrowCoder doesn't support fast coder currently.
- from pyflink.fn_execution import coder_impl_slow
- return coder_impl_slow.IterableCoderImpl(self._field_coder.get_impl(),
- self._separated_with_end_message)
- else:
- return coder_impl.IterableCoderImpl(self._field_coder.get_impl(),
- self._separated_with_end_message)
+ return coder_impl.IterableCoderImpl(self._field_coder.get_impl(),
+ self._separated_with_end_message)
class ValueCoder(LengthPrefixBaseCoder):
@@ -179,12 +173,7 @@ class ValueCoder(LengthPrefixBaseCoder):
super(ValueCoder, self).__init__(field_coder)
def get_impl(self):
- if isinstance(self._field_coder, (ArrowCoder, OverWindowArrowCoder)):
- # ArrowCoder and OverWindowArrowCoder doesn't support fast coder currently.
- from pyflink.fn_execution import coder_impl_slow
- return coder_impl_slow.ValueCoderImpl(self._field_coder.get_impl())
- else:
- return coder_impl.ValueCoderImpl(self._field_coder.get_impl())
+ return coder_impl.ValueCoderImpl(self._field_coder.get_impl())
#########################################################################
@@ -240,9 +229,7 @@ class ArrowCoder(FieldCoder):
self._timezone = timezone
def get_impl(self):
- # ArrowCoder doesn't support fast coder implementation currently.
- from pyflink.fn_execution import coder_impl_slow
- return coder_impl_slow.ArrowCoderImpl(self._schema, self._row_type, self._timezone)
+ return coder_impl.ArrowCoderImpl(self._schema, self._row_type, self._timezone)
def __repr__(self):
return 'ArrowCoder[%s]' % self._schema
@@ -257,9 +244,7 @@ class OverWindowArrowCoder(FieldCoder):
self._arrow_coder = ArrowCoder(schema, row_type, timezone)
def get_impl(self):
- # OverWindowArrowCoder doesn't support fast coder implementation currently.
- from pyflink.fn_execution import coder_impl_slow
- return coder_impl_slow.OverWindowArrowCoderImpl(self._arrow_coder.get_impl())
+ return coder_impl.OverWindowArrowCoderImpl(self._arrow_coder.get_impl())
def __repr__(self):
return 'OverWindowArrowCoder[%s]' % self._arrow_coder
diff --git a/flink-python/pyflink/fn_execution/stream_fast.pxd b/flink-python/pyflink/fn_execution/stream_fast.pxd
index ff9151b..110ca7d 100644
--- a/flink-python/pyflink/fn_execution/stream_fast.pxd
+++ b/flink-python/pyflink/fn_execution/stream_fast.pxd
@@ -25,11 +25,13 @@ cdef class LengthPrefixInputStream:
cdef class LengthPrefixOutputStream:
cdef void write(self, char*data, size_t length)
cpdef void flush(self)
+ cpdef void close(self)
cdef class InputStream:
cdef char*_input_data
cdef size_t _input_pos
+ cpdef bytes read(self, size_t size)
cdef long read_byte(self) except? -1
cdef int8_t read_int8(self) except? -1
cdef int16_t read_int16(self) except? -1
@@ -44,6 +46,7 @@ cdef class OutputStream:
cdef size_t buffer_size
cdef size_t pos
+ cpdef void write(self, bytes v)
cdef void write_byte(self, unsigned char val)
cdef void write_int8(self, int8_t v)
cdef void write_int16(self, int16_t v)
diff --git a/flink-python/pyflink/fn_execution/stream_fast.pyx b/flink-python/pyflink/fn_execution/stream_fast.pyx
index 3ef5225..883f8e6 100644
--- a/flink-python/pyflink/fn_execution/stream_fast.pyx
+++ b/flink-python/pyflink/fn_execution/stream_fast.pyx
@@ -31,11 +31,17 @@ cdef class LengthPrefixOutputStream:
pass
cpdef void flush(self):
pass
+ cpdef void close(self):
+ pass
cdef class InputStream:
def __init__(self):
self._input_pos = 0
+ cpdef bytes read(self, size_t size):
+ self._input_pos += size
+ return self._input_data[self._input_pos - size: self._input_pos]
+
cdef long read_byte(self) except? -1:
self._input_pos += 1
return <long> <unsigned char> self._input_data[self._input_pos - 1]
@@ -94,6 +100,21 @@ cdef class OutputStream:
self.buffer_size = 1024
self.pos = 0
+ cpdef void write(self, bytes v):
+ cdef size_t length
+ cdef char*b
+ length = len(v)
+ b = <char*> v
+ if self.buffer_size < self.pos + length:
+ self._extend(length)
+ if length < 8:
+ # This is faster than memcpy when the string is short.
+ for i in range(length):
+ self.buffer[self.pos + i] = b[i]
+ else:
+ memcpy(self.buffer + self.pos, b, length)
+ self.pos += length
+
cdef void write_byte(self, unsigned char v):
if self.buffer_size < self.pos + 1:
self._extend(1)
diff --git a/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java
index 7123cee..710a046 100644
--- a/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java
@@ -53,6 +53,16 @@ public interface PythonFunctionRunner {
Tuple2<byte[], Integer> pollResult() throws Exception;
/**
+ * Retrieves the Python function result, waiting if necessary until an element becomes
+ * available.
+ *
+ * @return the head of he Python function result buffer. f0 means the byte array buffer which
+ * stores the Python function result. f1 means the length of the Python function result byte
+ * array.
+ */
+ Tuple2<byte[], Integer> takeResult() throws Exception;
+
+ /**
* Forces to finish the processing of the current bundle of elements. It will flush the data
* cached in the data buffer for processing and retrieves the state mutations (if exists) made
* by the Python function. The call blocks until all of the outputs produced by this bundle have
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java
index 7abb6c1..895d511 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java
@@ -33,6 +33,7 @@ import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.ChainingStrategy;
import org.apache.flink.streaming.api.operators.sorted.state.BatchExecutionInternalTimeServiceManager;
import org.apache.flink.streaming.api.operators.sorted.state.BatchExecutionKeyedStateBackend;
+import org.apache.flink.streaming.api.runners.python.beam.BeamPythonFunctionRunner;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.table.functions.python.PythonEnv;
import org.apache.flink.util.Preconditions;
@@ -43,7 +44,11 @@ import java.lang.reflect.Field;
import java.util.HashMap;
import java.util.List;
import java.util.Objects;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import static org.apache.flink.streaming.api.utils.ClassLeakCleaner.cleanUpLeakingClasses;
import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.inBatchExecutionMode;
@@ -78,6 +83,8 @@ public abstract class AbstractPythonFunctionOperator<OUT> extends AbstractStream
/** Callback to be executed after the current bundle was finished. */
private transient Runnable bundleFinishedCallback;
+ private transient ExecutorService flushThreadPool;
+
/** The python config. */
private PythonConfig config;
@@ -133,6 +140,7 @@ public abstract class AbstractPythonFunctionOperator<OUT> extends AbstractStream
timestamp -> checkInvokeFinishBundleByTime(),
bundleCheckPeriod,
bundleCheckPeriod);
+ this.flushThreadPool = Executors.newSingleThreadExecutor();
} finally {
super.open();
}
@@ -158,6 +166,10 @@ public abstract class AbstractPythonFunctionOperator<OUT> extends AbstractStream
pythonFunctionRunner.close();
pythonFunctionRunner = null;
}
+ if (flushThreadPool != null) {
+ flushThreadPool.shutdown();
+ flushThreadPool = null;
+ }
} finally {
super.close();
@@ -295,7 +307,7 @@ public abstract class AbstractPythonFunctionOperator<OUT> extends AbstractStream
protected void emitResults() throws Exception {
Tuple2<byte[], Integer> resultTuple;
- while ((resultTuple = pythonFunctionRunner.pollResult()) != null) {
+ while ((resultTuple = pythonFunctionRunner.pollResult()) != null && resultTuple.f1 != 0) {
emitResult(resultTuple);
}
}
@@ -317,9 +329,37 @@ public abstract class AbstractPythonFunctionOperator<OUT> extends AbstractStream
protected void invokeFinishBundle() throws Exception {
if (elementCount > 0) {
- pythonFunctionRunner.flush();
- elementCount = 0;
+ AtomicBoolean flushThreadFinish = new AtomicBoolean(false);
+ AtomicReference<Throwable> exceptionReference = new AtomicReference<>();
+ flushThreadPool.submit(
+ () -> {
+ try {
+ pythonFunctionRunner.flush();
+ } catch (Throwable e) {
+ exceptionReference.set(e);
+ } finally {
+ flushThreadFinish.set(true);
+ // interrupt the progress of takeResult to avoid the main thread is
+ // blocked forever.
+ ((BeamPythonFunctionRunner) pythonFunctionRunner).notifyNoMoreResults();
+ }
+ });
+ Tuple2<byte[], Integer> resultTuple;
+ while (!flushThreadFinish.get()) {
+ resultTuple = pythonFunctionRunner.takeResult();
+ if (resultTuple.f1 != 0) {
+ emitResult(resultTuple);
+ emitResults();
+ }
+ }
emitResults();
+ Throwable flushThreadThrowable = exceptionReference.get();
+ if (flushThreadThrowable != null) {
+ throw new RuntimeException(
+ "Error while waiting for BeamPythonFunctionRunner flush",
+ flushThreadThrowable);
+ }
+ elementCount = 0;
lastFinishBundleTime = getProcessingTimeService().getCurrentProcessingTime();
// callback only after current bundle was fully finalized
if (bundleFinishedCallback != null) {
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
index 0e6fbba..ed8793c 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
@@ -347,6 +347,14 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
}
@Override
+ public Tuple2<byte[], Integer> takeResult() throws Exception {
+ byte[] result = resultBuffer.take();
+ this.reusableResultTuple.f0 = result;
+ this.reusableResultTuple.f1 = result.length;
+ return this.reusableResultTuple;
+ }
+
+ @Override
public void flush() throws Exception {
if (bundleStarted) {
try {
@@ -357,6 +365,11 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
}
}
+ /** Interrupts the progress of takeResult. */
+ public void notifyNoMoreResults() {
+ resultBuffer.add(new byte[0]);
+ }
+
private void finishBundle() {
try {
remoteBundle.close();