Call Predictive Models in Batch Transform Scripts

Score, classify, or rank records during a batch data transform. Pass structured feature columns from a data lake object (DLO) or data model object (DMO) to a configured predictive model in AI Models (formerly Einstein Studio). Then write the predictions to a target DLO or DMO.

• Set up Salesforce CLI and the Data Cloud Code Extension plugin. See Set Up Salesforce CLI for Code Extension.
• Initialize a script package and make sure that you can run local validation commands. See Write and Validate Custom Scripts.
• Use a Salesforce org that has at least one predictive model deployed and activated in AI Models.
• Note the API name of the predictive model that you want to call in your script. See Look Up the API Name.
• Configure an external client app to authenticate calls to predictive models. Complete only the Create a Salesforce App and Obtain Credentials sections of Get Started with the Agent API.

In the main() function of payload/entrypoint.py in your initialized script package, implement this flow:

1. Read source data.
2. Calculate the features that your model expects.
3. Build a prediction column.
4. Parse the prediction from the response.
5. Write the results to a target object.

For a reference implementation, see example/payload/entrypoint.py in the same package.

1. Read your source data from a DLO or DMO and specify the API name of the deployed predictive model to call. read_dlo() returns a DataFrame, a distributed table of rows and columns that Spark uses to hold your data.

2. Calculate the features that your model expects. Built-in transforms can’t derive trends such as usage decline, so use Python to compute the trends from the raw activity data.

3. Build the prediction column with einstein_predict_col(), a built-in user-defined function (UDF). The prediction column is a new column expression that scores every record when you add it to your DataFrame in the next step. In the feature mapping, map each feature name that your model expects to the DataFrame column that holds its value. Set the prediction type to match how you built the model. PredictionType supports REGRESSION, BINARY_CLASSIFICATION, CLASSIFICATION, MULTI_OUTCOME, and CLUSTERING. The function returns a struct with status, response, error_code, and error_message fields for each row.

4. Add the prediction column and parse the value from the JSON response. Store the prediction struct in a single column so that the model runs only once per row. In this code, .withColumn("pred", pred_col) adds the struct, and .drop("pred", "result_type") removes the intermediate columns after you read from them. The response field is a JSON string that you parse with get_json_object. Read the predicted value with the path for your model’s prediction type:

   • BINARY_CLASSIFICATION and CLASSIFICATION: $.results[0].prediction.classProbabilities[0].probability
   • REGRESSION: $.results[0].prediction.predictedValue

Each result also carries a result type at $.results[0].type, which is PredictionFailure when the model fails to score a row. Because per-row failures don’t stop the entire job, read the predicted value only when the row’s status field is SUCCESS and the result type isn’t PredictionFailure. Otherwise, set the value to null.

5. Write the enriched data back to your target DLO or DMO.

To run a one-off prediction without a DataFrame, use client.einstein_predict() with the model API name, the prediction type, and a mapping of feature names to literal values. The method returns the parsed response as a dictionary and raises EinsteinPredictionsCallError when the call fails.

For failed predictions, decide whether to keep the null fallback shown earlier or raise a controlled error. Don’t log sensitive data.

Before you deploy your script to Data 360, test it locally against your sandbox.

1. From a terminal in your script package root, log in to your org with the external client app credentials. Salesforce CLI opens your default browser for Salesforce login and then saves the session on your computer so that your local script run can authenticate.

Replace {MY_DOMAIN_URL} with your org domain and {CONSUMER_KEY} with the consumer key for your external client app.

2. Run the script locally to test it against data in Data 360.

Predictions run as API calls, with one call per row when you score each record in a DataFrame. To increase throughput for larger datasets:

• Repartition your DataFrame with .repartition() to spread the work across more parallel worker processes (executors). Without repartitioning, requests can run one after another on only a few workers and create a bottleneck.
• Choose a larger compute type so that more CPUs are available to share the load.

• Write and Validate Custom Scripts
• Considerations When Writing Code Extension
• Call Generative AI Models in Batch Transform Scripts
• Invoke Code Extension by Creating a Batch Data Transform
• Get Started with Agent API