tfma.extractors.BatchedInputExtractor

Creates an extractor for extracting features, labels, and example weights.

The extractor's PTransform extracts features, labels, and example weights from the batched features (i.e., Arrow RecordBatch) stored under tfma.ARROW_RECORD_BATCH_KEY in the incoming extract and adds it to the output extract under the keys tfma.FEATURES_KEY, tfma.LABELS_KEY, and tfma.EXAMPLE_WEIGHTS_KEY. If the eval_config contains a prediction_key and a corresponding key is found in the parse example, then predictions will also be extracted and stored under the tfma.PREDICTIONS_KEY. Any extracts that already exist will be merged with the values parsed by this extractor with this extractor's values taking precedence when duplicate keys are detected.

Note that the use of a prediction_key in an eval_config serves two use cases: (1) as a key into the dict of predictions output by predict extractor (2) as the key for a pre-computed prediction stored as a feature. The InputExtractor can be used to handle case (2). These cases are meant to be exclusive (i.e. if approach (2) is used then a predict extractor would not be configured and if (1) is used then a key matching the predictons would not be stored in the features). However, if a feature key happens to match the same name as the prediction output key then both paths may be executed. In this case, the value stored here will be replaced by the predict extractor (though it will still be popped from the features).

eval_config Eval config.

Batched extractor for extracting features, labels, and example weights.