topic Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column? in Machine Learning

What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

excavator-matt — Wed, 03 Sep 2025 09:32:33 GMT

We're trying to run the bundled sentence-transformers library from SBert in a notebook running Databricks ML 16.4 on an AWS g4dn.2xlarge [T4] instance.

However, we're experiencing out of memory crashes and are wondering what the optimal to run sentence vector encoding in Databricks is.

We have tried three different approaches, but neither really works.

1. Skip spark entirely

In this naive approach, we skip spark entirely and continue to run it in standard Python using the toPandas() function on the Spark DataFrame

projects_pdf = df_projects.toPandas()
max_seq_length = 256

sentence_model_name = "paraphrase-multilingual-mpnet-base-v2"
sentence_model = SentenceTransformer(sentence_model_name)
sentence_model.max_seq_length = max_seq_length

text_to_encode = projects_pdf["project_text"].tolist()
np_text_embeddings = sentence_model.encode(text_to_encode, batch_size=128, show_progress_bar=True, convert_to_numpy=True)

This runs and renders the progress bar nicely, but the problem is now converting back into Delta table.

projects_pdf["text_embeddings"] = np_text_embeddings.tolist()

projects_pdf.to_delta("europe_prod_catalog.ad_hoc.project_recommendation_stage", mode="overwrite")

This part will crash with memory issue ("The spark driver has stopped unexpectedly and is restarting. Your notebook will be automatically reattached.")

2. Use Pandas UDF

The second approach is stolen from StackOverflow and is based on Spark's pandas_udf, but does work four our volume of data.

from sentence_transformers import SentenceTransformer

import pandas as pd

import pyspark.sql.functions as F

from pyspark.sql.types import ArrayType, DoubleType, StringType

from sentence_transformers import SentenceTransformer

sentence_model_name = "paraphrase-multilingual-mpnet-base-v2"

max_seq_length = 256

mpnet_sentence_model = SentenceTransformer(sentence_model_name)

mpnet_sentence_model.max_seq_length = max_seq_length

@F.pandas_udf(returnType=ArrayType(DoubleType()))

def mpnet_encode(x: pd.Series) -> pd.Series:

return pd.Series(mpnet_sentence_model.encode(x, batch_size=128).tolist())

# apply udf and show

project_df_2 = projects_df.withColumn("project_text_embedding", mpnet_encode("project_text"))

project_df_2.write.mode("overwrite").saveAsTable("my_table)

This delays execution, but once you try to save it with saveAsTable, we get the same memory error ("The spark driver has stopped unexpectedly and is restarting. Your notebook will be automatically reattached."). I also couldn't get the progress bar to work here.

3. Use MLFlow Spark UDF

I am not entirely sure what MLFlow does and if it is any different from the previous approach, but I also tried using spark_udf.

from sentence_transformers import SentenceTransformer
import mlflow

sentence_model = SentenceTransformer("paraphrase-multilingual-mpnet-base-v2")
sentence_model.max_seq_length = 256
data = "MLflow is awesome!"
signature = mlflow.models.infer_signature(
model_input=data,
model_output=sentence_model.encode(data),
)

with mlflow.start_run() as run:
mlflow.sentence_transformers.log_model(
artifact_path="paraphrase-multilingual-mpnet-base-v2-256",
model=sentence_model,
signature=signature,
input_example=data,
)
model_uri = f"runs:/{run.info.run_id}/paraphrase-multilingual-mpnet-base-v2-256"
print(model_uri)

udf = mlflow.pyfunc.spark_udf(
spark,
model_uri=model_uri,
)

# Apply the Spark UDF to the DataFrame. This performs batch predictions across all rows in a distributed manner.
df_project_embedding = df_projects.withColumn("prediction", udf(df_projects["project_text"]))

This ticks, but you don't see if it makes any progress.

Conclusion

The current workaround is to go with the first approach and skip the spark part by storing it as a file in a Databricks Volume instead. However, this fundamentally tabular data (although it involves vector as a column) and having it in a volume loses all the benefits of Databricks.

Another aspect we considered was to create our own batching solution, but the point of Spark is that it should abstract big data handling, so it also seems wrong.

What is the ideal approach here?

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

Louis_Frolio — Mon, 08 Sep 2025 17:32:20 GMT

Spark is designed to handle very large datasets by distributing processing across a cluster, which is why working with Spark DataFrames unlocks these scalability benefits. In contrast, Python and Pandas are not inherently distributed; Pandas dataframes are eagerly evaluated and executed locally, so you can encounter memory issues when working with large datasets. For instance, exceeding around 95 GB of data in Pandas often leads to out-of-memory errors because only the driver node handles all computation, regardless of cluster size.

To bridge this gap, consider using the Pandas API on Spark, which is part of the Spark ecosystem. This API provides Pandas-equivalent syntax and functionality, while leveraging Spark’s distributed processing to handle larger data volumes efficiently. You can learn more here: https://docs.databricks.com/aws/en/pandas/pandas-on-spark.

In short, the Pandas API on Spark lets you write familiar Pandas-style code but benefit from distributed computation. It greatly reduces memory bottlenecks and scales to bigger datasets than native Pandas workflows allow.

Hope this helps, Louis.

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

excavator-matt — Mon, 06 Oct 2025 14:30:50 GMT

Sorry for the lack of reply. I had to switch tasks, but I hope to be able to test this. You suggest that Pandas On Spark is more efficient implemented than Pandas UDF?

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

Louis_Frolio — Tue, 07 Oct 2025 15:26:54 GMT

Yes because you will be using Spark dataframes which are distributed. Pandas dataframes are not distributed.

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

jamesl — Tue, 07 Oct 2025 21:34:12 GMT

If you didn't get this to work with Pandas API on Spark, you might also try importing and instantiating the SentenceTransformer model inside the pandas UDF for proper distributed execution.

Each executor runs code independently, and when Spark executes a pandas UDF the function is serialized and sent to worker nodes. If you instantiate the model globally (outside the UDF), only the driver knows about it, and Spark would then try to serialize the entire model object and send it to the workers. This could fail or lead to memory issues with complex objects like ML models.

By creating the model inside the UDF function, you ensure that each executor loads the model locally and has everything it needs to process the batch or partition of data it receives. Perhaps something like this...

@F.pandas_udf(returnType=ArrayType(DoubleType())) def mpnet_encode(x: pd.Series) -> pd.Series: # Import and instantiate model inside the UDF from sentence_transformers import SentenceTransformer model = SentenceTransformer("paraphrase-multilingual-mpnet-base-v2") model.max_seq_length = 256 return pd.Series(model.encode(x.tolist(), batch_size=128).tolist())

I hope that helps. Let us know if it works out for you.

-James

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

excavator-matt — Thu, 30 Oct 2025 15:59:03 GMT

@Louis_Frolio

I tried the Pandas on Spark approach.

How do I from Delta table into Pandas on Spark DataFrame. Is this the best way?

projects_df = spark.read.table("my_catalog.my_schema.my_project_table")

projects_spdf = ps.from_pandas(projects_df.toPandas())

It runs past the sentence transformers bit, but when I try to

projects_spdf["text_embeddings"] = np_text_embeddings.tolist()

I get this strange error.

/databricks/spark/python/pyspark/pandas/frame.py in ?(psdf, this_column_labels, that_column_labels) 13467 def assign_columns( 13468 psdf: DataFrame, this_column_labels: List[Label], that_column_labels: List[Label] 13469 ) -> Iterator[Tuple["Series", Label]]: > 13470 assert len(key) == len(that_column_labels) 13471 # Note that here intentionally uses `zip_longest` that combine 13472 # that_columns. 13473 for k, this_label, that_label in zip_longest(

At least it isn't a memory issue, but the attempt to run standard pandas at least went past this

Perhaps my issues is not so much sentence transformers but how to go from a massive list of arrays back into Delta table. That why I am a bit hesitant about @jamesl , but I could give it a try. Maybe there is lazy load issue somewhere.

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

excavator-matt — Thu, 30 Oct 2025 18:13:58 GMT

Also, upgrading to 17.3 ML still gives the same error.

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

Louis_Frolio — Fri, 31 Oct 2025 15:14:29 GMT

@excavator-matt I’d recommend a quick refresher on the Pandas API on Spark to understand the implementation details. This video breaks it down clearly: https://youtu.be/tdZDotqKtps?si=pcIzCUYs2s_TeQKx

Hope this helps. — Louis

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

excavator-matt — Wed, 25 Feb 2026 12:51:23 GMT

If anyone else is still interested in this topic, I tried running it with Databricks runtime 18.1 ML. We still have the same memory issue, but we now also get the nice error message for in approach 3 as

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 24.00 MiB. GPU 0 has a total capacity of 14.58 GiB of which 5.62 MiB is free

One thing I didn't discuss in this thread is to try to use the hosted models somehow. If your end goal is a vector index, you can now also bypass this step and simply select embedding model instead of precomputing it. This raises the question on how to reason about host. The hosted models are limited compared to the world wide web and I think those are pay per token instead of runtime.

Re: What is the most efficient way of running sentence-transformers on a Spark DataFrame column?

excavator-matt — Wed, 25 Feb 2026 13:35:14 GMT

Also, I forgot to mention the workaround solution for the first approach. If you write to parquet in a volume, you can then convert it back to a Delta table in a later cell.

Instead of this

projects_pdf.to_delta("europe_prod_catalog.ad_hoc.project_recommendation_stage", mode="overwrite")

You do this

# Avoid datetime64 timestamps error

def convert_datetime_columns_to_str(df😞

for col in df.columns:

if pd.api.types.is_datetime64_any_dtype(df[col]):

df[col] = df[col].astype(str)

return df

projects_pdf_fixed = convert_datetime_columns_to_str(projects_pdf)

projects_pdf_fixed.to_parquet("/Volumes/europe_prod_catalog/ad_hoc/temp/project_recommendation_embedding.parquet")

Then in next cell you do

project_embedded_df = spark.read.parquet("/Volumes/europe_prod_catalog/ad_hoc/temp/project_recommendation_embedding.parquet")

project_embedded_df.write.mode("overwrite").saveAsTable("europe_prod_catalog.ad_hoc.project_recommendation_embedding")