How to scale TFRS?

The documentation does not appear to offer much guidance on how to scale up a TFRS based solution, for example, here: https://www.tensorflow.org/recommenders/examples/basic_retrieval.

To start, let's consider a (common, I'm sure) scenario of N million users and M million items (such as movies). We would be envisioning at least millions (billions) of feature records also, but for the moment we could set that aside and contemplate just the 'basic retrieval' use-case/example.

There are at least 3 considerations while designing an actual recommender, in terms of scale:

1. How to feed input data into TFRS. I'm starting with Parquet from our datalake. Due to the size of this data and various filtering and augmentation of data sets, I'm looking at having to read Parquet in a Spark job, transforming the data, then persisting it into an intermediary Parquet format ready to be consumed by TFRS. Issue: extra I/O. I don't see a way to do all the transformations in TF/TFRS, especially joins.
2. The building of vocabularies. From the example:

unique_movie_titles = np.unique(np.concatenate(list(movie_titles)))
unique_user_ids = np.unique(np.concatenate(list(user_ids)))

I can uniquefy my data in Spark when prepping Parquet for TFRS. However, it appears that building a model in TFRS is oriented toward using numpy arrays:

user_model = tf.keras.Sequential([
  tf.keras.layers.experimental.preprocessing.StringLookup(
      vocabulary=unique_user_ids, mask_token=None),
  # We add an additional embedding to account for unknown tokens.
  tf.keras.layers.Embedding(len(unique_user_ids) + 1, embedding_dimension)
])

How are vocabularies expected to perform with millions of ID's? how will this affect memory and if this is not expected to scale, what alternatives are there to vocabulary/numpy array approach?
3. Getting the recommendations out, at scale. In my use-case, I'm going to ask the model to generate more recommendations than actually needed, per user, so that I can perform some post-filtering based on a set of a few rules. I need a scalable way to convert tensors en masse to Spark datasets which will write Parquet or CSV back into the datalake. This may mean extracting recs as Parquet into intermediary files then having separate code to turn that into the final, post-processed results. Issue: extra I/O and processing.

All in all, what this feels like is potentially a need for a tighter integration with Spark datasets where it is easy and seamless to load TF datasets from Spark and then just as easy to convert the resulting recommendations (e.g. a tfrs.layers.factorized_top_k.BruteForce index) to Spark datasets.

So far, I don't seem to see an easy way to do this (?). Some considerations:

• The LinkedIn's Spark-TFRecord library and the Spark Tensorflow Connector both require intermediary tfrecord files.
• I'm "fishing" for a more Spark oriented way to build a TFRS model and then extract results en masse for millions of input users. Perhaps Uber's Petastorm is an answer for the input data load side. But what about TF datasets to Spark conversion?