---
schema_version: "1"
id: 741dfb92f6f28a9bf18d7e49cf529494038df3df72e4fcd168abb3e9b8895622
title: Call Predictive Models in Batch Transform Scripts
collection: developer
version: current
locale: en-us
portal_url: https://docs.salesforce.com/developer/current/en-us/data-cloud/data-cloud-code-ext/guides/use-predictive-model-in-custom-script.md
canonical_url: https://developer.salesforce.com/docs/data/data-cloud-code-ext/guide/use-predictive-model-in-custom-script.html
source_site: developer.salesforce.com
content_type: documentation
content_hash: 55f22e135fddcdeca77f3c23c335b278c7840a1d08f1c1cc2f962d405c0c7e4c
products:
  - Data 360 Code Extension Guide
  - Data 360 Developer Documentation
---
# 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](https://help.salesforce.com/s/articleView?id=sf.c360_a_dc_editions.htm&type=5). |

| Permission Sets Needed |  |
| :---- | :---- |
| To call predictive models from custom scripts: | Permission set: <ul><li>Data Cloud Architect</li></ul> |

## Prerequisites

- Set up Salesforce CLI and the Data Cloud Code Extension plugin. See [Set Up Salesforce CLI for Code Extension](https://developer.salesforce.com/docs/data/data-cloud-code-ext/guide/set-up-sdk.md).
- Initialize a script package and make sure that you can run local validation commands. See [Write and Validate Custom Scripts](https://developer.salesforce.com/docs/data/data-cloud-code-ext/guide/author-custom-script.md).
- 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](https://developer.salesforce.com/docs/ai/agentforce/guide/models-api-names.html).
- Configure an external client app to authenticate calls to predictive models. Complete only the [Create a Salesforce App](https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app) and [Obtain Credentials](https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#obtain-credentials) sections of [Get Started with the Agent API](https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html).

## Use Predictive Models in Your Script

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.

**note:** 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.

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.

```python
from pyspark.sql.functions import (
    col,
    lit,
    when,
    datediff,
    current_date,
    round as spark_round,
    get_json_object,
)

from datacustomcode.client import Client, einstein_predict_col
from datacustomcode.einstein_predictions.types import PredictionType
from datacustomcode.io.writer.base import WriteMode

def main():
    client = Client()
    df = client.read_dlo("Customer_Activity__dll")
    prediction_model_api_name = "Customer_Churn_Prediction"
```

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.

```python
    df = (
        df
        .withColumn(
            "usage_decline_30_days",
            spark_round(((col("activitycount120days__c") - col("activitycount30days__c")) / col("activitycount120days__c")) * 100, 2),
        )
        .withColumn(
            "usage_decline_60_days",
            spark_round(((col("activitycount120days__c") - col("activitycount60days__c")) / col("activitycount120days__c")) * 100, 2),
        )
        .withColumn(
            "usage_decline_90_days",
            spark_round(((col("activitycount120days__c") - col("activitycount90days__c")) / col("activitycount120days__c")) * 100, 2),
        )
        .withColumn("days_since_last_login", datediff(current_date(), col("lastlogindate__c")))
        .withColumn(
            "feature_adoption_score",
            spark_round((col("featuresusedcount__c") / col("totalfeaturesavailable__c")) * 100, 0),
        )
        .withColumn("account_age", datediff(current_date(), col("signupdate__c")))
    )
```

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.

```python
    pred_col = einstein_predict_col(
        prediction_model_api_name,
        PredictionType.BINARY_CLASSIFICATION,
        {
            "UsageDecline30Days": col("usage_decline_30_days"),
            "UsageDecline60Days": col("usage_decline_60_days"),
            "UsageDecline90Days": col("usage_decline_90_days"),
            "DaysSinceLastLogin": col("days_since_last_login"),
            "SupportTicketCount": col("supportticketcount__c"),
            "FeatureAdoptionScore": col("feature_adoption_score"),
            "AccountAge": col("account_age"),
            "SubscriptionTier": col("subscriptiontier__c"),
        },
    )
```

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.

```python
    df = (
        df
        .withColumn("pred", pred_col)
        .withColumn("result_type", get_json_object(col("pred.response"), "$.results[0].type"))
        .withColumn(
            "churn_probability__c",
            when(
                (col("pred.status") == "SUCCESS")
                & (col("result_type") != "PredictionFailure"),
                get_json_object(
                    col("pred.response"),
                    "$.results[0].prediction.classProbabilities[0].probability",
                ).cast("double"),
            ).otherwise(lit(None).cast("double")),
        )
        .drop("pred", "result_type")
    )
```

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

```python
    df_output = df.select(
        col("customerid__c").alias("customer_id__c"),
        col("subscriptiontier__c").alias("subscription_tier__c"),
        col("usage_decline_90_days").alias("usage_decline_90_days__c"),
        col("days_since_last_login").alias("days_since_last_login__c"),
        col("feature_adoption_score").alias("feature_adoption_score__c"),
        col("account_age").alias("account_age__c"),
        col("churn_probability__c"),
    )
    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.

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.

```bash
sf org login web \
  --alias myorg \
  --instance-url https://{MY_DOMAIN_URL} \
  --client-id {CONSUMER_KEY} \
  --scopes "sfap_api api"
```

   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.

```bash
sf 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()`](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.repartition.html) 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.

## See Also

- [Write and Validate Custom Scripts](https://developer.salesforce.com/docs/data/data-cloud-code-ext/guide/author-custom-script.md)
- [Considerations When Writing Code Extension](https://developer.salesforce.com/docs/data/data-cloud-code-ext/guide/considerations-when-writing-custom-code.md)
- [Call Generative AI Models in Batch Transform Scripts](https://developer.salesforce.com/docs/data/data-cloud-code-ext/guide/use-llm-gateway-in-custom-script.md)
- [Invoke Code Extension by Creating a Batch Data Transform](https://developer.salesforce.com/docs/data/data-cloud-code-ext/guide/create-batch-transform.md)
- [Get Started with Agent API](https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html)