salesforce-data-customcode 6.0.5.dev1__tar.gz → 6.0.6.dev1__tar.gz

salesforce_data_customcode-6.0.5.dev1/README.md → salesforce_data_customcode-6.0.6.dev1/PKG-INFO

@@ -1,4 +1,29 @@
-# Data Cloud Custom Code SDK (BETA)
+Metadata-Version: 2.4
+Name: salesforce-data-customcode
+Version: 6.0.6.dev1
+Summary: Data Cloud Custom Code SDK
+License-Expression: Apache-2.0
+License-File: LICENSE.txt
+Requires-Python: >=3.10,<3.12
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Operating System :: Unix
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: click (>=8.1.8,<9.0.0)
+Requires-Dist: loguru (>=0.7.3,<0.8.0)
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: pydantic (==2.13.1)
+Requires-Dist: pyspark (==3.5.1)
+Requires-Dist: pyyaml (>=6.0,<7.0)
+Requires-Dist: requests (==2.33.1)
+Requires-Dist: salesforce-cdp-connector (>=1.0.19)
+Requires-Dist: setuptools_scm (>=7.1.0,<8.0.0)
+Description-Content-Type: text/markdown
+
+# Data Cloud Custom Code SDK
 
 This package provides a development kit for creating custom data transformations in [Data Cloud](https://www.salesforce.com/data/). It allows you to write your own data processing logic in Python while leveraging Data Cloud's infrastructure for data access and running data transformations, mapping execution into Data Cloud data structures like [Data Model Objects](https://help.salesforce.com/s/articleView?id=data.c360_a_data_model_objects.htm&type=5) and [Data Lake Objects](https://help.salesforce.com/s/articleView?id=sf.c360_a_data_lake_objects.htm&language=en_US&type=5).
 
@@ -11,7 +36,7 @@ Use of this project with Salesforce is subject to the [TERMS OF USE](./TERMS_OF_
 - **Python 3.11 only** (currently supported version - if your system version is different, we recommend using [pyenv](https://github.com/pyenv/pyenv) to configure 3.11)
 - JDK 17
 - Docker support like [Docker Desktop](https://docs.docker.com/desktop/)
-- A salesforce org with some DLOs or DMOs with data and this feature enabled (it is not GA)
+- A salesforce org with some DLOs or DMOs with data and this feature enabled
 - **One of the following** for authentication:
   - A Salesforce org already authenticated via the [Salesforce CLI](https://developer.salesforce.com/tools/salesforcecli)
     (simplest — no External Client App needed)
@@ -91,7 +116,7 @@ After modifying the `entrypoint.py` as needed, using any dependencies you add in
 ```zsh
 cd my_package
 datacustomcode scan ./payload/entrypoint.py
-datacustomcode deploy --path ./payload --name my_custom_script --cpu-size CPU_L
+datacustomcode deploy --path ./payload --name my_custom_script --cpu-size CPU_L --sf-cli-org myorg
 ```
 
 > [!TIP]
@@ -320,6 +345,7 @@ Options:
 - `--description TEXT`: Description of the transformation job (default: "")
 - `--network TEXT`: docker network (default: "default")
 - `--cpu-size TEXT`: CPU size for the deployment (default: `CPU_2XL`). Available options: CPU_L(Large), CPU_XL(Extra Large), CPU_2XL(2X Large), CPU_4XL(4X Large)
+- `--sf-cli-org TEXT`: Salesforce CLI org alias or username (e.g. `myorg`). Fetches credentials via `sf org display` — no `datacustomcode configure` step needed. Takes precedence over `--profile` if both are supplied.
 - `--function-invoke-opt TEXT`: Currently we support only `UnstructuredChunking` for functions.
 
 
@@ -372,6 +398,56 @@ datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
 ```
 
 
+## Testing Einstein Predictions
+
+You can use AI models configured in Einstein Studio to score your data while
+transforming it. As with the LLM Gateway, there are two flavors: a one-shot
+scalar call (`client.einstein_predict`) and a per-row column helper
+(`einstein_predict_col`). Below is a sample code example:
+
+```
+from datacustomcode.client import Client, einstein_predict_col
+from datacustomcode.einstein_predictions.types import PredictionType
+
+
+def main():
+  client = Client()
+  df = client.read_dlo("Input__dll")
+  # einstein_predict_col returns a struct
+  # {status, response, error_code, error_message} per row, so per-row
+  # failures don't abort the Spark job. `response` is the prediction
+  # payload as a JSON string. Pick the field you want with [].
+  df_scored = df.withColumn(
+    "prediction__c",
+    einstein_predict_col(
+        "my_regression_model",  # An AI model in your org
+        PredictionType.REGRESSION,
+        {"square_feet": col("square_feet__c"), "beds": col("beds__c")},
+    )["response"],
+  )
+
+  dlo_name = "Output_dll"
+  client.write_to_dlo(dlo_name, df_scored, write_mode=WriteMode.APPEND)
+
+  # One-shot scalar prediction returns the response payload as a dict
+  prediction = client.einstein_predict(
+    "my_regression_model",
+    PredictionType.REGRESSION,
+    {"square_feet": 1800, "beds": 3},
+  )
+
+if __name__ == "__main__":
+  main()
+```
+
+Testing this code locally uses the same External Client App setup described in
+[Testing LLM Gateway](#testing-llm-gateway). Once your `myorg` alias is set up,
+run:
+```
+datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
+```
+
+
 ## Docker usage
 
 The SDK provides Docker-based development options that allow you to test your code in an environment that closely resembles Data Cloud's execution environment.
@@ -483,7 +559,7 @@ sf --version
 **Browser-based (recommended for developer orgs and sandboxes):**
 ```zsh
 # Production / Developer Edition
-sf org login web --alias myorg
+sf org login web --alias myorg --instance-url <your-instance-url>
 
 # Sandbox
 sf org login web --alias mysandbox --instance-url https://test.salesforce.com
@@ -513,3 +589,4 @@ If you're using OAuth Tokens authentication, the initial configure will retrieve
 
 - [Troubleshooting](./docs/troubleshooting.md)
 - [For Contributors](./FOR_CONTRIBUTORS.md)
+

salesforce_data_customcode-6.0.5.dev1/PKG-INFO → salesforce_data_customcode-6.0.6.dev1/README.md

@@ -1,29 +1,4 @@
-Metadata-Version: 2.4
-Name: salesforce-data-customcode
-Version: 6.0.5.dev1
-Summary: Data Cloud Custom Code SDK
-License-Expression: Apache-2.0
-License-File: LICENSE.txt
-Requires-Python: >=3.10,<3.12
-Classifier: Development Status :: 4 - Beta
-Classifier: Operating System :: Unix
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: click (>=8.1.8,<9.0.0)
-Requires-Dist: loguru (>=0.7.3,<0.8.0)
-Requires-Dist: numpy
-Requires-Dist: pandas
-Requires-Dist: pydantic (==2.13.1)
-Requires-Dist: pyspark (==3.5.1)
-Requires-Dist: pyyaml (>=6.0,<7.0)
-Requires-Dist: requests (==2.33.1)
-Requires-Dist: salesforce-cdp-connector (>=1.0.19)
-Requires-Dist: setuptools_scm (>=7.1.0,<8.0.0)
-Description-Content-Type: text/markdown
-
-# Data Cloud Custom Code SDK (BETA)
+# Data Cloud Custom Code SDK
 
 This package provides a development kit for creating custom data transformations in [Data Cloud](https://www.salesforce.com/data/). It allows you to write your own data processing logic in Python while leveraging Data Cloud's infrastructure for data access and running data transformations, mapping execution into Data Cloud data structures like [Data Model Objects](https://help.salesforce.com/s/articleView?id=data.c360_a_data_model_objects.htm&type=5) and [Data Lake Objects](https://help.salesforce.com/s/articleView?id=sf.c360_a_data_lake_objects.htm&language=en_US&type=5).
 
@@ -36,7 +11,7 @@ Use of this project with Salesforce is subject to the [TERMS OF USE](./TERMS_OF_
 - **Python 3.11 only** (currently supported version - if your system version is different, we recommend using [pyenv](https://github.com/pyenv/pyenv) to configure 3.11)
 - JDK 17
 - Docker support like [Docker Desktop](https://docs.docker.com/desktop/)
-- A salesforce org with some DLOs or DMOs with data and this feature enabled (it is not GA)
+- A salesforce org with some DLOs or DMOs with data and this feature enabled
 - **One of the following** for authentication:
   - A Salesforce org already authenticated via the [Salesforce CLI](https://developer.salesforce.com/tools/salesforcecli)
     (simplest — no External Client App needed)
@@ -116,7 +91,7 @@ After modifying the `entrypoint.py` as needed, using any dependencies you add in
 ```zsh
 cd my_package
 datacustomcode scan ./payload/entrypoint.py
-datacustomcode deploy --path ./payload --name my_custom_script --cpu-size CPU_L
+datacustomcode deploy --path ./payload --name my_custom_script --cpu-size CPU_L --sf-cli-org myorg
 ```
 
 > [!TIP]
@@ -345,6 +320,7 @@ Options:
 - `--description TEXT`: Description of the transformation job (default: "")
 - `--network TEXT`: docker network (default: "default")
 - `--cpu-size TEXT`: CPU size for the deployment (default: `CPU_2XL`). Available options: CPU_L(Large), CPU_XL(Extra Large), CPU_2XL(2X Large), CPU_4XL(4X Large)
+- `--sf-cli-org TEXT`: Salesforce CLI org alias or username (e.g. `myorg`). Fetches credentials via `sf org display` — no `datacustomcode configure` step needed. Takes precedence over `--profile` if both are supplied.
 - `--function-invoke-opt TEXT`: Currently we support only `UnstructuredChunking` for functions.
 
 
@@ -397,6 +373,56 @@ datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
 ```
 
 
+## Testing Einstein Predictions
+
+You can use AI models configured in Einstein Studio to score your data while
+transforming it. As with the LLM Gateway, there are two flavors: a one-shot
+scalar call (`client.einstein_predict`) and a per-row column helper
+(`einstein_predict_col`). Below is a sample code example:
+
+```
+from datacustomcode.client import Client, einstein_predict_col
+from datacustomcode.einstein_predictions.types import PredictionType
+
+
+def main():
+  client = Client()
+  df = client.read_dlo("Input__dll")
+  # einstein_predict_col returns a struct
+  # {status, response, error_code, error_message} per row, so per-row
+  # failures don't abort the Spark job. `response` is the prediction
+  # payload as a JSON string. Pick the field you want with [].
+  df_scored = df.withColumn(
+    "prediction__c",
+    einstein_predict_col(
+        "my_regression_model",  # An AI model in your org
+        PredictionType.REGRESSION,
+        {"square_feet": col("square_feet__c"), "beds": col("beds__c")},
+    )["response"],
+  )
+
+  dlo_name = "Output_dll"
+  client.write_to_dlo(dlo_name, df_scored, write_mode=WriteMode.APPEND)
+
+  # One-shot scalar prediction returns the response payload as a dict
+  prediction = client.einstein_predict(
+    "my_regression_model",
+    PredictionType.REGRESSION,
+    {"square_feet": 1800, "beds": 3},
+  )
+
+if __name__ == "__main__":
+  main()
+```
+
+Testing this code locally uses the same External Client App setup described in
+[Testing LLM Gateway](#testing-llm-gateway). Once your `myorg` alias is set up,
+run:
+```
+datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
+```
+
+
 ## Docker usage
 
 The SDK provides Docker-based development options that allow you to test your code in an environment that closely resembles Data Cloud's execution environment.
@@ -508,7 +534,7 @@ sf --version
 **Browser-based (recommended for developer orgs and sandboxes):**
 ```zsh
 # Production / Developer Edition
-sf org login web --alias myorg
+sf org login web --alias myorg --instance-url <your-instance-url>
 
 # Sandbox
 sf org login web --alias mysandbox --instance-url https://test.salesforce.com
@@ -538,4 +564,3 @@ If you're using OAuth Tokens authentication, the initial configure will retrieve
 
 - [Troubleshooting](./docs/troubleshooting.md)
 - [For Contributors](./FOR_CONTRIBUTORS.md)
-

{salesforce_data_customcode-6.0.5.dev1 → salesforce_data_customcode-6.0.6.dev1}/pyproject.toml

@@ -5,7 +5,7 @@ requires = ["poetry-core>=2.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]
 [project]
 authors = []
 classifiers = [
-  "Development Status :: 4 - Beta",
+  "Development Status :: 5 - Production/Stable",
   "Operating System :: Unix",
   "Programming Language :: Python",
   "Programming Language :: Python :: 3",
@@ -18,7 +18,7 @@ license = "Apache-2.0"
 name = "salesforce-data-customcode"
 readme = "README.md"
 requires-python = ">=3.10,<3.12"
-version = "6.0.5.dev1"
+version = "6.0.6.dev1"
 
 [tool.black]
 exclude = '''

{salesforce_data_customcode-6.0.5.dev1 → salesforce_data_customcode-6.0.6.dev1}/src/datacustomcode/__init__.py

@@ -17,10 +17,13 @@ __all__ = [
     "AuthType",
     "Client",
     "Credentials",
+    "DefaultSparkEinsteinPredictions",
     "DefaultSparkLLMGateway",
     "PrintDataCloudWriter",
     "QueryAPIDataCloudReader",
+    "SparkEinsteinPredictions",
     "SparkLLMGateway",
+    "einstein_predict_col",
     "llm_gateway_generate_text_col",
 ]
 
@@ -59,4 +62,16 @@ def __getattr__(name: str):
         from datacustomcode.client import llm_gateway_generate_text_col
 
         return llm_gateway_generate_text_col
+    elif name == "SparkEinsteinPredictions":
+        from datacustomcode.einstein_predictions import SparkEinsteinPredictions
+
+        return SparkEinsteinPredictions
+    elif name == "DefaultSparkEinsteinPredictions":
+        from datacustomcode.einstein_predictions import DefaultSparkEinsteinPredictions
+
+        return DefaultSparkEinsteinPredictions
+    elif name == "einstein_predict_col":
+        from datacustomcode.client import einstein_predict_col
+
+        return einstein_predict_col
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

{salesforce_data_customcode-6.0.5.dev1 → salesforce_data_customcode-6.0.6.dev1}/src/datacustomcode/client.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
+    Any,
     ClassVar,
     Dict,
     Optional,
@@ -24,6 +25,7 @@ from typing import (
 )
 
 from datacustomcode.config import config
+from datacustomcode.einstein_predictions_config import spark_einstein_predictions_config
 from datacustomcode.file.path.default import DefaultFindFilePath
 from datacustomcode.io.reader.base import BaseDataCloudReader
 from datacustomcode.llm_gateway_config import spark_llm_gateway_config
@@ -34,6 +36,8 @@ if TYPE_CHECKING:
 
     from pyspark.sql import Column, DataFrame as PySparkDataFrame
 
+    from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
+    from datacustomcode.einstein_predictions.types import PredictionType
     from datacustomcode.io.reader.base import BaseDataCloudReader
     from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
     from datacustomcode.llm_gateway.spark_base import SparkLLMGateway
@@ -99,6 +103,70 @@ def llm_gateway_generate_text_col(
     return gateway.llm_gateway_generate_text_col(template, values, model_id=model_id)
 
 
+def _build_spark_einstein_predictions() -> "SparkEinsteinPredictions":
+    """Instantiate the SDK-configured :class:`SparkEinsteinPredictions`.
+
+    Raises:
+        RuntimeError: If no ``spark_einstein_predictions_config`` has been loaded.
+    """
+    cfg = spark_einstein_predictions_config.spark_einstein_predictions_config
+    if cfg is None:
+        raise RuntimeError(
+            "spark_einstein_predictions_config is not configured. Add a "
+            "'spark_einstein_predictions_config' section to config.yaml."
+        )
+    return cfg.to_object()
+
+
+def einstein_predict_col(
+    model_api_name: str,
+    prediction_type: "PredictionType",
+    features: Dict[str, "Column"],
+    settings: Optional[Dict[str, Any]] = None,
+) -> "Column":
+    """Build a Spark Column that runs an Einstein prediction per row.
+
+    The returned Column yields a struct ``{status, response, error_code,
+    error_message}`` for each row. Use ``[...]`` (or ``getField``) to pick the
+    field you want, e.g. ``einstein_predict_col(...)["response"]``. ``response``
+    holds the prediction response payload as a JSON string. Per-row failures
+    populate ``status`` / ``error_code`` / ``error_message`` so a single bad row
+    does not abort the whole Spark job.
+
+    Example:
+
+        >>> from datacustomcode.einstein_predictions.types import PredictionType
+        >>> result = einstein_predict_col(
+        ...     "my_regression_model",
+        ...     PredictionType.REGRESSION,
+        ...     {"square_feet": col("square_feet__c"), "beds": col("beds__c")},
+        ... )
+        >>> df.withColumn("prediction__c", result["response"])
+        >>> # …or keep the struct around and inspect failures:
+        >>> df.withColumn("pred", result).select(
+        ...     "pred.status", "pred.response", "pred.error_message"
+        ... )
+
+    Args:
+        model_api_name: API name of the Einstein model to invoke.
+        prediction_type: The :class:`PredictionType` of the model.
+        features: A mapping from model feature column name to a Spark ``Column``
+            supplying that feature's per-row value.
+        settings: Optional prediction settings forwarded to the model.
+
+    Returns:
+        A Spark ``Column`` of ``StructType`` with fields ``status``,
+        ``response``, ``error_code``, and ``error_message`` (all nullable
+        strings). On success, ``status == "SUCCESS"`` and ``response`` holds
+        the JSON-serialized prediction payload; on failure, ``status ==
+        "ERROR"`` and the ``error_*`` fields carry diagnostic detail.
+    """
+    predictions = Client()._get_spark_einstein_predictions()
+    return predictions.einstein_predict_col(
+        model_api_name, prediction_type, features, settings=settings
+    )
+
+
 class DataCloudObjectType(Enum):
     DLO = "dlo"
     DMO = "dmo"
@@ -158,6 +226,8 @@ class Client:
         reader: A custom reader to use for reading Data Cloud objects.
         writer: A custom writer to use for writing Data Cloud objects.
         spark_llm_gateway: Optional custom :class:`SparkLLMGateway`.
+        spark_einstein_predictions: Optional custom
+            :class:`SparkEinsteinPredictions`.
 
     Example:
     >>> client = Client()
@@ -172,6 +242,7 @@ class Client:
     _writer: BaseDataCloudWriter
     _file: DefaultFindFilePath
     _spark_llm_gateway: Optional[SparkLLMGateway]
+    _spark_einstein_predictions: Optional[SparkEinsteinPredictions]
     _data_layer_history: dict[DataCloudObjectType, set[str]]
     _code_type: str
 
@@ -181,12 +252,14 @@ class Client:
         writer: Optional[BaseDataCloudWriter] = None,
         spark_provider: Optional[BaseSparkSessionProvider] = None,
         spark_llm_gateway: Optional[SparkLLMGateway] = None,
+        spark_einstein_predictions: Optional[SparkEinsteinPredictions] = None,
         code_type: str = "script",
     ) -> Client:
 
         if cls._instance is None:
             cls._instance = super().__new__(cls)
             cls._instance._spark_llm_gateway = spark_llm_gateway
+            cls._instance._spark_einstein_predictions = spark_einstein_predictions
             # Initialize Readers and Writers from config
             # and/or provided reader and writer
             if reader is None or writer is None:
@@ -358,6 +431,49 @@ class Client:
             self._spark_llm_gateway = _build_spark_llm_gateway()
         return self._spark_llm_gateway
 
+    def einstein_predict(
+        self,
+        model_api_name: str,
+        prediction_type: "PredictionType",
+        features: Dict[str, Any],
+        settings: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Issue a one-shot Einstein prediction. This is the scalar counterpart
+        to :func:`einstein_predict_col`: it runs **once** — not per row. Use the
+        column helper method instead when you want to fan a prediction out
+        across every row of a DataFrame.
+
+        Example:
+
+            >>> from datacustomcode.einstein_predictions.types import PredictionType
+            >>> response = Client().einstein_predict(
+            ...     "my_regression_model",
+            ...     PredictionType.REGRESSION,
+            ...     {"square_feet": 1800, "beds": 3},
+            ... )
+
+        Args:
+            model_api_name: API name of the Einstein model to invoke.
+            prediction_type: The :class:`PredictionType` of the model.
+            features: A mapping from model feature column name to a single
+                scalar value (``str`` / ``float`` / ``bool``).
+            settings: Optional prediction settings forwarded to the model.
+
+        Returns:
+            The prediction response payload as a plain Python ``dict``.
+
+        Raises:
+            EinsteinPredictionsCallError: If the prediction call fails.
+        """
+        return self._get_spark_einstein_predictions().einstein_predict(
+            model_api_name, prediction_type, features, settings=settings
+        )
+
+    def _get_spark_einstein_predictions(self) -> SparkEinsteinPredictions:
+        if self._spark_einstein_predictions is None:
+            self._spark_einstein_predictions = _build_spark_einstein_predictions()
+        return self._spark_einstein_predictions
+
     def _validate_data_layer_history_does_not_contain(
         self, data_cloud_object_type: DataCloudObjectType
     ) -> None:

{salesforce_data_customcode-6.0.5.dev1 → salesforce_data_customcode-6.0.6.dev1}/src/datacustomcode/config.yaml

@@ -24,6 +24,9 @@ einstein_predictions_config:
   options:
     credentials_profile: default
 
+spark_einstein_predictions_config:
+  type_config_name: DefaultSparkEinsteinPredictions
+
 llm_gateway_config:
   type_config_name: DefaultLLMGateway
   options:

{salesforce_data_customcode-6.0.5.dev1 → salesforce_data_customcode-6.0.6.dev1}/src/datacustomcode/einstein_predictions/__init__.py

@@ -14,9 +14,17 @@
 # limitations under the License.
 
 from datacustomcode.einstein_predictions.base import EinsteinPredictions
+from datacustomcode.einstein_predictions.errors import EinsteinPredictionsCallError
 from datacustomcode.einstein_predictions.impl.default import DefaultEinsteinPredictions
+from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
+from datacustomcode.einstein_predictions.spark_default import (
+    DefaultSparkEinsteinPredictions,
+)
 
 __all__ = [
     "DefaultEinsteinPredictions",
+    "DefaultSparkEinsteinPredictions",
     "EinsteinPredictions",
+    "EinsteinPredictionsCallError",
+    "SparkEinsteinPredictions",
 ]

salesforce_data_customcode-6.0.6.dev1/src/datacustomcode/einstein_predictions/errors.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Exceptions raised by Einstein Predictions implementations."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+
+class EinsteinPredictionsCallError(RuntimeError):
+    """Raised when an Einstein Predictions call returns an error."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status: Optional[object] = None,
+        error_code: Optional[str] = None,
+        error_message: Optional[str] = None,
+    ) -> None:
+        super().__init__(message)
+        self.status = status
+        self.error_code = error_code
+        self.error_message = error_message

salesforce_data_customcode-6.0.6.dev1/src/datacustomcode/einstein_predictions/spark_base.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+)
+
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
+if TYPE_CHECKING:
+    from pyspark.sql import Column
+
+    from datacustomcode.einstein_predictions.types import PredictionType
+
+
+class SparkEinsteinPredictions(ABC, UserExtendableNamedConfigMixin):
+    CONFIG_NAME: str
+
+    def __init__(self, **kwargs: Any) -> None:
+        pass
+
+    @abstractmethod
+    def einstein_predict(
+        self,
+        model_api_name: str,
+        prediction_type: PredictionType,
+        features: Dict[str, Any],
+        settings: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Issue a one-shot Einstein prediction and return the response data.
+
+        ``features`` maps each model feature column name to a single scalar
+        value (``str``/``float``/``bool``). The value is wrapped into a
+        single-element prediction column of the appropriate type.
+        """
+
+    @abstractmethod
+    def einstein_predict_col(
+        self,
+        model_api_name: str,
+        prediction_type: PredictionType,
+        features: Dict[str, "Column"],
+        settings: Optional[Dict[str, Any]] = None,
+    ) -> "Column":
+        """Build a Spark ``Column`` that invokes Einstein predict per row and
+        yields a struct ``{status, response, error_code, error_message}``.
+
+        ``features`` maps each model feature column name to a Spark ``Column``
+        supplying that feature's per-row value. Select an individual field,
+        e.g. ``einstein_predict_col(...)["response"]``. ``response`` holds the
+        prediction response payload as a JSON string. Returning a struct means
+        a single failing row leaves the rest of the DataFrame intact — callers
+        can inspect ``status`` / ``error_code`` per row instead of having the
+        Spark job abort.
+        """