Skip to main content

ai_query

Invokes a Databricks Model Serving endpoint and returns the response.

For the corresponding Databricks SQL function, see ai_query function.

Syntax

Python
from pyspark.sql import functions as dbf

dbf.ai_query(endpoint=<endpoint>, request=<request>, return_type=<return_type>)

Parameters

Parameter

Type

Description

endpoint

str

The name of the Databricks Model Serving endpoint to invoke.

request

pyspark.sql.Column or str

A column containing the input data to send to the endpoint.

return_type

str or pyspark.sql.types.DataType, optional

The expected return type, either as a DDL string (for example, "STRING", "ARRAY<FLOAT>") or a PySpark pyspark.sql.types.DataType instance (for example, StringType(), ArrayType(FloatType())). When omitted the type is inferred from the endpoint.

Returns

pyspark.sql.Column: A new column containing the endpoint response.