| ################################################################################ |
| # 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. |
| ################################################################################ |
| # cython: language_level = 3 |
| # cython: infer_types = True |
| # cython: profile=True |
| # cython: boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True |
| |
| from libc.stdlib cimport realloc |
| from libc.string cimport memcpy |
| |
| from pyflink.fn_execution.utils.operation_utils import PeriodicThread |
| |
| cdef class BeamInputStream(LengthPrefixInputStream): |
| def __cinit__(self, input_stream, size): |
| self._input_buffer_size = size |
| self._input_stream = input_stream |
| self._parse_input_stream(input_stream) |
| |
| cdef size_t read(self, char** data): |
| cdef size_t length = 0 |
| cdef bint has_prefix = True |
| cdef size_t shift = 0 |
| cdef char bits |
| # read the var-int size |
| while has_prefix: |
| bits = self._input_data[self._input_pos] & 0x7F |
| length |= bits << shift |
| shift += 7 |
| if not (self._input_data[self._input_pos] & 0x80): |
| has_prefix = False |
| self._input_pos += 1 |
| data[0] = self._input_data + self._input_pos |
| self._input_pos += length |
| self._input_stream.pos = self._input_pos |
| return length |
| |
| cdef size_t available(self): |
| return self._input_buffer_size - self._input_pos |
| |
| cdef void _parse_input_stream(self, BInputStream input_stream): |
| self._input_data = input_stream.allc |
| self._input_pos = input_stream.pos |
| |
| 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 |
| self._output_data = <char*> realloc(self._output_data, self._output_buffer_size) |
| self._output_stream.buffer_size = self._output_buffer_size |
| self._output_stream.data = self._output_data |
| # write variable prefix length |
| while size: |
| bits = size & 0x7F |
| size >>= 7 |
| if size: |
| bits |= 0x80 |
| self._output_data[self._output_pos] = bits |
| self._output_pos += 1 |
| |
| if length < 8: |
| # This is faster than memcpy when the string is short. |
| for i in range(length): |
| self._output_data[self._output_pos + i] = data[i] |
| else: |
| memcpy(self._output_data + self._output_pos, data, length) |
| self._output_pos += length |
| self._output_stream.pos = self._output_pos |
| |
| cpdef void flush(self): |
| self._output_stream.flush() |
| self._output_pos = 0 |
| |
| 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 |
| |
| cpdef bint maybe_flush(self): |
| if self._output_pos > 10_000_000: |
| 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 |
| |
| cpdef bint maybe_flush(self): |
| if self._flush_event: |
| self.flush() |
| self._flush_event = False |
| elif BeamSizeBasedOutputStream.maybe_flush(self): |
| self._flush_event = False |
