Query a served model with ai_query()

Preview

This feature is in Public Preview.

This article describes how to query a model serving endpoint from SQL with ai_query().

What is ai_query()?

The ai_query() function is a built-in Databricks SQL function, part of AI functions, that allows the model hosted by a model serving endpoint to be accessible from SQL queries. For syntax and design patterns, see ai_query function.

This function is only available in workspaces and regions where Model Serving is available and enabled. It only supports endpoints serving only models governed by Unity Catalog.

Requirements

  • The SQL commands in this guide must be run in the Databricks SQL query editor. They can not be run directly in a Databricks notebook.

  • This function is only available on pro and serverless SQL warehouses.

  • These features are in public preview. To enroll in the public preview, please populate and submit the AI Functions Public Preview enrollment form.

  • An existing model serving endpoint. See Create and manage model serving endpoints.

  • You must enable AWS PrivateLink to use this feature on pro SQL warehouses.

Query the endpoint with ai_query()

You can query the model behind the endpoint using ai_query() on serverless or pro SQL warehouses. For scoring request and response formats see Send scoring requests to serving endpoints.

Note

This function is not supported in notebook environments, including Databricks notebooks.

The following example queries the model behind the sentiment-analysis endpoint with the text dataset and specifies the return type of the request.

select text, ai_query("sentiment-analysis",  text, "returnType", 'STRUCT<label:STRING, score:DOUBLE>') as predict from test.default.c4_200k limit 10