ai_query issue

[smpa01]

As per  the doc , one can execute an ai_query in python. But I am sot sure why I can't

Works in SQL

select 1 as colA, ai_query('databricks-llama-4-maverick', 'what is 2+2') as test

But not in spark

df = spark.sql("select 1 as colA")
df = df.selectExpr("""
ai_query('databricks-llama-4-maverick', 'what is 2+2') as result
                   """)

df.show()

I am getting

AnalysisException: [UNRESOLVED_ROUTINE] Cannot resolve function `ai_query` on search path [`system`.`builtin`, `system`.`session`, `catalog`.`default`].; line 1 pos 0

I have also converted the ai_query to a sql udf but even then it does not work.

[Vinay_M_R · Accepted solution]

Hello @smpa01,

I tried to reproduce this issue internally and it worked for me as expected in spark as well. I am sharing below screenshot for your reference:

I used DBR 15.4 LTS and serverless it worked on both