Решение проблем сериализации в PySpark

Если вы когда-либо пробовали использовать библиотеки, такие как numpy, scipy, scikit-learn, transformers, PyTorch или Tensorflow в PySpark, вы знаете, насколько сложно это может быть из-за ошибок сериализации, поскольку ядро ​​Java Spark пытается распределить работу по кластеру. Это означает, что вам нужно использовать набор инструментов Spark ML, что не так желательно, как типичный стек Python ML.

Решение? Используйте RDD вместо DataFrames, а затем используйте свои инструменты Python через RDD.mapPartitions() вместо RDD.map().

Например, следующий код, использующий RDD.map(), не будет работать в PySpark 3.x. Это вызывает ошибку сериализации, когда Spark пытается Javify класс TfidfVectorizer scikit-learn и терпит неудачу.

# Create and fit a TF-IDF vectorizer on a small sample of documents
tf_vectorizer = TfidfVectorizer()
tf_vectorizer = tf_vectorizer.fit(bio_corpus)
def vectorize_post(post:dict, tf_vectorizer:TfidfVectorizer) -> dict:
   """TF-IDF vectorize blog posts"""
   post["body_vector"] = tf_vectorizer.transform(post["body"])
   return post
# TF-IDF encode many documents using the vectorizer and RDD.map()
vector_posts = posts.map(lambda x: vectorize_post(x, tf_vectorizer))

Однако следующий код, который выполняет то же самое, будет запущен с использованием RDD.mapPartitions().

tf_vectorizer = TfidfVectorizer()
tf_vectorizer = tf_vectorizer.fit(bio_corpus)
def vectorize_post_partitions(partition: typing.Iterator[dict], tf_vectorizer:TfidfVectorizer) -> typing.Iterator[dict]:
   """TF-IDF vectorize all blog posts within a partition"""
   # Vectorize and emit each post in the partition
   for post in partition:
         post["body_vector"] = tf_vectorizer.transform(post["body"])
         yield post
# Ensure there are enough partitions to achieve parallel processing
posts = posts.partition(64)  # Ex. 64 cores in cluster
vector_posts = posts.mapPartitions(lambda x: vectorize_post_partitions(x, tf_vectorizer))

Результат двух заданий одинаков, и до тех пор, пока вы создаете достаточное количество разделов, эффективность сохраняется. Я не могу говорить о точной эффективности, но на практике это работает достаточно хорошо.

Пусть ваших перегородок будет много, пусть ваши работы завершатся! С новым годом!

P.S. Этот блог больше не о продвижении себя в качестве консультанта, потому что теперь я являюсь соучредителем и техническим директором Deep Discovery, где мы используем нейронные сети Graph (GNN) для борьбы с отмыванием денег! :)

Рассел Джерни
Технический директор Deep Discovery