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.

Edition Table
Available in: Developer, Enterprise, Performance, and Unlimited Editions. See Data 360 edition availability.

Permission Sets Needed
To call predictive models from custom scripts:	Permission set: Data Cloud Architect

Prerequisites

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.

Use Predictive Models in Your Script

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

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

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

These steps use a customer churn scenario. The script reads customer activity, calculates usage-trend features, predicts each customer’s churn probability, and writes the results for your sales teams. Replace the DLO and DMO names, the model API name, the features that you calculate and map, and the output columns with the values for your own model and data.

Note

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.

1from pyspark.sql.functions import (
2    col,
3    lit,
4    when,
5    datediff,
6    current_date,
7    round as spark_round,
8    get_json_object,
9)
10
11from datacustomcode.client import Client, einstein_predict_col
12from datacustomcode.einstein_predictions.types import PredictionType
13from datacustomcode.io.writer.base import WriteMode
14
15def main():
16    client = Client()
17    df = client.read_dlo("Customer_Activity__dll")
18    prediction_model_api_name = "Customer_Churn_Prediction"

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.

1df = (
2        df
3        .withColumn(
4            "usage_decline_30_days",
5            spark_round(((col("activitycount120days__c") - col("activitycount30days__c")) / col("activitycount120days__c")) * 100, 2),
6        )
7        .withColumn(
8            "usage_decline_60_days",
9            spark_round(((col("activitycount120days__c") - col("activitycount60days__c")) / col("activitycount120days__c")) * 100, 2),
10        )
11        .withColumn(
12            "usage_decline_90_days",
13            spark_round(((col("activitycount120days__c") - col("activitycount90days__c")) / col("activitycount120days__c")) * 100, 2),
14        )
15        .withColumn("days_since_last_login", datediff(current_date(), col("lastlogindate__c")))
16        .withColumn(
17            "feature_adoption_score",
18            spark_round((col("featuresusedcount__c") / col("totalfeaturesavailable__c")) * 100, 0),
19        )
20        .withColumn("account_age", datediff(current_date(), col("signupdate__c")))
21    )

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.

1pred_col = einstein_predict_col(
2        prediction_model_api_name,
3        PredictionType.BINARY_CLASSIFICATION,
4        {
5            "UsageDecline30Days": col("usage_decline_30_days"),
6            "UsageDecline60Days": col("usage_decline_60_days"),
7            "UsageDecline90Days": col("usage_decline_90_days"),
8            "DaysSinceLastLogin": col("days_since_last_login"),
9            "SupportTicketCount": col("supportticketcount__c"),
10            "FeatureAdoptionScore": col("feature_adoption_score"),
11            "AccountAge": col("account_age"),
12            "SubscriptionTier": col("subscriptiontier__c"),
13        },
14    )

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.

1df = (
2        df
3        .withColumn("pred", pred_col)
4        .withColumn("result_type", get_json_object(col("pred.response"), "$.results[0].type"))
5        .withColumn(
6            "churn_probability__c",
7            when(
8                (col("pred.status") == "SUCCESS")
9                & (col("result_type") != "PredictionFailure"),
10                get_json_object(
11                    col("pred.response"),
12                    "$.results[0].prediction.classProbabilities[0].probability",
13                ).cast("double"),
14            ).otherwise(lit(None).cast("double")),
15        )
16        .drop("pred", "result_type")
17    )

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

1df_output = df.select(
2        col("customerid__c").alias("customer_id__c"),
3        col("subscriptiontier__c").alias("subscription_tier__c"),
4        col("usage_decline_90_days").alias("usage_decline_90_days__c"),
5        col("days_since_last_login").alias("days_since_last_login__c"),
6        col("feature_adoption_score").alias("feature_adoption_score__c"),
7        col("account_age").alias("account_age__c"),
8        col("churn_probability__c"),
9    )
10    client.write_to_dlo("Customer_Churn_Risk__dll", df_output, write_mode=WriteMode.OVERWRITE)

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.

Validate Predictive Model Calls Locally

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

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.

1sf org login web \
2  --alias myorg \
3  --instance-url https://{MY_DOMAIN_URL} \
4  --client-id {CONSUMER_KEY} \
5  --scopes "sfap_api api"

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

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

1sf data-code-extension script run --entrypoint ./payload/entrypoint.py --target-org myorg

Improve Throughput for Large Datasets

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.

Call Predictive Models in Batch Transform Scripts

Prerequisites

Use Predictive Models in Your Script

Validate Predictive Model Calls Locally

Improve Throughput for Large Datasets

See Also