sdks/python/apache_beam/runners/worker/operations.pxd

#
# 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.
#

cimport cython

from apache_beam.runners.common cimport DoFnRunner
from apache_beam.runners.common cimport Receiver
from apache_beam.runners.worker cimport opcounters
from apache_beam.utils.windowed_value cimport WindowedValue
from apache_beam.utils.windowed_value cimport WindowedBatch
#from libcpp.string cimport string

cdef WindowedValue _globally_windowed_value
cdef type _global_window_type


cdef class ConsumerSet(Receiver):
  cdef list consumers
  cdef readonly opcounters.OperationCounters opcounter
  cdef public step_name
  cdef public output_index
  cdef public coder
  cdef public object output_sampler
  cdef public object element_sampler
  cdef public object execution_context

  cpdef update_counters_start(self, WindowedValue windowed_value)
  cpdef update_counters_finish(self)
  cpdef update_counters_batch(self, WindowedBatch windowed_batch)

cdef class SingletonElementConsumerSet(ConsumerSet):
  cdef Operation consumer

  cpdef receive(self, WindowedValue windowed_value)
  cpdef receive_batch(self, WindowedBatch windowed_batch)
  cpdef flush(self)

cdef class GeneralPurposeConsumerSet(ConsumerSet):
  cdef list element_consumers
  cdef list passthrough_batch_consumers
  cdef dict other_batch_consumers
  cdef bint has_batch_consumers
  cdef list _batched_elements
  cdef object producer_batch_converter

  cpdef receive(self, WindowedValue windowed_value)
  cpdef receive_batch(self, WindowedBatch windowed_batch)
  cpdef flush(self)

cdef class Operation(object):
  cdef readonly name_context
  cdef readonly operation_name
  cdef readonly spec
  cdef object consumers
  cdef readonly counter_factory
  cdef public metrics_container
  cdef public execution_context
  # Public for access by Fn harness operations.
  # TODO(robertwb): Cythonize FnHarness.
  cdef public list receivers
  cdef readonly bint debug_logging_enabled
  # For legacy workers.
  cdef bint setup_done

  cdef public step_name  # initialized lazily

  cdef readonly object state_sampler

  cdef readonly object scoped_start_state
  cdef readonly object scoped_process_state
  cdef readonly object scoped_finish_state

  cdef readonly object data_sampler

  cpdef start(self)
  cpdef process(self, WindowedValue windowed_value)
  cpdef finish(self)
  cpdef teardown(self)
  cpdef output(self, WindowedValue windowed_value, int output_index=*)
  cpdef execution_time_monitoring_infos(self, transform_id)
  cpdef user_monitoring_infos(self, transform_id)
  cpdef pcollection_count_monitoring_infos(self, tag_to_pcollection_id)
  cpdef monitoring_infos(self, transform_id, tag_to_pcollection_id)


cdef class ReadOperation(Operation):
  @cython.locals(windowed_value=WindowedValue)
  cpdef start(self)


cdef class ImpulseReadOperation(Operation):
  cdef object source
  @cython.locals(windowed_value=WindowedValue)
  cpdef process(self, WindowedValue impulse)


cdef class DoOperation(Operation):
  cdef DoFnRunner dofn_runner
  cdef object tagged_receivers
  cdef object side_input_maps
  cdef object user_state_context
  cdef public dict timer_inputs
  cdef dict timer_specs
  cdef public object input_info
  cdef object fn
  cdef object scoped_timer_processing_state


cdef class SdfProcessSizedElements(DoOperation):
  cdef object lock
  cdef object element_start_output_bytes

cdef class SdfTruncateSizedRestrictions(DoOperation):
  pass

cdef class CombineOperation(Operation):
  cdef object phased_combine_fn


cdef class PGBKCVOperation(Operation):
  cdef public object combine_fn
  cdef public object combine_fn_add_input
  cdef public object combine_fn_compact
  cdef public bint is_default_windowing
  cdef public object timestamp_combiner
  cdef dict table
  cdef long max_keys
  cdef long key_count

  cpdef add_key_value(self, wkey, value, timestamp)
  cpdef output_key(self, wkey, value, timestamp)


cdef class FlattenOperation(Operation):
  pass