salesforce-data-customcode 6.0.5.dev1__py3-none-any.whl → 6.0.6.dev2__py3-none-any.whl

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.
+        """

datacustomcode/einstein_predictions/spark_default.py

@@ -0,0 +1,239 @@
+# 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
+
+import json
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+)
+
+from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
+from datacustomcode.einstein_predictions.types import (
+    PredictionColumn,
+    PredictionRequest,
+    PredictionType,
+)
+
+if TYPE_CHECKING:
+    from pyspark.sql import Column
+
+    from datacustomcode.einstein_predictions.base import EinsteinPredictions
+    from datacustomcode.einstein_predictions.types import PredictionResponse
+
+
+_STATUS_SUCCESS = "SUCCESS"
+_STATUS_ERROR = "ERROR"
+
+
+class DefaultSparkEinsteinPredictions(SparkEinsteinPredictions):
+
+    CONFIG_NAME = "DefaultSparkEinsteinPredictions"
+
+    def __init__(
+        self,
+        einstein_predictions: Optional["EinsteinPredictions"] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        if einstein_predictions is None:
+            einstein_predictions = _build_underlying_predictions()
+        self._einstein_predictions: "EinsteinPredictions" = einstein_predictions
+
+    def einstein_predict(
+        self,
+        model_api_name: str,
+        prediction_type: PredictionType,
+        features: Dict[str, Any],
+        settings: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        return _invoke_predictions(
+            self._einstein_predictions,
+            model_api_name,
+            prediction_type,
+            features,
+            settings,
+        )
+
+    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 per-row UDF that returns a struct ``{status, response,
+        error_code, error_message}`` so per-row failures do not abort the
+        Spark job. Callers select the field they want, e.g.
+        ``einstein_predict_col(...)["response"]``. ``response`` carries the
+        prediction response payload serialized as a JSON string.
+        """
+        from pyspark.sql.functions import struct, udf
+        from pyspark.sql.types import (
+            StringType,
+            StructField,
+            StructType,
+        )
+
+        feature_names = list(features.keys())
+        values_col = struct(*[features[name].alias(name) for name in feature_names])
+
+        predictions = self._einstein_predictions
+        result_schema = StructType(
+            [
+                StructField("status", StringType(), True),
+                StructField("response", StringType(), True),
+                StructField("error_code", StringType(), True),
+                StructField("error_message", StringType(), True),
+            ]
+        )
+
+        def _predict(values_row: Any) -> Dict[str, Optional[str]]:
+            if values_row is None:
+                return {
+                    "status": _STATUS_ERROR,
+                    "response": None,
+                    "error_code": None,
+                    "error_message": "features column was null for this row",
+                }
+            row_features = (
+                values_row.asDict()
+                if hasattr(values_row, "asDict")
+                else dict(values_row)
+            )
+            return _invoke_predictions_as_struct(
+                predictions,
+                model_api_name,
+                prediction_type,
+                row_features,
+                settings,
+            )
+
+        return udf(_predict, result_schema)(values_col)
+
+
+def _build_underlying_predictions() -> "EinsteinPredictions":
+    from datacustomcode.einstein_predictions_config import einstein_predictions_config
+
+    cfg = einstein_predictions_config.einstein_predictions_config
+    if cfg is None:
+        raise RuntimeError(
+            "einstein_predictions_config is not configured. Add an "
+            "'einstein_predictions_config' section to config.yaml."
+        )
+    return cfg.to_object()
+
+
+def _feature_to_prediction_column(name: str, value: Any) -> PredictionColumn:
+    """Wrap a single scalar feature value into a one-element prediction column.
+
+    The value's Python type selects the prediction column value type. ``bool``
+    is checked before ``int``/``float`` because ``bool`` is a subclass of
+    ``int`` in Python.
+    """
+    if isinstance(value, bool):
+        return PredictionColumn(column_name=name, boolean_values=[value])
+    if isinstance(value, (int, float)):
+        return PredictionColumn(column_name=name, double_values=[float(value)])
+    return PredictionColumn(column_name=name, string_values=[str(value)])
+
+
+def _build_request(
+    model_api_name: str,
+    prediction_type: PredictionType,
+    features: Dict[str, Any],
+    settings: Optional[Dict[str, Any]],
+) -> PredictionRequest:
+    prediction_columns = [
+        _feature_to_prediction_column(name, value) for name, value in features.items()
+    ]
+    return PredictionRequest(
+        prediction_type=prediction_type,
+        model_api_name=model_api_name,
+        prediction_columns=prediction_columns,
+        settings=settings,
+    )
+
+
+def _call_predictions(
+    predictions: "EinsteinPredictions",
+    model_api_name: str,
+    prediction_type: PredictionType,
+    features: Dict[str, Any],
+    settings: Optional[Dict[str, Any]],
+) -> "PredictionResponse":
+    """Build the request and dispatch it to the underlying predictions resource."""
+    request = _build_request(model_api_name, prediction_type, features, settings)
+    return predictions.predict(request)
+
+
+def _invoke_predictions(
+    predictions: "EinsteinPredictions",
+    model_api_name: str,
+    prediction_type: PredictionType,
+    features: Dict[str, Any],
+    settings: Optional[Dict[str, Any]],
+) -> Dict[str, Any]:
+    from datacustomcode.einstein_predictions.errors import EinsteinPredictionsCallError
+
+    response = _call_predictions(
+        predictions, model_api_name, prediction_type, features, settings
+    )
+    if not response.is_success:
+        error_code = _extract_error_code(response)
+        raise EinsteinPredictionsCallError(
+            f"Einstein Predictions call failed: "
+            f"status_code={response.status_code}, "
+            f"error_code={error_code!r}, message={response.data!r}",
+            status=response.status_code,
+            error_code=error_code,
+            error_message=str(response.data) if response.data else None,
+        )
+    return response.data or {}
+
+
+def _invoke_predictions_as_struct(
+    predictions: "EinsteinPredictions",
+    model_api_name: str,
+    prediction_type: PredictionType,
+    features: Dict[str, Any],
+    settings: Optional[Dict[str, Any]],
+) -> Dict[str, Optional[str]]:
+    response = _call_predictions(
+        predictions, model_api_name, prediction_type, features, settings
+    )
+    if not response.is_success:
+        return {
+            "status": _STATUS_ERROR,
+            "response": None,
+            "error_code": _extract_error_code(response),
+            "error_message": str(response.data) if response.data else None,
+        }
+    return {
+        "status": _STATUS_SUCCESS,
+        "response": json.dumps(response.data) if response.data is not None else None,
+        "error_code": None,
+        "error_message": None,
+    }
+
+
+def _extract_error_code(response: "PredictionResponse") -> Optional[str]:
+    if response.data:
+        error_code = response.data.get("errorCode")
+        if error_code is not None:
+            return str(error_code)
+    return None

datacustomcode/einstein_predictions_config.py

@@ -21,11 +21,17 @@ from typing import (
     Union,
 )
 
-from datacustomcode.common_config import BaseConfig, default_config_file
+from datacustomcode.common_config import (
+    BaseConfig,
+    BaseObjectConfig,
+    default_config_file,
+)
 from datacustomcode.einstein_platform_config import CredentialsObjectConfig
 from datacustomcode.einstein_predictions.base import EinsteinPredictions
+from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
 
 _E = TypeVar("_E", bound=EinsteinPredictions)
+_S = TypeVar("_S", bound=SparkEinsteinPredictions)
 
 
 class EinsteinPredictionsObjectConfig(CredentialsObjectConfig[_E], Generic[_E]):
@@ -54,6 +60,44 @@ class EinsteinPredictionsConfig(BaseConfig):
         return self
 
 
+class SparkEinsteinPredictionsObjectConfig(BaseObjectConfig, Generic[_S]):
+    type_to_create: ClassVar[Type[SparkEinsteinPredictions]] = SparkEinsteinPredictions  # type: ignore[type-abstract]
+
+    def to_object(self) -> SparkEinsteinPredictions:
+        type_ = self.type_to_create.subclass_from_config_name(self.type_config_name)
+        return type_(**self.options)
+
+
+class SparkEinsteinPredictionsConfig(BaseConfig):
+    spark_einstein_predictions_config: Union[
+        SparkEinsteinPredictionsObjectConfig[SparkEinsteinPredictions], None
+    ] = None
+
+    def update(
+        self, other: "SparkEinsteinPredictionsConfig"
+    ) -> "SparkEinsteinPredictionsConfig":
+        def merge(
+            config_a: Union[SparkEinsteinPredictionsObjectConfig, None],
+            config_b: Union[SparkEinsteinPredictionsObjectConfig, None],
+        ) -> Union[SparkEinsteinPredictionsObjectConfig, None]:
+            if config_a is not None and config_a.force:
+                return config_a
+            if config_b:
+                return config_b
+            return config_a
+
+        self.spark_einstein_predictions_config = merge(
+            self.spark_einstein_predictions_config,
+            other.spark_einstein_predictions_config,
+        )
+        return self
+
+
 # Global Einstein Predictions config instance
 einstein_predictions_config = EinsteinPredictionsConfig()
 einstein_predictions_config.load(default_config_file())
+
+
+# Global Spark Einstein Predictions config instance
+spark_einstein_predictions_config = SparkEinsteinPredictionsConfig()
+spark_einstein_predictions_config.load(default_config_file())

datacustomcode/io/reader/base.py

@@ -41,3 +41,45 @@ class BaseDataCloudReader(BaseDataAccessLayer):
         name: str,
         schema: Union[AtomicType, StructType, str, None] = None,
     ) -> PySparkDataFrame: ...
+
+    def read_dlo_deltas(self) -> PySparkDataFrame:
+        """Read the streaming change feed (deltas) for a Data Lake Object.
+
+        This is the streaming counterpart to :meth:`read_dlo`. It returns a
+        streaming DataFrame over the change feed the Data Cloud runtime
+        publishes for a streaming (``DELTA_SYNC``) transform. Concrete
+        streaming behavior is provided by the deployed Data Cloud runtime; the
+        base implementation raises :class:`NotImplementedError` so local
+        readers that do not support streaming fail clearly.
+
+        Returns:
+            A streaming PySpark DataFrame over the DLO change feed.
+
+        Raises:
+            NotImplementedError: If the active reader does not support streaming
+                deltas (e.g. the local development readers).
+        """
+        raise NotImplementedError(
+            "read_dlo_deltas is only supported when running in the Data Cloud "
+            "streaming runtime; the local reader does not support streaming "
+            "deltas."
+        )
+
+    def read_dmo_deltas(self) -> PySparkDataFrame:
+        """Read the streaming change feed (deltas) for a Data Model Object.
+
+        Streaming counterpart to :meth:`read_dmo`. See :meth:`read_dlo_deltas`
+        for behavior and the local-development caveat.
+
+        Returns:
+            A streaming PySpark DataFrame over the DMO change feed.
+
+        Raises:
+            NotImplementedError: If the active reader does not support streaming
+                deltas (e.g. the local development readers).
+        """
+        raise NotImplementedError(
+            "read_dmo_deltas is only supported when running in the Data Cloud "
+            "streaming runtime; the local reader does not support streaming "
+            "deltas."
+        )

datacustomcode/io/writer/base.py

@@ -22,6 +22,7 @@ from datacustomcode.io.base import BaseDataAccessLayer
 
 if TYPE_CHECKING:
     from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession
+    from pyspark.sql.streaming import StreamingQuery
 
 
 class WriteMode(str, Enum):
@@ -57,3 +58,35 @@ class BaseDataCloudWriter(BaseDataAccessLayer):
     def write_to_dmo(
         self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
     ) -> None: ...
+
+    def write_dlo_deltas(
+        self, name: str, dataframe: PySparkDataFrame
+    ) -> StreamingQuery:
+        """Write a streaming DataFrame of deltas to a Data Lake Object.
+
+        Streaming counterpart to :meth:`write_to_dlo`. Starts a streaming query
+        that writes each micro-batch to the target DLO via the Data Cloud
+        streaming sink and returns the resulting ``StreamingQuery`` handle. The
+        runtime owns the trigger and checkpoint location; callers pass only the
+        table name. Concrete streaming behavior is provided by the deployed
+        Data Cloud runtime; the base implementation raises
+        :class:`NotImplementedError`.
+
+        Args:
+            name: Target Data Lake Object name.
+            dataframe: Streaming PySpark DataFrame produced from a
+                ``read_dlo_deltas`` / ``read_dmo_deltas`` source.
+
+        Returns:
+            The started ``StreamingQuery``; the caller drives its lifecycle
+            (typically ``query.awaitTermination()``).
+
+        Raises:
+            NotImplementedError: If the active writer does not support streaming
+                deltas (e.g. the local development writers).
+        """
+        raise NotImplementedError(
+            "write_dlo_deltas is only supported when running in the Data Cloud "
+            "streaming runtime; the local writer does not support streaming "
+            "deltas."
+        )

datacustomcode/run.py

@@ -42,6 +42,15 @@ def _set_config_option(config_obj, key: str, value: Optional[str]) -> None:
         config_obj.options[key] = value
 
 
+def _read_streaming_source(config_json: dict) -> Optional[str]:
+    """Return the streaming source name from config.json's ``streamingSource``."""
+    source = config_json.get("streamingSource")
+    if not isinstance(source, dict):
+        return None
+    name = source.get("name")
+    return str(name) if name else None
+
+
 def _update_config_options(profile: Optional[str], sf_cli_org: Optional[str]):
     if sf_cli_org:
         config_key = "sf_cli_org"
@@ -125,6 +134,8 @@ def run_entrypoint(
         _set_config_option(config.reader_config, "dataspace", dataspace)
         _set_config_option(config.writer_config, "dataspace", dataspace)
 
+        config.streaming_source = _read_streaming_source(config_json)
+
     _update_config_options(profile, sf_cli_org)
 
     for dependency in dependencies:

datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py

@@ -0,0 +1,49 @@
+"""Streaming BYOC transform: read a DLO change feed and write the deltas back.
+
+This example is the streaming counterpart to a normal batch entrypoint. Instead
+of a batch ``Client`` with ``read_dlo`` / ``write_to_dlo`` (which read and write
+a bounded snapshot), it uses a :class:`StreamingClient` and its streaming delta
+methods:
+
+* ``client.read_dlo_deltas()`` returns a *streaming* DataFrame over the
+  Change Data Feed of the source DLO. Each row carries the source columns plus
+  change-feed metadata columns (``_record_type``, ``_commit_*``).
+* ``client.write_dlo_deltas(name, df)`` starts a streaming query that writes
+  each micro-batch to the target DLO and returns the ``StreamingQuery`` handle.
+  The runtime owns the trigger, and checkpoint location — the caller only
+  chooses the table.
+
+The transform in between is ordinary PySpark. Because the source is a change
+feed, keep the metadata columns on the DataFrame you hand to
+``write_dlo_deltas`` — the sink relies on them to merge changes correctly.
+
+This entrypoint only runs inside the Data Cloud streaming (``DELTA_SYNC``)
+runtime; the local ``datacustomcode run`` readers/writers raise
+``NotImplementedError`` for the delta methods.
+"""
+
+from pyspark.sql.functions import col, upper
+
+from datacustomcode.client import StreamingClient
+
+
+def main():
+    client = StreamingClient()
+
+    # Streaming DataFrame over the source DLO's change feed.
+    deltas = client.read_dlo_deltas()
+
+    # Ordinary PySpark transform.
+    transformed = deltas.withColumn("description__c", upper(col("description__c")))
+
+    # Start the streaming write. write_dlo_deltas returns the StreamingQuery;
+    # the trigger and checkpoint location are provided by the runtime.
+    query = client.write_dlo_deltas("Account_std_copy__dll", transformed)
+
+    # Drive the query's lifecycle. In the streaming runtime this blocks until
+    # the job is stopped by the platform.
+    query.awaitTermination()
+
+
+if __name__ == "__main__":
+    main()

datacustomcode/templates/script/payload/entrypoint.py

@@ -38,6 +38,23 @@ def main():
         >>> generated_text = client.llm_gateway_generate_text(
         ...     prompt, model_id
         ... )
+
+    You can also score rows with an Einstein Studio model. The per-row helper
+    returns the same ``{status, response, error_code, error_message}`` struct,
+    where ``response`` is the prediction payload as a JSON string.
+
+    Example:
+
+        >>> from datacustomcode.client import einstein_predict_col
+        >>> from datacustomcode.einstein_predictions.types import PredictionType
+            df_scored = df.withColumn(
+            ...     "prediction__c",
+            ...     einstein_predict_col(
+            ...         "my_regression_model",
+            ...         PredictionType.REGRESSION,
+            ...         {"beds": col("beds__c"), "baths": col("baths__c")},
+            ...     )["response"],
+            ... )
     """
 
     # Drop specific columns related to relationships