I’m trying to generate 35k+ embeddings in Databricks. What I’ve tried so far:

• Per-row UDF (very slow).
• Replaced UDF with rdd.mapPartitions to batch API calls, create one Azure client per partition, and call client.embed_documents(texts) in batches. This avoids per-row Python UDF overhead and improves embedding throughput.
• Measured embedding execution vs Delta write time; embedding materialization is fine but the Delta write (saveAsTable / commit) is now the dominant, slow step. I used persist() to avoid double computation when calling count() before write.

Minimal embedding function I tested (simplified):

from pyspark.sql import Row, DataFrame
from pyspark.sql.types import StructField, StructType, ArrayType, FloatType
import os, time
from langchain_openai import AzureOpenAIEmbeddings

def embed_with_map_partitions_simple(df: DataFrame, column_names: str | list[str], batch_size: int = 128, repartition: int | None = None) -> DataFrame:
    if isinstance(column_names, str):
        column_names = [column_names]
    if repartition:
        df = df.repartition(repartition)
    spark = df.sparkSession
    new_fields = list(df.schema.fields) + [StructField(f"{c}_embedding", ArrayType(FloatType()), True) for c in column_names]
    new_schema = StructType(new_fields)

    def partition_embed(rows_iter):
        rows = list(rows_iter)
        if not rows:
            return iter(())
        client = AzureOpenAIEmbeddings(
            azure_endpoint=os.getenv("openai_api_base"),
            azure_deployment=os.getenv("openai_deployment_name"),
            api_key=os.getenv("openai_api_key"),
            api_version=os.getenv("openai_api_version")
        )
        n = len(rows)
        embeddings_per_column = {c: [None]*n for c in column_names}
        for col in column_names:
            for i in range(0, n, batch_size):
                batch = rows[i:i+batch_size]
                texts = [getattr(r, col, "") or "" for r in batch]
                try:
                    batch_emb = client.embed_documents(texts)
                except Exception:
                    batch_emb = [[0.0]*3072 for _ in texts]
                for j, emb in enumerate(batch_emb):
                    embeddings_per_column[col][i+j] = emb
        for idx, r in enumerate(rows):
            d = r.asDict()
            for c in column_names:
                d[f"{c}_embedding"] = embeddings_per_column[c][idx] or [0.0]*3072
            yield Row(**d)

    return spark.createDataFrame(df.rdd.mapPartitions(partition_embed), schema=new_schema)

Question: can Databricks advise best practices to reduce Delta write/commit time for this workflow (recommended write options, file sizing/num files, transaction tuning, or cluster/io settings)? Also any guidance on safely persisting large transformed DF before writing and on Stitch/OPTIMIZE usage would be helpful.

Thanks.