def _RunBeamImpl(self, inputs, outputs, preprocessing_fn,
input_dataset_metadata, raw_examples_data_format,
transform_output_path, compute_statistics,
materialize_output_paths):
"""Perform data preprocessing with FlumeC++ runner.
Args:
inputs: A dictionary of labelled input values.
outputs: A dictionary of labelled output values.
preprocessing_fn: The tf.Transform preprocessing_fn.
input_dataset_metadata: A DatasetMetadata object for the input data.
raw_examples_data_format: A string describing the raw data format.
transform_output_path: An absolute path to write the output to.
compute_statistics: A bool indicating whether or not compute statistics.
materialize_output_paths: Paths to materialized outputs.
Raises:
RuntimeError: If reset() is not being invoked between two run().
ValueError: If the schema is empty.
Returns:
Status of the execution.
"""
raw_examples_file_format = common.GetSoleValue(
inputs, labels.EXAMPLES_FILE_FORMAT_LABEL, strict=False)
analyze_and_transform_data_paths = common.GetValues(
inputs, labels.ANALYZE_AND_TRANSFORM_DATA_PATHS_LABEL)
transform_only_data_paths = common.GetValues(
inputs, labels.TRANSFORM_ONLY_DATA_PATHS_LABEL)
stats_use_tfdv = common.GetSoleValue(
inputs, labels.TFT_STATISTICS_USE_TFDV_LABEL)
per_set_stats_output_paths = common.GetValues(
outputs, labels.PER_SET_STATS_OUTPUT_PATHS_LABEL)
temp_path = common.GetSoleValue(outputs, labels.TEMP_OUTPUT_LABEL)
tf.logging.info('Analyze and transform data patterns: %s',
list(enumerate(analyze_and_transform_data_paths)))
tf.logging.info('Transform data patterns: %s',
list(enumerate(transform_only_data_paths)))
tf.logging.info('Transform materialization output paths: %s',
list(enumerate(materialize_output_paths)))
tf.logging.info('Transform output path: %s', transform_output_path)
feature_spec = input_dataset_metadata.schema.as_feature_spec()
try:
analyze_input_columns = tft.get_analyze_input_columns(
preprocessing_fn, feature_spec)
transform_input_columns = (tft.get_transform_input_columns(
preprocessing_fn, feature_spec))
except AttributeError:
# If using TFT 1.12, fall back to assuming all features are used.
analyze_input_columns = feature_spec.keys()
transform_input_columns = feature_spec.keys()
# Use the same dataset (same columns) for AnalyzeDataset and computing
# pre-transform stats so that the data will only be read once for these
# two operations.
if compute_statistics:
analyze_input_columns = list(
set(
list(analyze_input_columns) +
list(transform_input_columns)))
analyze_input_dataset_metadata = copy.deepcopy(input_dataset_metadata)
transform_input_dataset_metadata = copy.deepcopy(
input_dataset_metadata)
if input_dataset_metadata.schema is not _RAW_EXAMPLE_SCHEMA:
analyze_input_dataset_metadata.schema = dataset_schema.from_feature_spec(
{
feature: feature_spec[feature]
for feature in analyze_input_columns
})
transform_input_dataset_metadata.schema = (
dataset_schema.from_feature_spec({
feature: feature_spec[feature]
for feature in transform_input_columns
}))
can_process_jointly = not bool(per_set_stats_output_paths
or materialize_output_paths)
analyze_data_list = self._MakeDatasetList(
analyze_and_transform_data_paths, raw_examples_file_format,
raw_examples_data_format, analyze_input_dataset_metadata,
can_process_jointly)
transform_data_list = self._MakeDatasetList(
list(analyze_and_transform_data_paths) +
list(transform_only_data_paths), raw_examples_file_format,
raw_examples_data_format, transform_input_dataset_metadata,
can_process_jointly)
desired_batch_size = self._GetDesiredBatchSize(
raw_examples_data_format)
with self._CreatePipeline(outputs) as p:
with tft_beam.Context(
temp_dir=temp_path,
desired_batch_size=desired_batch_size,
passthrough_keys={_TRANSFORM_INTERNAL_FEATURE_FOR_KEY},
use_deep_copy_optimization=True):
# pylint: disable=expression-not-assigned
# pylint: disable=no-value-for-parameter
analyze_decode_fn = (self._GetDecodeFunction(
raw_examples_data_format,
analyze_input_dataset_metadata.schema))
for (idx, dataset) in enumerate(analyze_data_list):
dataset.encoded = (p
| 'ReadAnalysisDataset[{}]'.format(idx)
>> self._ReadExamples(dataset))
dataset.decoded = (
dataset.encoded
| 'DecodeAnalysisDataset[{}]'.format(idx) >>
self._DecodeInputs(analyze_decode_fn))
input_analysis_data = (
[dataset.decoded for dataset in analyze_data_list]
| 'FlattenAnalysisDatasets' >> beam.Flatten())
transform_fn = ((input_analysis_data, input_dataset_metadata)
| 'AnalyzeDataset' >>
tft_beam.AnalyzeDataset(preprocessing_fn))
# Write the raw/input metadata.
(input_dataset_metadata
| 'WriteMetadata' >> tft_beam.WriteMetadata(
os.path.join(transform_output_path,
tft.TFTransformOutput.RAW_METADATA_DIR), p))
# WriteTransformFn writes transform_fn and metadata to subdirectories
# tensorflow_transform.SAVED_MODEL_DIR and
# tensorflow_transform.TRANSFORMED_METADATA_DIR respectively.
(transform_fn | 'WriteTransformFn' >>
tft_beam.WriteTransformFn(transform_output_path))
if compute_statistics or materialize_output_paths:
# Do not compute pre-transform stats if the input format is raw proto,
# as StatsGen would treat any input as tf.Example.
if (compute_statistics and not self._IsDataFormatProto(
raw_examples_data_format)):
# Aggregated feature stats before transformation.
pre_transform_feature_stats_path = os.path.join(
transform_output_path, tft.TFTransformOutput.
PRE_TRANSFORM_FEATURE_STATS_PATH)
# TODO(b/70392441): Retain tf.Metadata (e.g., IntDomain) in
# schema. Currently input dataset schema only contains dtypes,
# and other metadata is dropped due to roundtrip to tensors.
schema_proto = schema_utils.schema_from_feature_spec(
analyze_input_dataset_metadata.schema.
as_feature_spec())
([
dataset.decoded
if stats_use_tfdv else dataset.encoded
for dataset in analyze_data_list
]
| 'FlattenPreTransformAnalysisDatasets' >>
beam.Flatten()
| 'GenerateAggregatePreTransformAnalysisStats' >>
self._GenerateStats(pre_transform_feature_stats_path,
schema_proto,
use_deep_copy_optimization=True,
use_tfdv=stats_use_tfdv))
transform_decode_fn = (self._GetDecodeFunction(
raw_examples_data_format,
transform_input_dataset_metadata.schema))
# transform_data_list is a superset of analyze_data_list, we pay the
# cost to read the same dataset (analyze_data_list) again here to
# prevent certain beam runner from doing large temp materialization.
for (idx, dataset) in enumerate(transform_data_list):
dataset.encoded = (
p
| 'ReadTransformDataset[{}]'.format(idx) >>
self._ReadExamples(dataset))
dataset.decoded = (
dataset.encoded
| 'DecodeTransformDataset[{}]'.format(idx) >>
self._DecodeInputs(transform_decode_fn))
(dataset.transformed, metadata) = (
((dataset.decoded,
transform_input_dataset_metadata), transform_fn)
| 'TransformDataset[{}]'.format(idx) >>
tft_beam.TransformDataset())
if materialize_output_paths or not stats_use_tfdv:
dataset.transformed_and_encoded = (
dataset.transformed
| 'EncodeTransformedDataset[{}]'.format(idx) >>
beam.ParDo(self._EncodeAsExamples(), metadata))
if compute_statistics:
# Aggregated feature stats after transformation.
_, metadata = transform_fn
post_transform_feature_stats_path = os.path.join(
transform_output_path, tft.TFTransformOutput.
POST_TRANSFORM_FEATURE_STATS_PATH)
# TODO(b/70392441): Retain tf.Metadata (e.g., IntDomain) in
# schema. Currently input dataset schema only contains dtypes,
# and other metadata is dropped due to roundtrip to tensors.
transformed_schema_proto = schema_utils.schema_from_feature_spec(
metadata.schema.as_feature_spec())
([(dataset.transformed if stats_use_tfdv else
dataset.transformed_and_encoded)
for dataset in transform_data_list]
| 'FlattenPostTransformAnalysisDatasets' >>
beam.Flatten()
| 'GenerateAggregatePostTransformAnalysisStats' >>
self._GenerateStats(post_transform_feature_stats_path,
transformed_schema_proto,
use_tfdv=stats_use_tfdv))
if per_set_stats_output_paths:
assert len(transform_data_list) == len(
per_set_stats_output_paths)
# TODO(b/67632871): Remove duplicate stats gen compute that is
# done both on a flattened view of the data, and on each span
# below.
bundles = zip(transform_data_list,
per_set_stats_output_paths)
for (idx, (dataset,
output_path)) in enumerate(bundles):
if stats_use_tfdv:
data = dataset.transformed
else:
data = dataset.transformed_and_encoded
(data
| 'GeneratePostTransformStats[{}]'.format(idx)
>> self._GenerateStats(
output_path,
transformed_schema_proto,
use_tfdv=stats_use_tfdv))
if materialize_output_paths:
assert len(transform_data_list) == len(
materialize_output_paths)
bundles = zip(transform_data_list,
materialize_output_paths)
for (idx, (dataset,
output_path)) in enumerate(bundles):
(dataset.transformed_and_encoded
| 'Materialize[{}]'.format(idx) >>
self._WriteExamples(raw_examples_file_format,
output_path))
return _Status.OK()
def _RunBeamImpl(self, inputs: Mapping[Text, Any],
outputs: Mapping[Text, Any], preprocessing_fn: Any,
input_dataset_metadata: dataset_metadata.DatasetMetadata,
raw_examples_data_format: Text, transform_output_path: Text,
compute_statistics: bool,
materialize_output_paths: Sequence[Text]) -> _Status:
"""Perform data preprocessing with FlumeC++ runner.
Args:
inputs: A dictionary of labelled input values.
outputs: A dictionary of labelled output values.
preprocessing_fn: The tf.Transform preprocessing_fn.
input_dataset_metadata: A DatasetMetadata object for the input data.
raw_examples_data_format: A string describing the raw data format.
transform_output_path: An absolute path to write the output to.
compute_statistics: A bool indicating whether or not compute statistics.
materialize_output_paths: Paths to materialized outputs.
Raises:
RuntimeError: If reset() is not being invoked between two run().
ValueError: If the schema is empty.
Returns:
Status of the execution.
"""
raw_examples_file_format = common.GetSoleValue(
inputs, labels.EXAMPLES_FILE_FORMAT_LABEL, strict=False)
analyze_and_transform_data_paths = common.GetValues(
inputs, labels.ANALYZE_AND_TRANSFORM_DATA_PATHS_LABEL)
transform_only_data_paths = common.GetValues(
inputs, labels.TRANSFORM_ONLY_DATA_PATHS_LABEL)
stats_use_tfdv = common.GetSoleValue(inputs,
labels.TFT_STATISTICS_USE_TFDV_LABEL)
per_set_stats_output_paths = common.GetValues(
outputs, labels.PER_SET_STATS_OUTPUT_PATHS_LABEL)
temp_path = common.GetSoleValue(outputs, labels.TEMP_OUTPUT_LABEL)
input_cache_dir = common.GetSoleValue(
inputs, labels.CACHE_INPUT_PATH_LABEL, strict=False)
output_cache_dir = common.GetSoleValue(
outputs, labels.CACHE_OUTPUT_PATH_LABEL, strict=False)
tf.logging.info('Analyze and transform data patterns: %s',
list(enumerate(analyze_and_transform_data_paths)))
tf.logging.info('Transform data patterns: %s',
list(enumerate(transform_only_data_paths)))
tf.logging.info('Transform materialization output paths: %s',
list(enumerate(materialize_output_paths)))
tf.logging.info('Transform output path: %s', transform_output_path)
feature_spec = schema_utils.schema_as_feature_spec(
_GetSchemaProto(input_dataset_metadata)).feature_spec
try:
analyze_input_columns = tft.get_analyze_input_columns(
preprocessing_fn, feature_spec)
transform_input_columns = (
tft.get_transform_input_columns(preprocessing_fn, feature_spec))
except AttributeError:
# If using TFT 1.12, fall back to assuming all features are used.
analyze_input_columns = feature_spec.keys()
transform_input_columns = feature_spec.keys()
# Use the same dataset (same columns) for AnalyzeDataset and computing
# pre-transform stats so that the data will only be read once for these
# two operations.
if compute_statistics:
analyze_input_columns = list(
set(list(analyze_input_columns) + list(transform_input_columns)))
if input_dataset_metadata.schema is _RAW_EXAMPLE_SCHEMA:
analyze_input_dataset_metadata = input_dataset_metadata
transform_input_dataset_metadata = input_dataset_metadata
else:
analyze_input_dataset_metadata = dataset_metadata.DatasetMetadata(
dataset_schema.from_feature_spec(
{feature: feature_spec[feature]
for feature in analyze_input_columns}))
transform_input_dataset_metadata = dataset_metadata.DatasetMetadata(
dataset_schema.from_feature_spec(
{feature: feature_spec[feature]
for feature in transform_input_columns}))
can_process_jointly = not bool(per_set_stats_output_paths or
materialize_output_paths or output_cache_dir)
analyze_data_list = self._MakeDatasetList(
analyze_and_transform_data_paths, raw_examples_file_format,
raw_examples_data_format, analyze_input_dataset_metadata,
can_process_jointly)
transform_data_list = self._MakeDatasetList(
list(analyze_and_transform_data_paths) +
list(transform_only_data_paths), raw_examples_file_format,
raw_examples_data_format, transform_input_dataset_metadata,
can_process_jointly)
desired_batch_size = self._GetDesiredBatchSize(raw_examples_data_format)
with self._CreatePipeline(outputs) as p:
with tft_beam.Context(
temp_dir=temp_path,
desired_batch_size=desired_batch_size,
passthrough_keys={_TRANSFORM_INTERNAL_FEATURE_FOR_KEY},
use_deep_copy_optimization=True):
# pylint: disable=expression-not-assigned
# pylint: disable=no-value-for-parameter
_ = (
p | self._IncrementColumnUsageCounter(
len(feature_spec.keys()), len(analyze_input_columns),
len(transform_input_columns)))
(new_analyze_data_dict, input_cache, flat_data_required) = (
p | self._OptimizeRun(input_cache_dir, output_cache_dir,
analyze_data_list, feature_spec,
preprocessing_fn, self._GetCacheSource()))
# Removing unneeded datasets if they won't be needed for
# materialization. This means that these datasets won't be included in
# the statistics computation or profiling either.
if not materialize_output_paths:
analyze_data_list = [
d for d in new_analyze_data_dict.values() if d is not None
]
analyze_decode_fn = (
self._GetDecodeFunction(raw_examples_data_format,
analyze_input_dataset_metadata.schema))
for (idx, dataset) in enumerate(analyze_data_list):
dataset.encoded = (
p | 'ReadAnalysisDataset[{}]'.format(idx) >>
self._ReadExamples(dataset))
dataset.decoded = (
dataset.encoded
| 'DecodeAnalysisDataset[{}]'.format(idx) >>
self._DecodeInputs(analyze_decode_fn))
input_analysis_data = {}
for key, dataset in six.iteritems(new_analyze_data_dict):
if dataset is None:
input_analysis_data[key] = None
else:
input_analysis_data[key] = dataset.decoded
if flat_data_required:
flat_input_analysis_data = (
[dataset.decoded for dataset in analyze_data_list]
| 'FlattenAnalysisDatasets' >> beam.Flatten(pipeline=p))
else:
flat_input_analysis_data = None
if input_cache:
tf.logging.info('Analyzing data with cache.')
transform_fn, cache_output = (
(flat_input_analysis_data, input_analysis_data, input_cache,
input_dataset_metadata)
| 'AnalyzeDataset' >> tft_beam.AnalyzeDatasetWithCache(
preprocessing_fn, pipeline=p))
# Write the raw/input metadata.
(input_dataset_metadata
| 'WriteMetadata' >> tft_beam.WriteMetadata(
os.path.join(transform_output_path,
tft.TFTransformOutput.RAW_METADATA_DIR), p))
# WriteTransformFn writes transform_fn and metadata to subdirectories
# tensorflow_transform.SAVED_MODEL_DIR and
# tensorflow_transform.TRANSFORMED_METADATA_DIR respectively.
(transform_fn |
'WriteTransformFn' >> tft_beam.WriteTransformFn(transform_output_path))
if output_cache_dir is not None and cache_output is not None:
# TODO(b/37788560): Possibly make this part of the beam graph.
tf.io.gfile.makedirs(output_cache_dir)
tf.logging.info('Using existing cache in: %s', input_cache_dir)
if input_cache_dir is not None:
# Only copy cache that is relevant to this iteration. This is
# assuming that this pipeline operates on rolling ranges, so those
# cache entries may also be relevant for future iterations.
for span_cache_dir in input_analysis_data:
full_span_cache_dir = os.path.join(input_cache_dir,
span_cache_dir)
if tf.io.gfile.isdir(full_span_cache_dir):
self._CopyCache(full_span_cache_dir,
os.path.join(output_cache_dir, span_cache_dir))
(cache_output
| 'WriteCache' >> analyzer_cache.WriteAnalysisCacheToFS(
p, output_cache_dir, sink=self._GetCacheSink()))
if compute_statistics or materialize_output_paths:
# Do not compute pre-transform stats if the input format is raw proto,
# as StatsGen would treat any input as tf.Example.
if (compute_statistics and
not self._IsDataFormatProto(raw_examples_data_format)):
# Aggregated feature stats before transformation.
pre_transform_feature_stats_path = os.path.join(
transform_output_path,
tft.TFTransformOutput.PRE_TRANSFORM_FEATURE_STATS_PATH)
schema_proto = _GetSchemaProto(analyze_input_dataset_metadata)
([
dataset.decoded if stats_use_tfdv else dataset.encoded
for dataset in analyze_data_list
]
| 'FlattenPreTransformAnalysisDatasets' >> beam.Flatten(pipeline=p)
| 'GenerateAggregatePreTransformAnalysisStats' >>
self._GenerateStats(
pre_transform_feature_stats_path,
schema_proto,
use_deep_copy_optimization=True,
use_tfdv=stats_use_tfdv))
transform_decode_fn = (
self._GetDecodeFunction(raw_examples_data_format,
transform_input_dataset_metadata.schema))
# transform_data_list is a superset of analyze_data_list, we pay the
# cost to read the same dataset (analyze_data_list) again here to
# prevent certain beam runner from doing large temp materialization.
for (idx, dataset) in enumerate(transform_data_list):
dataset.encoded = (
p
| 'ReadTransformDataset[{}]'.format(idx) >>
self._ReadExamples(dataset))
dataset.decoded = (
dataset.encoded
| 'DecodeTransformDataset[{}]'.format(idx) >>
self._DecodeInputs(transform_decode_fn))
(dataset.transformed,
metadata) = (((dataset.decoded, transform_input_dataset_metadata),
transform_fn)
| 'TransformDataset[{}]'.format(idx) >>
tft_beam.TransformDataset())
if materialize_output_paths or not stats_use_tfdv:
dataset.transformed_and_encoded = (
dataset.transformed
| 'EncodeTransformedDataset[{}]'.format(idx) >> beam.ParDo(
self._EncodeAsExamples(), metadata))
if compute_statistics:
# Aggregated feature stats after transformation.
_, metadata = transform_fn
post_transform_feature_stats_path = os.path.join(
transform_output_path,
tft.TFTransformOutput.POST_TRANSFORM_FEATURE_STATS_PATH)
# TODO(b/70392441): Retain tf.Metadata (e.g., IntDomain) in
# schema. Currently input dataset schema only contains dtypes,
# and other metadata is dropped due to roundtrip to tensors.
transformed_schema_proto = _GetSchemaProto(metadata)
([(dataset.transformed
if stats_use_tfdv else dataset.transformed_and_encoded)
for dataset in transform_data_list]
| 'FlattenPostTransformAnalysisDatasets' >> beam.Flatten()
| 'GenerateAggregatePostTransformAnalysisStats' >>
self._GenerateStats(
post_transform_feature_stats_path,
transformed_schema_proto,
use_tfdv=stats_use_tfdv))
if per_set_stats_output_paths:
assert len(transform_data_list) == len(per_set_stats_output_paths)
# TODO(b/67632871): Remove duplicate stats gen compute that is
# done both on a flattened view of the data, and on each span
# below.
bundles = zip(transform_data_list, per_set_stats_output_paths)
for (idx, (dataset, output_path)) in enumerate(bundles):
if stats_use_tfdv:
data = dataset.transformed
else:
data = dataset.transformed_and_encoded
(data
| 'GeneratePostTransformStats[{}]'.format(idx) >>
self._GenerateStats(
output_path,
transformed_schema_proto,
use_tfdv=stats_use_tfdv))
if materialize_output_paths:
assert len(transform_data_list) == len(materialize_output_paths)
bundles = zip(transform_data_list, materialize_output_paths)
for (idx, (dataset, output_path)) in enumerate(bundles):
(dataset.transformed_and_encoded
| 'Materialize[{}]'.format(idx) >> self._WriteExamples(
raw_examples_file_format, output_path))
return _Status.OK()
def build_pipeline(df: pd.DataFrame, job_label: str, runner: str, project: str,
region: str, output_dir: str, compression: str,
num_shards: int, dataflow_options: dict,
integer_label: bool) -> beam.Pipeline:
"""Runs TFRecorder Beam Pipeline.
Args:
df: Pandas DataFrame
job_label: User description for the beam job.
runner: Beam Runner: (e.g. DataflowRunner, DirectRunner).
project: GCP project ID (if DataflowRunner)
region: GCP compute region (if DataflowRunner)
output_dir: GCS or Local Path for output.
compression: gzip or None.
num_shards: Number of shards.
dataflow_options: Dataflow Runner Options (optional)
integer_label: Flags if label is already an integer.
Returns:
beam.Pipeline
Note: These inputs must be validated upstream (by client.create_tfrecord())
"""
job_name = _get_job_name(job_label)
job_dir = _get_job_dir(output_dir, job_name)
options = _get_pipeline_options(runner, job_name, job_dir, project, region,
dataflow_options)
#with beam.Pipeline(runner, options=options) as p:
p = beam.Pipeline(options=options)
with tft_beam.Context(temp_dir=os.path.join(job_dir, 'tft_tmp')):
converter = tft.coders.CsvCoder(constants.IMAGE_CSV_COLUMNS,
constants.IMAGE_CSV_METADATA.schema)
extract_images_fn = beam_image.ExtractImagesDoFn(
constants.IMAGE_URI_KEY)
flatten_rows = ToCSVRows()
# Each element in the image_csv_data PCollection will be a dict
# including the image_csv_columns and the image features created from
# extract_images_fn.
image_csv_data = (
p
| 'ReadFromDataFrame' >> beam.Create(df.values.tolist())
| 'ToCSVRows' >> beam.ParDo(flatten_rows)
| 'DecodeCSV' >> beam.Map(converter.decode)
| 'ReadImage' >> beam.ParDo(extract_images_fn))
# Split dataset into train and validation.
train_data, val_data, test_data, discard_data = (
image_csv_data | 'SplitDataset' >> beam.Partition(
_partition_fn, len(constants.SPLIT_VALUES)))
train_dataset = (train_data, constants.RAW_METADATA)
val_dataset = (val_data, constants.RAW_METADATA)
test_dataset = (test_data, constants.RAW_METADATA)
# TensorFlow Transform applied to all datasets.
preprocessing_fn = functools.partial(_preprocessing_fn,
integer_label=integer_label)
transformed_train_dataset, transform_fn = (
train_dataset
| 'AnalyzeAndTransformTrain' >>
tft_beam.AnalyzeAndTransformDataset(preprocessing_fn))
transformed_train_data, transformed_metadata = transformed_train_dataset
transformed_data_coder = tft.coders.ExampleProtoCoder(
transformed_metadata.schema)
transformed_val_data, _ = (
(val_dataset, transform_fn)
| 'TransformVal' >> tft_beam.TransformDataset())
transformed_test_data, _ = (
(test_dataset, transform_fn)
| 'TransformTest' >> tft_beam.TransformDataset())
# Sinks for TFRecords and metadata.
tfr_writer = functools.partial(_get_write_to_tfrecord,
output_dir=job_dir,
compress=compression,
num_shards=num_shards)
_ = (transformed_train_data
| 'EncodeTrainData' >> beam.Map(transformed_data_coder.encode)
| 'WriteTrainData' >> tfr_writer(prefix='train'))
_ = (transformed_val_data
| 'EncodeValData' >> beam.Map(transformed_data_coder.encode)
| 'WriteValData' >> tfr_writer(prefix='val'))
_ = (transformed_test_data
| 'EncodeTestData' >> beam.Map(transformed_data_coder.encode)
| 'WriteTestData' >> tfr_writer(prefix='test'))
_ = (discard_data
| 'DiscardDataWriter' >> beam.io.WriteToText(
os.path.join(job_dir, 'discarded-data')))
# Output transform function and metadata
_ = (transform_fn
| 'WriteTransformFn' >> tft_beam.WriteTransformFn(job_dir))
# Output metadata schema
_ = (transformed_metadata
| 'WriteMetadata' >> tft_beam.WriteMetadata(job_dir, pipeline=p))
return p