Get and parse raw responses of AI Functions

Prerequisites

• AI Center is activated. For more information, see AI Center.

• The engine version is esr-4.10.0 or later.

• The Spark configuration spark.emr.serverless.ai.function.enable true is added when you submit the job or create the session.

Raw functions

A raw function takes exactly the same arguments as its standard counterpart. The difference is the return value: the standard function returns a parsed result, whereas the raw function returns the original JSON response from the model service. For argument details, see AI function.

ai_query_raw()

Syntax

ai_query_raw(prompt [, service_name] [, options] [, data] [, data_type])

The return type is STRING. The positions of the generated text, usage, and request identifiers depend on the response format of the model service.

SQL example

SELECT ai_query_raw(
  'Summarize the following text in one sentence: AI Function integrates model capabilities into Spark.',
  'my_qwen_service',
  '{"temperature":0.2}'
) AS raw_response;

PySpark example

from emr_serverless_spark_ai import ai_query_raw

result_df = source_df.select(
    ai_query_raw(
        "prompt",
        service_name="my_qwen_service",
        options='{"temperature":0.2}',
    ).alias("raw_response")
)

ai_embedding_raw()

Syntax

ai_embedding_raw(
    input,
    dimension=None,
    service_name=None,
    options=None,
)

The input parameter can be a STRING or an ARRAY<STRING>, and the return type is always STRING. The raw text embedding result uses a unified envelope whose root contains the local_results and calls fields:

{
  "local_results": [ ],
  "calls": [ ]
}

• local_results: Input elements that are processed locally without calling the model service, such as NULL and empty strings. In each element, input_index is the position in the input of the current Spark row (0 for a scalar input), embedding is the locally generated result (currently an empty array), and reason is the reason, for example null_input or empty_input.

• calls: The physical model calls involved in the current Spark row. In each call, input_indices maps the results back to the positions in the original input of the current row, call_id is the physical call identifier generated by the AI Function engine, and response is the raw JSON response from the model service, which keeps the usage, dense, sparse, and unknown extension fields.

SELECT ai_embedding_raw(
  array('First text', NULL, '', 'Fourth text'),
  service_name => 'qwen3.7-text-embedding'
);

SELECT ai_embedding_raw(
  text,
  dimension => 1024,
  service_name => 'qwen3.7-text-embedding',
  options => '{
    "api_mode":"dashscope",
    "text_type":"document",
    "output_type":"dense&sparse"
  }'
)
FROM source_table;

ai_embedding_multimodal_raw()

Syntax

ai_embedding_multimodal_raw(
  data [, service_name] [, options] [, data_type]
)

The return type is STRING. The image input method is the same as ai_embedding_multimodal(). The response keeps the request identifier, usage, vectors, and extension fields returned by the model service.

Example

SELECT ai_embedding_multimodal_raw(
  image_col,
  service_name => 'multimodal-embedding-model',
  data_type => 'binary'
) AS raw_response
FROM image_table;

Parse a raw response

A raw function returns JSON text as a STRING. Use the JSON functions of Spark SQL to read the fields you need:

• To read a single field, use get_json_object() with a JSON path.

• To work with the whole response, use from_json() with a STRUCT schema and then access individual fields.

WITH result AS (
  SELECT ai_embedding_raw(
    text,
    service_name => 'qwen3.7-text-embedding',
    options => '{"api_mode":"dashscope"}'
  ) AS raw_response
  FROM source_table
)
SELECT
  get_json_object(raw_response, '$.calls[0].call_id') AS call_id,
  get_json_object(raw_response, '$.calls[0].response.request_id') AS request_id
FROM result;

import json
from emr_serverless_spark_ai import ai_embedding_raw

raw_df = source_df.select(
    ai_embedding_raw(
        "text",
        service_name="qwen3.7-text-embedding",
        options=json.dumps({"api_mode": "dashscope"}),
    ).alias("raw_response")
)

Best practices

• Prefer the standard functions for data processing. Use a raw function only when you need the original service fields.

• When you use a raw response, keep the complete JSON and read only the fields you need. Allow for unknown extension fields.

• To get a stable JSON structure, set api_mode explicitly to avoid the auto mode switching APIs when parameters change.

• The options of the embedding functions accept both dimension and dimensions to specify the vector dimension. Use only one of them and do not set both at the same time.

References

• AI function

• Model service