Scaling ML inference on Databricks: Liquid or partitioned? Salted or not?

introduction

Continuous variables for four different products. The machine learning pipeline is built on Databricks and has two main components.

1. Preparing features in SQL using serverless computing.
2. Use job clusters to infer ensembles of hundreds of models and control computational power.

On the first attempt, the 420 core cluster took nearly 10 hours to process just 18 partitions.

The purpose is Adjust data flow to maximize cluster utilization and ensure scalability. Inference is performed on four sets of ML models, one set for each product. However, we will focus on the following points: How to save data according to the layout How much parallel processing can be exploited? for reasoning. It does not focus on the inner workings of the reasoning itself.

If you have too few file partitions, your cluster will take a long time to scan large files, and unless you repartition at that point (which means adding network delay and data shuffling), you may end up inferring a large set of rows within each partition as well. It also increases execution time.

However, companies have limited patience to ship ML pipelines that directly impact their organizations. Therefore, testing is limited.

This article reviews the landscape of feature data, provides an overview of ML inference, and presents results and discussion of inference performance based on four dataset processing scenarios.

1. Partitioned table, no salt, no row limit within partition (Unsalted, partitioned)
2. Partitioned table, salted, 1M row limit (Salty flavor, with partitions)
3. Liquid cluster table, no salt, no row limit on partitions (salt-free/liquid)
4. Liquid cluster table, salt, with 1M row limit (salty and liquid)

data landscape

A dataset contains features that a set of ML models uses for inference. It has approximately 550 million rows and contains 4 products identified by attributes ProductLine:

• Product A: ~10.45 million (1.9%)
• Product B: ~4.4 million (0.8%)
• Product C: ~100 million (17.6%)
• Product D: ~354 million (79.7%)

Then it has another low cardinality attribute. attrB, It contains only two distinct values ​​and is used as a filter to extract a subset of the dataset for all parts of the ML system.

moreover, RunDate Records the date the feature was generated. These are for addition only. Finally, the dataset is read using the following query:

SELECT
  Id,
  ProductLine,
  AttrB,
  AttrC,
  RunDate,
  {model_features}
FROM
  catalog.schema.FeatureStore
WHERE
  ProductLine = :product AND
  AttrB = :attributeB AND
  RunDate = :RunDate

Implementing salt

The salting here is dynamically generated. Its purpose is to distribute data according to volume. This means that larger products are allocated more buckets and smaller products are allocated fewer buckets. For example, product D should receive about 80% of the bucket given its proportion in the data landscape.

We do this to ensure that inference execution time is predictable and to maximize cluster utilization.

# Calculate percentage of each (ProductLine, AttrB) based on row counts
brand_cat_counts = df_demand_price_grid_load.groupBy(
   "ProductLine", "AttrB"
).count()
total_count = df_demand_price_grid_load.count()
brand_cat_percents = brand_cat_counts.withColumn(
   "percent", F.col("count") / F.lit(total_count)
)

# Collect percentages as dicts with string keys (this will later determine
# the number of salt buckets each product receives
brand_cat_percent_dict = {
   f"{row['ProductLine']}|{row['AttrB']}": row['percent']
   for row in brand_cat_percents.collect()
}

# Collect counts as dicts with string keys (this will help
# to add an additional bucket if counts is not divisible by the number of 
# buckets for the product
brand_cat_count_dict = {
   f"{row['ProductLine']}|{row['AttrB']}": row['count']
   for row in brand_cat_percents.collect()
}

# Helper to flatten key-value pairs for create_map
def dict_to_map_expr(d):
   expr = []
   for k, v in d.items():
       expr.append(F.lit(k))
       expr.append(F.lit(v))
   return expr

percent_case = F.create_map(*dict_to_map_expr(brand_cat_percent_dict))
count_case = F.create_map(*dict_to_map_expr(brand_cat_count_dict))

# Add string key column in pyspark
df_demand_price_grid_load = df_demand_price_grid_load.withColumn(
   "product_cat_key",
   F.concat_ws("|", F.col("ProductLine"), F.col("AttrB"))
)

df_demand_price_grid_load = df_demand_price_grid_load.withColumn(
   "percent", percent_case.getItem(F.col("product_cat_key"))
).withColumn(
   "product_count", count_case.getItem(F.col("product_cat_key"))
)

# Set min/max buckets
min_buckets = 10
max_buckets = 1160

# Calculate buckets per row based on (BrandName, price_delta_cat) percentage
df_demand_price_grid_load = df_demand_price_grid_load.withColumn(
   "buckets_base",
   (F.lit(min_buckets) + (F.col("percent") * (max_buckets - min_buckets))).cast("int")
)

# Add an extra bucket if brand_count is not divisible by buckets_base
df_demand_price_grid_load = df_demand_price_grid_load.withColumn(
   "buckets",
   F.when(
       (F.col("product_count") % F.col("buckets_base")) != 0,
       F.col("buckets_base") + 1
   ).otherwise(F.col("buckets_base"))
)

# Generate salt per row based on (ProductLine, AttrB) bucket count
df_demand_price_grid_load = df_demand_price_grid_load.withColumn(
   "salt",
   (F.rand(seed=42) * F.col("buckets")).cast("int")
)

# Perform the repartition using the core attributes and the salt column
df_demand_price_grid_load = df_demand_price_grid_load.repartition(
   1200, "AttrB", "ProductLine", "salt"
).drop("product_cat_key", "percent", "brand_count", "buckets_base", "buckets", "salt")

Finally, save the dataset to a feature table and add the maximum number of rows per partition. This is to prevent Spark from generating partitions with too many rows. It may be generated even if the salt has already been calculated.

Why force 1 million rows? The main focus is on model inference time, not file size. After several tests with 1M, 1.5M, and 2M, the first test yielded the best performance in this example. Again, this project has a very limited budget and time, so you need to make the most of your resources.

df_demand_price_grid_load.write\
   .mode("overwrite")\
   .option("replaceWhere", f"RunDate = '{params['RunDate']}'")\
   .option("maxRecordsPerFile", 1_000_000) \
   .partitionBy("RunDate", "price_delta_cat", "BrandName") \
   .saveAsTable(f"{params['catalog_revauto']}.{params['schema_revenueautomation']}.demand_features_price_grid")

Why not rely on Spark’s Adaptive Query Execution (AQE)?

Recall that the main focus is on inference time, not measurements tuned for regular Spark SQL queries such as file size. Using only AQE was actually our first attempt. As you can see from the results, the execution time was very undesirable and the cluster utilization was not maximized considering the data ratio.

machine learning inference

I have a pipeline with four tasks, one for each product. All tasks follow these general steps:

• Loads functionality from the corresponding product
• Loads a subset of ML models for the corresponding product
• Perform inference on half of the sliced ​​subset. AttrB
• Perform inference on the other half of the slice AttrB
• Save the data to the results table

To avoid overwhelming this article with numbers, I will focus on one of the inference stages, but the other stage is very similar in structure and results. Additionally, in Figure 2, you can see the DAG of the inference you are evaluating.

It looks very simple, but the execution time depends on how your data is stored and the size of your cluster.

Cluster configuration

The inference stage we are analyzing has one cluster per product, tailored to the project’s infrastructure limitations and data distribution.

• Product A: 35 workers (Standard_DS14v2, 420 cores)
• Product B: 5 workers (Standard_DS14v2, 70 cores)
• Product C: 1 worker (Standard_DS14v2, 14 cores)
• Product D: 1 worker (Standard_DS14v2, 14 cores)

Additionally, AdaptiveQueryExecution is enabled by default, and Spark determines how to best store data depending on the specified context.

Results and discussion

For each scenario, you can see the number of file partitions per product and the average number of rows per partition, which tells you how many rows the ML system performs inference per Spark task. Additionally, we present Spark UI metrics to observe runtime performance and explore the distribution of data during inference. The Spark UI part only runs for the largest product, D, to avoid including too much information. Additionally, in some scenarios, product D inference becomes a runtime bottleneck. Another reason why it was the main focus of the results.

Unsalted cutlet partitioned

Figure 3 shows that the average file partition contains tens of millions of rows. This means that a single executor will take a significant amount of execution time. On average, the largest product, C, has over 45 million rows in a single partition. Product B is the smallest, with an average number of rows of about 12 million.

Figure 4. shows the number of partitions per product, with a total of 26 for all partitions. Looking at product D, we see that it is severely short of 18 partitions on the 420 cores available, and on average runs about 40 million lines of inference across all partitions.

Look at Figure 5. The cluster spent a total of 9.9 hours, but the job had to be killed because it was expensive and blocked others from testing, and it wasn’t finished yet.

The summary statistics in Figure 6 for completed tasks show that the Product D partition was highly skewed. The maximum input size was approximately 56M and the execution time was 7.8 hours.

Salt-free/liquid

In this scenario, we observe very similar results in terms of average number of rows per file partition and number of partitions per product, as shown in Figure 7 and Figure 8, respectively.

Product D has 19 file partitions, which is still well short of 420 cores.

Since we can already predict that this experiment will be very expensive, we decided to skip the inference test in this scenario. Again, in an ideal situation I would bring it forward, but I have an outstanding ticket on my board.

salty and partitioned

After applying the salting and repartitioning process, the average number of records per partition for products A and B is approximately 2.5 million, and for products C and D approximately 1 million, as shown in Figure 9.

Additionally, looking at Figure 10, the number of file partitions increases to approximately 860 in Product D, with 430 at each inference stage.

As a result, the execution time for product D inference with 360 tasks is 3 hours, as shown in Figure 11.

Looking at the summary statistics in Figure 12, the execution time is around 1.7 and the distribution seems balanced, but the task takes up to 3 hours and is worth further investigation in the future.

One of the big advantages is that Salt distributes the data according to the product ratio. You can increase the number of shuffle partitions if more resources are available. repartition() Add workers according to the percentage of data. This ensures that the process scales predictably.

salty and liquid

This scenario combines the two most powerful levers we have considered so far.

Salting, which controls file size and parallelism, and Liquid Clustering, which keeps related data in the same location without strict partition boundaries.

After applying the same salting strategy and limit of 1 million rows per partition, the liquid clustered table exhibits an average partition size very similar to the salted and partitioned case, as shown in Figure 13. Products C and D remain close to the 1 million row goal, while Products A and B are slightly above that threshold.

However, the main difference is in how these partitions are distributed and consumed by Spark. As shown in Figure 14, product D again reaches a large number of file partitions and provides enough parallelism to saturate the available cores during inference.

Unlike its partitioned counterpart, liquid clustering allows Spark to adapt its file layout over time while still benefiting from salt. This results in fewer extreme outliers in both input size and task duration, and a more even distribution of work among performers.

The summary statistics in Figure 15 show that the majority of the tasks are completed within a narrow execution window and the maximum duration of the tasks is shorter than in the salty partitioned scenario. This indicates reduced skew and improved load balancing across the cluster.

An important side effect is that liquid clustering maintains data locality for filtered columns without enforcing strict partition boundaries. This allows Spark to still benefit from data skipping, but the salt prevents a single executor from being overwhelmed by tens of millions of rows.

whole, salty and liquid emerges as the most robust setup. This maximizes parallelism, minimizes skew, and reduces operational risk as inference workloads grow or cluster configuration changes.

Important points

• Inference scalability is often limited by data layout rather than model complexity. If the file partitions are not sized appropriately, hundreds of cores can sit idle while a few executors process tens of millions of rows.
• Partitioning alone is not sufficient for large-scale inference. If you don’t control file size, partitioned tables can produce large partitions and long-running skewed tasks.
• Salting is an effective tool for breaking parallelism. Introducing a salt key and enforcing a per-partition row limit significantly increases the number of executable tasks and stabilizes the execution time.
• Liquid clustering complements salting by reducing skew without having hard boundaries. This allows Spark to adapt its file layout over time, making the system more resilient as data grows.