def transform_data(train_data_file: str, test_data_file: str, working_dir: str):
  """Transform the data and write out as a TFRecord of Example protos.

  Read in the data using the CSV reader, and transform it using a
  preprocessing pipeline that scales numeric data and converts categorical data
  from strings to int64 values indices, by creating a vocabulary for each
  category.

  Args:
    train_data_file: File containing training data
    test_data_file: File containing test data
    working_dir: Directory to write transformed data and metadata to
  """

  def preprocessing_fn(inputs):
    """Preprocess input columns into transformed columns."""
    # Since we are modifying some features and leaving others unchanged, we
    # start by setting `outputs` to a copy of `inputs.
    outputs = inputs.copy()

    # Scale numeric columns to have range [0, 1].
    for key in NUMERIC_FEATURE_KEYS:
      outputs[key] = tft.scale_to_0_1(inputs[key])

    for key in OPTIONAL_NUMERIC_FEATURE_KEYS:
      # This is a RaggedTensor because it is optional. Here we fill in a default
      # value when it is missing, after scaling it.
      outputs[key] = tft.scale_to_0_1(inputs[key]).to_tensor(
          default_value=0.0, shape=[None, 1]
      )

    # For all categorical columns except the label column, we generate a
    # vocabulary, and convert the string feature to a one-hot encoding.
    for key in CATEGORICAL_FEATURE_KEYS:
      integerized = tft.compute_and_apply_vocabulary(
          tf.strings.strip(inputs[key]),
          num_oov_buckets=NUM_OOV_BUCKETS,
          vocab_filename=key,
      )
      depth = (
          tft.experimental.get_vocabulary_size_by_name(key) + NUM_OOV_BUCKETS
      )
      one_hot_encoded = tf.one_hot(
          integerized,
          depth=tf.cast(depth, tf.int32),
          on_value=1,
          off_value=0,
          dtype=tf.int64,
      )
      # Saving one-hot encoded outputs as sparse in order to avoid large dense
      # (mostly empty) tensors. This is especially important when saving
      # transformed data to disk.
      outputs[key] = tf.sparse.from_dense(
          tf.reshape(one_hot_encoded, [-1, depth])
      )
      tft.experimental.annotate_sparse_output_shape(outputs[key], depth)

    # For the label column we provide the mapping from string to index.
    table_keys = ['>50K', '<=50K']
    with tf.init_scope():
      initializer = tf.lookup.KeyValueTensorInitializer(
          keys=table_keys,
          values=tf.cast(tf.range(len(table_keys)), tf.int64),
          key_dtype=tf.string,
          value_dtype=tf.int64,
      )
      table = tf.lookup.StaticHashTable(initializer, default_value=-1)
    # Remove trailing periods for test data when the data is read with tf.data.
    label_str = tf.strings.regex_replace(inputs[LABEL_KEY], r'\.', '')
    label_str = tf.strings.strip(label_str)
    data_labels = table.lookup(label_str)
    transformed_label = tf.one_hot(
        indices=data_labels, depth=len(table_keys), on_value=1.0, off_value=0.0
    )
    outputs[LABEL_KEY] = tf.reshape(transformed_label, [-1, len(table_keys)])

    return outputs

  # The "with" block will create a pipeline, and run that pipeline at the exit
  # of the block.
  with beam.Pipeline() as pipeline:
    with tft_beam.Context(temp_dir=tempfile.mkdtemp()):
      # Create a TFXIO to read the census data with the schema. To do this we
      # need to list all columns in order since the schema doesn't specify the
      # order of columns in the csv.
      # We first read CSV files and use BeamRecordCsvTFXIO whose .BeamSource()
      # accepts a PCollection[bytes] because we need to patch the records first
      # (see "FixCommasTrainData" below). Otherwise, tfxio.CsvTFXIO can be used
      # to both read the CSV files and parse them to TFT inputs:
      # csv_tfxio = tfxio.CsvTFXIO(...)
      # raw_data = (pipeline | 'ToRecordBatches' >> csv_tfxio.BeamSource())
      csv_tfxio = tfxio.BeamRecordCsvTFXIO(
          physical_format='text',
          column_names=ORDERED_CSV_COLUMNS,
          schema=_SCHEMA,
      )

      # Read in raw data and convert using CSV TFXIO.  Note that we apply
      # some Beam transformations here, which will not be encoded in the TF
      # graph since we don't do the from within tf.Transform's methods
      # (AnalyzeDataset, TransformDataset etc.).  These transformations are just
      # to get data into a format that the CSV TFXIO can read, in particular
      # removing spaces after commas.
      raw_data = (
          pipeline
          | 'ReadTrainData'
          >> beam.io.ReadFromText(
              train_data_file, coder=beam.coders.BytesCoder()
          )
          | 'FixCommasTrainData'
          >> beam.Map(lambda line: line.replace(b', ', b','))
          | 'DecodeTrainData' >> csv_tfxio.BeamSource()
      )

      # Combine data and schema into a dataset tuple.  Note that we already used
      # the schema to read the CSV data, but we also need it to interpret
      # raw_data.
      raw_dataset = (raw_data, csv_tfxio.TensorAdapterConfig())

      # The TFXIO output format is chosen for improved performance.
      transformed_dataset, transform_fn = (
          raw_dataset
          | tft_beam.AnalyzeAndTransformDataset(
              preprocessing_fn, output_record_batches=True
          )
      )

      # Extract transformed RecordBatches, encode and write them to the given
      # directory.
      _ = (
          transformed_dataset
          | 'EncodeTrainData' >> tft_beam.EncodeTransformedDataset()
          | 'WriteTrainData'
          >> beam.io.WriteToTFRecord(
              os.path.join(working_dir, TRANSFORMED_TRAIN_DATA_FILEBASE)
          )
      )

      # Now apply transform function to test data.  In this case we remove the
      # trailing period at the end of each line, and also ignore the header line
      # that is present in the test data file.
      raw_test_data = (
          pipeline
          | 'ReadTestData'
          >> beam.io.ReadFromText(
              test_data_file,
              skip_header_lines=1,
              coder=beam.coders.BytesCoder(),
          )
          | 'FixCommasTestData'
          >> beam.Map(lambda line: line.replace(b', ', b','))
          | 'RemoveTrailingPeriodsTestData' >> beam.Map(lambda line: line[:-1])
          | 'DecodeTestData' >> csv_tfxio.BeamSource()
      )

      raw_test_dataset = (raw_test_data, csv_tfxio.TensorAdapterConfig())

      # The TFXIO output format is chosen for improved performance.
      transformed_test_dataset = (
          raw_test_dataset,
          transform_fn,
      ) | tft_beam.TransformDataset(output_record_batches=True)

      # Extract transformed RecordBatches, encode and write them to the given
      # directory.
      _ = (
          transformed_test_dataset
          | 'EncodeTestData' >> tft_beam.EncodeTransformedDataset()
          | 'WriteTestData'
          >> beam.io.WriteToTFRecord(
              os.path.join(working_dir, TRANSFORMED_TEST_DATA_FILEBASE)
          )
      )

      # Will write a SavedModel and metadata to working_dir, which can then
      # be read by the tft.TFTransformOutput class.
      _ = transform_fn | 'WriteTransformFn' >> tft_beam.WriteTransformFn(
          working_dir
      )
Transform the data and write out as a TFRecord of Example protos.

  Read in the data using the CSV reader, and transform it using a
  preprocessing pipeline that scales numeric data and converts categorical data
  from strings to int64 values indices, by creating a vocabulary for each
  category.

  Args:
    train_data_file: File containing training data
    test_data_file: File containing test data
    working_dir: Directory to write transformed data and metadata to
  
examples/census_example_v2.py