Ai query parallel calls

joaoaugustofb
New Contributor

‎02-04-2026 05:12 AM

I’m trying to optimize ai_query calls on a table and wanted to get some ideas.

So far, I’ve tried repartitioning the DataFrame before running spark.sql(ai_query), but I didn’t see any meaningful performance gains. I also experimented with running multiple instances of the same notebook in parallel, but the improvements were marginal.

Has anyone tried a different approach that worked better? Any suggestions on how to improve performance or scale this more efficiently?

Labels:
0 Kudos

pavannaidu
Databricks Employee
Databricks Employee

‎02-06-2026 02:17 PM

When you are using ai_query(), there are two main aspects to performance: 

  1. Model serving endpoint
  2. SQL warehouse / Compute cluster  

Very likely, the performance is throttled by the model-serving endpoint's concurrency limit. Reference: https://docs.databricks.com/aws/en/machine-learning/model-serving/model-serving-limits

Can you share more about your model serving endpoint?