Unfortunately, you're right, there's no straight forward way to pack vocabulary computes as you would with other TFT analyzers.

Depending on the size of your vocabularies you could join some of them (though with this method your vocabulary range for each feature will not be continuous, and you will need to be careful with frequency_threshold/top_k to make sure specific feature vocabularies don't get completely filtered out).

raw_data = [
      {'A': 'hello', 'B': 'world'},
      {'A': 'world', 'B': 'hello'},
      {'A': '!', 'B': '!'},
  ]

raw_data_metadata = dataset_metadata.DatasetMetadata(
    dataset_schema.from_feature_spec({
        'A': tf.io.FixedLenFeature([], tf.string),
        'B': tf.io.FixedLenFeature([], tf.string),
    }))

def preprocessing_fn(inputs):
    a = tf.strings.join(['A', inputs['A']])
    b = tf.strings.join(['B', inputs['B']])
    a_b = tf.concat((a, b), axis=-1)
    vocab = tft.vocabulary(a_b, vocab_filename='a_b_vocab')
    return {
        'a_int': tft.apply_vocabulary(a, vocab),
        'b_int': tft.apply_vocabulary(b, vocab),
    }

...
... = tft_beam.AnalyzeAndTransformDataset(preprocessing_fn)
...

tft_output = tft.TFTransformOutput(transform_fn_dir)
print(tft_output.vocabulary_by_name('a_b_vocab'))

The vocabulary contents are:
[b'Bworld', b'Bhello', b'B!', b'Aworld', b'Ahello', b'A!']

And the transformed data is:
{'a_int': 4, 'b_int': 0}
{'a_int': 3, 'b_int': 1}
{'a_int': 5, 'b_int': 2}

@zoyahav, thanks for the tips! I think that could help in certain cases for some vocabulary features.

I do have a dataset currently with 102 string features and definitely have run into the issue with not even being able to start the dataflow job because the graph size is too large. In the short term, I will try to do my best to reduce the number of tft.vocabulary analyzers (e.g. pack them per your suggestion), but in the long term what are the plans for making the analyzers more scalable on dataflow?

@cyc, Following up on this issue, Beam 2.24 has been released and should hopefully help with this issue. Could you try it and let us know? Also, for Dataflow, the V2 runner may very likely help. To try it, add --experiments=use_runner_v2

I'm following up since the thread went quiet to make sure that this was resolved for you.

1. Were you able to try Beam >= 2.24, and was it an improvement?
2. Were you able to try Dataflow Runner V2, together with Dataflow Shuffle, and was it an improvement?

Closing this due to inactivity.Please feel free to reopen.Thanks.

I have encountered the same issue, but with Flink. I am also using tft.compute_and_apply_vocabulary on dozens of features, and the Flink DAG produced by TFT is enormous. Unless --execution_mode_for_batch=BATCH_FORCED is used the pipeline just hangs, and even with that option it takes an excessively long time to run. Increasing the number of task managers and the parallelism actually causes it to take longer to finish.