gigl.common.beam.RecordBatchToTFExampleCoderFn

class gigl.common.beam.coders.RecordBatchToTFExampleCoderFn

Bases: DoFn

Encode pyarrow.RecordBatch to serialized tf.train.Example(s)

Methods

__init__
default_label
default_type_hints
display_data	Returns the display data associated to a pipeline component.
finish_bundle	Called after a bundle of elements is processed on a worker.
from_callable
from_runner_api	Converts from an FunctionSpec to a Fn object.
get_function_arguments
get_input_batch_type	Determine the batch type expected as input to process_batch.
get_output_batch_type	Determine the batch type produced by this DoFn's process_batch implementation and/or its process implementation with @yields_batch.
get_type_hints	Gets and/or initializes type hints for this object.
infer_output_type
process	Note that transformed_metadata needs to be passed in as side input, i.e., as an argument of process function, instead of being passed to class init, since it could potentially materialize (depending on whether it is read from file or built by tft_beam.AnalyzeDataset) after the class is constructed.
process_batch
register_pickle_urn	Registers and implements the given urn via pickling.
register_urn	Registers a urn with a constructor.
setup	Called to prepare an instance for processing bundles of elements.
start_bundle	Called before a bundle of elements is processed on a worker.
teardown	Called to use to clean up this instance before it is discarded.
to_runner_api	Returns an FunctionSpec encoding this Fn.
to_runner_api_parameter	Returns the urn and payload for this Fn.
unbounded_per_element	A decorator on process fn specifying that the fn performs an unbounded amount of work per input element.
with_input_types
with_output_types
yields_batches	A decorator to apply to process indicating it yields batches.
yields_elements	A decorator to apply to process_batch indicating it yields elements.

BundleFinalizerParam

alias of _BundleFinalizerParam

RestrictionParam

alias of _RestrictionDoFnParam

StateParam

alias of _StateDoFnParam

TimerParam

alias of _TimerDoFnParam

WatermarkEstimatorParam

alias of _WatermarkEstimatorParam

__init__()

__weakref__

list of weak references to the object (if defined)

display_data() → dict

Returns the display data associated to a pipeline component.

It should be reimplemented in pipeline components that wish to have static display data.

Returns:

Dict[str, Any]: A dictionary containing key:value pairs. The value might be an integer, float or string value; a DisplayDataItem for values that have more data (e.g. short value, label, url); or a HasDisplayData instance that has more display data that should be picked up. For example:

{
  'key1': 'string_value',
  'key2': 1234,
  'key3': 3.14159265,
  'key4': DisplayDataItem('apache.org', url='http://apache.org'),
  'key5': subComponent
}

finish_bundle()

Called after a bundle of elements is processed on a worker.

classmethod from_runner_api(fn_proto: Type[RunnerApiFnT], context: beam_runner_api_pb2.FunctionSpec) → RunnerApiFnT

Converts from an FunctionSpec to a Fn object.

Prefer registering a urn with its parameter type and constructor.

get_input_batch_type(input_element_type) → TypeConstraint | type | None

Determine the batch type expected as input to process_batch.

The default implementation of get_input_batch_type simply observes the input typehint for the first parameter of process_batch. A Batched DoFn may override this method if a dynamic approach is required.

Args:

input_element_type: The element type of the input PCollection this

DoFn is being applied to.

Returns:

None if this DoFn cannot accept batches, else a Beam typehint or a native Python typehint.

get_output_batch_type(input_element_type) → TypeConstraint | type | None

Determine the batch type produced by this DoFn’s process_batch implementation and/or its process implementation with @yields_batch.

The default implementation of this method observes the return type annotations on process_batch and/or process. A Batched DoFn may override this method if a dynamic approach is required.

Args:

input_element_type: The element type of the input PCollection this

DoFn is being applied to.

Returns:

None if this DoFn will never yield batches, else a Beam typehint or a native Python typehint.

get_type_hints()

Gets and/or initializes type hints for this object.

If type hints have not been set, attempts to initialize type hints in this order: - Using self.default_type_hints(). - Using self.__class__ type hints.

process(element: RecordBatch, transformed_metadata: DatasetMetadata, *args, **kwargs) → Iterable[bytes]

Note that transformed_metadata needs to be passed in as side input, i.e., as an argument of process function, instead of being passed to class init, since it could potentially materialize (depending on whether it is read from file or built by tft_beam.AnalyzeDataset) after the class is constructed.

Args:

element (pa.RecordBatch): A batch of records, e.g., a batch of transformed features transformed_metadata (tensorflow_transform.tf_metadata.dataset_metadata.DatasetMetadata):

containing the schema needed by RecordBatchToExamplesEncoder for encoding

Yields:

bytes: serialized tf.Example

classmethod register_pickle_urn(pickle_urn)

Registers and implements the given urn via pickling.

classmethod register_urn(urn, parameter_type, fn=None)

Registers a urn with a constructor.

For example, if ‘beam:fn:foo’ had parameter type FooPayload, one could write RunnerApiFn.register_urn(‘bean:fn:foo’, FooPayload, foo_from_proto) where foo_from_proto took as arguments a FooPayload and a PipelineContext. This function can also be used as a decorator rather than passing the callable in as the final parameter.

A corresponding to_runner_api_parameter method would be expected that returns the tuple (‘beam:fn:foo’, FooPayload)

setup()

Called to prepare an instance for processing bundles of elements.

This is a good place to initialize transient in-memory resources, such as network connections. The resources can then be disposed in DoFn.teardown.

start_bundle()

Called before a bundle of elements is processed on a worker.

Elements to be processed are split into bundles and distributed to workers. Before a worker calls process() on the first element of its bundle, it calls this method.

teardown()

Called to use to clean up this instance before it is discarded.

A runner will do its best to call this method on any given instance to prevent leaks of transient resources, however, there may be situations where this is impossible (e.g. process crash, hardware failure, etc.) or unnecessary (e.g. the pipeline is shutting down and the process is about to be killed anyway, so all transient resources will be released automatically by the OS). In these cases, the call may not happen. It will also not be retried, because in such situations the DoFn instance no longer exists, so there’s no instance to retry it on.

Thus, all work that depends on input elements, and all externally important side effects, must be performed in DoFn.process or DoFn.finish_bundle.

to_runner_api(context: PipelineContext) → beam_runner_api_pb2.FunctionSpec

Returns an FunctionSpec encoding this Fn.

Prefer overriding self.to_runner_api_parameter.

to_runner_api_parameter(context)

Returns the urn and payload for this Fn.

The returned urn(s) should be registered with register_urn.

static unbounded_per_element()

A decorator on process fn specifying that the fn performs an unbounded amount of work per input element.

static yields_batches(fn)

A decorator to apply to process indicating it yields batches.

By default process is assumed to both consume and produce individual elements at a time. This decorator indicates that process produces “batches”, which are collections of multiple logical Beam elements.

static yields_elements(fn)

A decorator to apply to process_batch indicating it yields elements.

By default process_batch is assumed to both consume and produce “batches”, which are collections of multiple logical Beam elements. This decorator indicates that process_batch produces individual elements at a time. process_batch is always expected to consume batches.