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.

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.