[KongchangAI]
· 2 min read· 1,145 words

Running Decision Models Directly on Governed Data in Databricks: A Practical Guide to ai_query

Running Decision Models Directly on Governed Data in Databricks: A Practical Guide to ai_query

Databricks enables SQL-native inference on governed data by deploying open-weight decision models on serverless GPUs via ai_query.

This solution bridges the gap between data governance and AI inference: open-weight decision models are deployed to serverless GPUs via Databricks AI Runtime and Model Serving, then invoked directly in SQL or Lakeflow pipelines using the `ai_query` function, returning classification labels and probability scores. The entire workflow is packaged as an importable notebook that runs in three steps. The standout value is that inference happens in place — data stays within governed boundaries, preserving access controls, lineage tracking, and audit logs, making it well-suited for compliance-sensitive industries. The author frames this as a practical demo for data teams rather than a major product launch, and notes that model performance, GPU costs, and large-scale concurrency still require independent validation.

A Direct Path from Governed Data to Decision Models

When deploying AI decision models, enterprises have long faced an uncomfortable reality: models and data live in separate worlds. Data is strictly governed within a data lakehouse, while model inference requires a lengthy round-trip — exporting data, transmitting it, running inference, then writing results back. This introduces compliance risks and slows down iteration cycles.

A recently shared Databricks workflow proposes a more direct approach: deploy open-weight decision models on serverless GPUs, then invoke them directly from SQL queries or data pipeline tasks, so that inference happens where the data already lives.

Workflow diagram showing decision models running in Databricks

Core Mechanism: AI Runtime + Model Serving + ai_query

The technical backbone of this solution has three components that work together to close the loop from data to decision.

Model Serving on Serverless GPUs

The first step is deploying an open-weight model — something like SemIf-OpenJev — onto Serverless GPU infrastructure using Databricks AI Runtime and Model Serving. The value of serverless GPUs is that they eliminate the burden of manually provisioning and managing compute clusters. Resources spin up on demand and are released when done, so teams don't pay for idle GPU capacity or worry about low-level scheduling.

"Jev-style" decision models refer to models specifically designed for structured classification and probability scoring. Rather than generating free-form text, they output structured judgments that downstream systems can directly consume — such as which category a record belongs to, or the confidence score for matching a given criterion.

Open-weight models are models whose weight parameters are publicly released, allowing users to download and deploy them independently. The key difference from closed-source API services (like GPT-4 or Claude) is that data never needs to leave your own environment to be sent to a third-party provider. For compliance-sensitive use cases, this is critical — enterprises can run inference within their own controlled compute environment, avoiding the transmission of sensitive data to external API endpoints over the network. Open-weight models also offer predictable inference costs: rather than per-token billing, you pay for actual compute resources consumed, which is often more economical for batch scoring workloads. Common open-weight model families include Meta's Llama and Mistral AI's Mistral/Mixtral series. Variants purpose-built for structured classification tasks are typically fine-tuned on top of these base architectures to improve accuracy and inference speed on specific decision tasks.

Calling Models Directly from SQL with ai_query

What makes this approach truly elegant is the ai_query function. Once a model serving endpoint is deployed, analysts and engineers can call it directly within SQL queries — feeding an entire column of data to the model and getting back structured classification results and probability scores.

This means model inference is seamlessly embedded into familiar SQL workflows. For teams already orchestrating data pipelines with Lakeflow, decision models can exist as an ordinary step in the pipeline, sitting alongside other data transformation logic, without requiring a separate inference serving infrastructure.

ai_query is a built-in Databricks SQL function that allows users to invoke deployed Model Serving endpoints directly within standard SQL statements. Its calling convention resembles a regular SQL function call: SELECT ai_query('endpoint_name', input_column) FROM my_table. The function handles the underlying engineering details — HTTP request construction, response parsing, error retries, and more. The core value of this design is that it lowers the cognitive barrier to using models — analysts don't need to learn a Python SDK, REST API, or any orchestration framework. SQL proficiency alone is sufficient to embed models into existing data processing logic. Lakeflow (Databricks' data pipeline orchestration tool) also natively supports SQL tasks containing ai_query as nodes within a DAG, allowing decision model inference to be scheduled, monitored, and dependency-tracked just like any ordinary data transformation.

One Importable Notebook, Three Steps to Run

Another selling point emphasized in the original post is how low the barrier to entry is. The entire workflow is packaged into a single importable Notebook, reduced to three steps:

  • Pick your model — choose the model to deploy
  • Select Serverless GPU — specify the serverless GPU
  • Click Run All — run all cells with one click

Once the notebook runs successfully, users can use it as a starting point — swapping in their own data and decision criteria. In other words, the notebook functions more like a template than a one-off demo. Adapting it to a real business use case simply means replacing the sample data with your own tables and adjusting the classification logic to reflect your business rules.

Why "Data Stays Within Governed Boundaries" Matters

The most noteworthy value of this solution is that inference happens in place, on governed data.

In compliance-constrained industries — finance, healthcare, government — every data movement carries risk and audit cost. In traditional approaches, scoring sensitive data with an external model typically requires moving that data out of the governed environment, which may violate data residency requirements and expand the potential breach surface.

By hosting the model service inside the data lakehouse and invoking it directly through SQL, data never leaves the governed boundary. Permissions, lineage tracking, and auditing — all core data governance capabilities — remain fully intact, while AI capabilities are layered on top as simple "function calls." This architectural convergence may have more practical impact on enterprise AI adoption than improvements to model performance alone.

Data Residency is a core requirement in many national and industry regulatory frameworks, mandating that certain types of data must be stored and processed within designated geographic regions or system boundaries. The EU GDPR, China's Data Security Law, and financial sector data localization requirements all impose varying restrictions on cross-border or cross-system data movement. In the traditional "external API scoring" model, every inference call means data leaves the enterprise's controlled environment — potentially triggering data residency compliance issues and creating data flow records in audit logs that are difficult to trace. Databricks' Unity Catalog provides column-level access control, data lineage tracking, and operational audit logs. When inference occurs within the same platform via ai_query, these governance capabilities fully cover model invocation behavior, making it auditable and traceable to answer questions like "who ran what inference, on what data, and at what time."

A Brief Assessment

Based on publicly available information, this looks more like a hands-on demonstration for data teams than a major product announcement. Its significance lies in lowering the barrier to embedding decision models into governed data workflows: serverless GPUs handle elastic compute, ai_query handles the invocation experience, and a single notebook handles the onboarding cost.

One caveat: the original material does not cover specific model performance, GPU costs, or behavior under large-scale concurrent workloads. Interested teams will still need to validate these on their own data. But architecturally, this "bring the model to the data" paradigm represents a pragmatic direction for enterprise AI engineering.

#Databricks ai_query#open-weight decision models#serverless GPU inference#governed data AI#Lakeflow pipeline#Model Serving SQL#data residency compliance#Unity Catalog AI
Share:

Related articles