E-MapReduce:Call AI functions from PySpark

Install the Python library

The official runtime environments of EMR Serverless Spark (esr-4.9.1 and later) include the emr_serverless_spark_ai library, so no additional installation is required. After you import it with import emr_serverless_spark_ai.functions as ai, you can call functions in the ai.ai_query(...) form.

If your job uses a custom runtime environment, the environment replaces the default Python environment entirely and does not include emr_serverless_spark_ai. Add the library manually as follows:

1. Download emr_serverless_spark_ai-1.0.0+2d43d393-py3-none-any.whl, and then upload it to File Management in your workspace or to OSS.

   Note

   If the package is not compatible with your engine version, contact EMR Serverless Spark technical support for a matching emr_serverless_spark_ai package.

2. On the Runtime Environment Management page, edit the target runtime environment, click Add Library, set the source type to Workspace Resource or OSS Resource, and then select the .whl file that you uploaded in the previous step.

3. Save the configuration and wait until the runtime environment is reinitialized. After its status becomes Ready, jobs that use the environment can import emr_serverless_spark_ai.

For more information about creating and managing runtime environments, see Manage runtime environments.

Initialize a SparkSession and prepare sample data

All examples in this topic are based on the following SparkSession and sample data.

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, lit
import emr_serverless_spark_ai.functions as ai

spark = SparkSession.builder.appName("ai-function-demo").getOrCreate()

reviews = spark.createDataFrame(
    [
        ("R001", "The food was great and service was excellent!"),
        ("R002", "Terrible experience, very disappointing."),
        ("R003", "It was okay, nothing special."),
    ],
    ["review_id", "review"],
)

General-purpose queries: ai_query

ai_query sends a custom prompt to the model. If you pass a string as the first argument, the string is interpreted as a column name. To pass a fixed prompt, use lit() to wrap it as a literal, for example ai.ai_query(lit("Introduce yourself")).

# Use the default model
reviews.select("review", ai.ai_query("review").alias("response")).show(truncate=False)

# Specify a model service
reviews.select(
    "review",
    ai.ai_query("review", service_name="qwen3.7-plus").alias("response"),
).show(truncate=False)

# Pass model parameters such as system_prompt and temperature through options
reviews.select(
    "review",
    ai.ai_query(
        "review",
        service_name="qwen3.7-plus",
        options='{"system_prompt": "Summarize the review in one sentence", "temperature": 0}',
    ).alias("summary"),
).show(truncate=False)

To prevent a failed call from failing the entire job, use try_ai_query instead. Its behavior is identical to ai_query, except that it returns NULL for a row that fails.

reviews.select(
    "review",
    ai.try_ai_query("review", service_name="qwen3.7-plus").alias("response"),
).show(truncate=False)

Task-specific functions

Task-specific functions use prompts that are optimized for common NLP scenarios, so their input and output formats are more stable.

Sentiment analysis: ai_sentiment returns one of three values: positive, negative, and neutral.

reviews.select("review", ai.ai_sentiment("review").alias("sentiment")).show(truncate=False)

Sample result:

+----------------------------------------------+----------+
|review                                        |sentiment |
+----------------------------------------------+----------+
|The food was great and service was excellent! |positive  |
|Terrible experience, very disappointing.      |negative  |
|It was okay, nothing special.                 |neutral   |
+----------------------------------------------+----------+

Text classification: The second argument of ai_classify is a list of labels. You can pass a Python list directly.

reviews.select(
    "review",
    ai.ai_classify("review", ["complaint", "praise", "suggestion"]).alias("category"),
).show(truncate=False)

Entity extraction: The second argument of ai_extract is a list of fields to extract. The function returns a STRUCT column.

reviews.select(
    "review",
    ai.ai_extract("review", ["sentiment_word", "aspect"]).alias("entities"),
).show(truncate=False)

Translation: The second argument of ai_translate is the target language, which can be a literal value or a column reference.

# Constant target language
reviews.select("review", ai.ai_translate("review", "Chinese").alias("zh")).show(truncate=False)

# Target language resolved per row
with_lang = reviews.withColumn("target_lang", lit("Japanese"))
with_lang.select(
    "review",
    ai.ai_translate("review", col("target_lang")).alias("translated"),
).show(truncate=False)

Text summarization: The second argument of ai_summarize is the target summary length, in tokens.

reviews.select("review", ai.ai_summarize("review", 20).alias("summary")).show(truncate=False)

Similarity: The second argument of ai_similarity is the reference text or reference column. A string is treated as a literal value, and col('...') is treated as a column reference.

reviews.select(
    "review",
    ai.ai_similarity("review", "The food was delicious").alias("score"),
).show(truncate=False)

Grammar correction: ai_fix_grammar is often used to clean up user-generated content.

typo_df = spark.createDataFrame([("I has went to the store yesterday",)], ["text"])
typo_df.select("text", ai.ai_fix_grammar("text").alias("corrected")).show(truncate=False)

URI mode (recommended)

Pass image URIs by using the ai_query multimodal arguments data and data_type. In the following example, the path_df DataFrame has an image_path column that stores the OSS URIs of images.

path_df = spark.createDataFrame([("oss://bucket/img1.jpg",)], ["image_path"])
path_df.select(
    "image_path",
    ai.ai_query(
        lit("Identify the license plate number. If it is illegible, return 'Unrecognized'."),
        service_name="qwen3.7-plus",
        data="image_path",
        data_type="uri",
    ).alias("plate_number"),
).show(truncate=False)

To generate image embeddings, use ai_embedding_multimodal, which supports scenarios such as multimodal RAG and image-based search:

path_df.select(
    "image_path",
    ai.ai_embedding_multimodal("image_path", data_type="uri").alias("embedding"),
).show(truncate=False)

Binary mode (for testing only)

If the input column is of the BINARY type, pass it directly to data and keep data_type at its default value binary. Common sources include:

# Source 1: the content column returned by read_files
image_df = ai.read_files("oss://bucket/images/")
image_df.select(
    "path",
    ai.ai_query(
        lit("Describe this image in one sentence."),
        service_name="qwen3.7-plus",
        data="content",
    ).alias("description"),
).show(truncate=False)

# Source 2: a BINARY field in a Paimon Blob table
blob_df = spark.sql("SELECT fileName, picture FROM blob_tbl")
blob_df.select(
    "fileName",
    ai.ai_query(
        lit("Describe this image in one sentence."),
        service_name="qwen3.7-plus",
        data="picture",
    ).alias("description"),
).show(truncate=False)

For information about creating and writing to a Paimon Blob table, see Paimon Blob.

List input: video frames and multi-image understanding

data argument also accepts an array column, so that a single call can process multiple files: array<string> corresponds to data_type="uri_list", and array<binary> corresponds to data_type="binary_list".

• Video frame understanding (default): By default, list input is processed as video frames. A group of images is treated as video frames for temporal understanding, which suits scenarios where you extract frames from a video and then analyze the key frame sequence.

• Multi-image understanding: If the list elements are independent images, such as for comparison, quality inspection, or multi-image Q&A, configure options with "list_data_type":"image" to submit the request through the general multi-image API.

# By default, list input is analyzed as video frames
frames_df = spark.createDataFrame([
    (["oss://bucket/frames/f1.jpg",
      "oss://bucket/frames/f2.jpg",
      "oss://bucket/frames/f3.jpg"],),
], ["frame_uris"])
frames_df.select(
    ai.ai_query(
        lit("Describe the action sequence in the video based on these key frames."),
        service_name="qwen3.7-plus",
        data="frame_uris",
        data_type="uri_list",
    ).alias("video_description"),
).show(truncate=False)

# Multi-image understanding: set list_data_type=image
multi_image_df = spark.createDataFrame([
    (["oss://bucket/standard.jpg", "oss://bucket/compare.jpg"],),
], ["image_uris"])
multi_image_df.select(
    ai.ai_query(
        lit("Compare the differences between these two images."),
        service_name="qwen3.7-plus",
        data="image_uris",
        data_type="uri_list",
        options='{"list_data_type":"image"}',
    ).alias("comparison"),
).show(truncate=False)

array<string> columns can also be expanded from a directory by using list_files. For example, if each row in a business table stores the frame directory of a video, you can use list_files to expand it into an array of frame paths and feed it directly into video frame understanding:

# The clip_dir column in the videos table stores the frame base directory of each video
videos = spark.createDataFrame(
    [("v001", "oss://bucket/frames/v001/")],
    ["video_id", "clip_dir"],
)

videos.select(
    "video_id",
    ai.ai_query(
        lit("Describe what happens in the video based on these key frames."),
        service_name="qwen3.7-plus",
        data=ai.list_files("clip_dir", suffix="jpg,jpeg,png", recursive=True),
        data_type="uri_list",
    ).alias("video_description"),
).show(truncate=False)

Note the following two details:

• options accepts the image_url_options field, which applies to every image in the list, for example {"list_data_type":"image","image_url_options":{"detail":"high"}}. This object cannot contain the url field.

• NULL elements in the list are ignored. An empty list, or a list in which all elements are NULL, is handled based on data_error_mode, and returns NULL by default.

If the input is a single video file, you can also pass its URI with data_type="uri", and control the frame extraction rate by using the options field fps in the configuration.

Handle input read failures

If reading the data input fails, for example because the URI format is invalid, the file does not exist, or a read exception occurs, use the options configuration field data_error_mode to control the behavior.

path_df.select(
    "image_path",
    ai.ai_query(
        lit("Identify the license plate number. If it is illegible, return 'Unrecognized'."),
        service_name="qwen3.7-plus",
        data="image_path",
        data_type="uri",
        options='{"data_error_mode":"fail"}',
    ).alias("plate_number"),
).show(truncate=False)

read_files: a table-valued function

read_files is a table-valued function that returns a DataFrame. It is convenient for reading all files under one large directory:

image_df = ai.read_files("oss://bucket/images/")
image_df.printSchema()
# root
#  |-- path: string (nullable = true)
#  |-- content: binary (nullable = true)
#  |-- modificationTime: timestamp (nullable = true)
#  |-- length: long (nullable = true)

Read only metadata, not file content: content is a BINARY column, and loading complete files row by row into executor memory is expensive. If you need only metadata such as the path, modification time, and size, omit it in select by not selecting content. Spark then prunes the column and skips reading the actual file content:

# Read only the path and file size, without loading content
meta_df = image_df.select("path", "length", "modificationTime")
meta_df.show(truncate=False)

Select the content column only when you need binary mode for validation or debugging. In production, we recommend that you use URI mode and pass only path to the ai_query or ai_embedding_multimodal data argument.

list_files: a scalar function

list_files is a regular scalar function that takes one row and returns one row, with output of type array<string>. It suits scenarios in which one path maps to multiple files. For example, all frames extracted from a video belong to the same business record. In such cases the business table usually stores only a directory path, which must be expanded one-to-many into multiple image inputs. Traditionally this required a UDF. EMR Serverless Spark provides the list_files function for this purpose. Together with data_type="uri_list", you can feed the result directly into video frame understanding:

# The clip_dir column in the videos table stores the frame directory of each video
videos = spark.createDataFrame(
    [
        ("v001", "oss://bucket/frames/v001/"),
        ("v002", "oss://bucket/frames/v002/"),
    ],
    ["video_id", "clip_dir"],
)

# List all frame images in each directory to get an array<string> column
with_frames = videos.select(
    "video_id",
    ai.list_files("clip_dir", suffix="jpg,jpeg,png", recursive=True).alias("frame_uris"),
)

# Pass the column to uri_list to analyze each video as frames
with_frames.select(
    "video_id",
    ai.ai_query(
        lit("Describe what happens in the video based on these key frames."),
        service_name="qwen3.7-plus",
        data="frame_uris",
        data_type="uri_list",
    ).alias("video_description"),
).show(truncate=False)

list_files supports the following common optional arguments:

• suffix: A comma-separated allowlist of file name extensions, for example "jpg,jpeg,png".

• recursive: Specifies whether to list subdirectories recursively. Default value: False.

• ignore_error: Specifies whether to return an empty array if the directory does not exist or listing fails. Default value: False. If listing fails, an error is thrown.

Which function to use

• To expand one large directory into multiple rows and process each file as an independent row, use read_files.

• If each row in the business table is already bound to a directory that contains multiple files, and you want to keep the row granularity of the business table and process the file list as an array column, use list_files.

Is there any difference between ai.ai_query("text") and ai.ai_query(col("text"))?

No. The first positional argument is a data argument and follows the PySpark column name convention, so a string is automatically interpreted as a column name. This is consistent with the behavior of df.select("text").

How do I write named arguments in PySpark if they are used in SQL?

Use Python keyword arguments. For example, ai_query('t', service_name => 'qwen', data => path, data_type => 'uri') in SQL corresponds to ai.ai_query("t", service_name="qwen", data="path", data_type="uri") in PySpark. Argument names match the SQL named arguments exactly, for example service_name rather than model_service_name.

What should I pass to the options argument?

options must be a valid JSON object string, and the root node must be an object. Known fields, such as system_prompt, temperature, and dimensions, are recognized by EMR Serverless Spark. Unknown fields are passed through to the model service as-is, so that you can use extended capabilities on the model side. A JSON syntax error, a non-object root node, or trailing content causes an error during local validation.

How do I prevent a single failure from failing the entire job?

Every AI function has a twin version with the try_ prefix, for example try_ai_query, try_ai_embedding, and try_ai_embedding_multimodal. These functions return NULL for a row that fails, and other rows still produce results. They suit noisy data or jobs that are sensitive to SLA.