do not fail a spark job if a single row fails on UDF

Author: atulikumwenayo

README.md

@@ -316,13 +316,16 @@ from datacustomcode.client import Client, llm_gateway_generate_text_col
 def main():
   client = Client()
   df = client.read_dlo("Input__dll")
+  # llm_gateway_generate_text_col returns a struct
+  # {status, response, error_code, error_message} per row, so per-row
+  # failures don't abort the Spark job. Pick the field you want with [].
   df_generated = df.withColumn(
     "greeting__c",
     llm_gateway_generate_text_col(
         "In one sentence, greet {name} from {city}.",
         {"name": col("name__c"), "city": col("homecity__c")},
         model_id="sfdc_ai__DefaultGPT4Omni", # An AI model in your org
-    ),
+    )["response"],
   )
 
   dlo_name = "Output_dll"

src/datacustomcode/client.py

@@ -62,15 +62,23 @@ def llm_gateway_generate_text_col(
 ) -> "Column":
     """Build a Spark Column that runs the LLM Gateway 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. ``llm_gateway_generate_text_col(...)["response"]``.
+    Per-row failures populate ``status`` / ``error_code`` / ``error_message``
+    so a single bad row does not abort the whole Spark job.
+
     Example:
 
-        >>> df.withColumn(
-        ...     "greeting__c",
-        ...     llm_gateway_generate_text_col(
-        ...         "In one sentence, greet {name} from {city}.",
-        ...         {"name": col("name__c"), "city": col("homecity__c")},
-        ...         model_id="sfdc_ai__DefaultGPT4Omni",
-        ...     ),
+        >>> result = llm_gateway_generate_text_col(
+        ...     "In one sentence, greet {name} from {city}.",
+        ...     {"name": col("name__c"), "city": col("homecity__c")},
+        ...     model_id="sfdc_ai__DefaultGPT4Omni",
+        ... )
+        >>> df.withColumn("greeting__c", result["response"])
+        >>> # …or keep the struct around and inspect failures:
+        >>> df.withColumn("llm", result).select(
+        ...     "llm.status", "llm.response", "llm.error_message"
         ... )
 
     Args:
@@ -81,7 +89,11 @@ def llm_gateway_generate_text_col(
         model_id: LLM model id. Defaults to ``sfdc_ai__DefaultGPT4Omni``.
 
     Returns:
-        A Spark ``Column`` that, when evaluated, produces the generated text.
+        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 generated text; on failure, ``status == "ERROR"`` and the
+        ``error_*`` fields carry diagnostic detail.
     """
     gateway = Client()._get_spark_llm_gateway()
     return gateway.llm_gateway_generate_text_col(template, values, model_id=model_id)

src/datacustomcode/llm_gateway/spark_base.py

@@ -50,4 +50,13 @@ def llm_gateway_generate_text_col(
         values: Union[Dict[str, "Column"], "Column"],
         model_id: Optional[str] = None,
     ) -> "Column":
-        """Build a Spark ``Column`` that invokes the LLM Gateway per row."""
+        """Build a Spark ``Column`` that invokes the LLM Gateway per row and
+        yields a struct ``{status, response, error_code, error_message}``.
+
+        Select an individual field, e.g.
+        ``llm_gateway_generate_text_col(...)["response"]``. Returning a struct
+        means a single failing row doesn't abort the Spark job.
+        Failing row leaves the rest of the DataFrame intact — callers can
+        inspect ``status`` / ``error_code`` per row instead of having the
+        Spark job abort.
+        """

src/datacustomcode/llm_gateway/spark_default.py

@@ -28,10 +28,16 @@
     from pyspark.sql import Column
 
     from datacustomcode.llm_gateway.base import LLMGateway
+    from datacustomcode.llm_gateway.types.generate_text_response import (
+        GenerateTextResponse,
+    )
 
 
 _DEFAULT_LLM_MODEL_ID = "sfdc_ai__DefaultGPT4Omni"
 
+_STATUS_SUCCESS = "SUCCESS"
+_STATUS_ERROR = "ERROR"
+
 
 class DefaultSparkLLMGateway(SparkLLMGateway):
 
@@ -60,29 +66,50 @@ def llm_gateway_generate_text_col(
         values: Union[Dict[str, "Column"], "Column"],
         model_id: Optional[str] = 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.
+        ``llm_gateway_generate_text_col(...)["response"]``.
+        """
         from pyspark.sql.functions import struct, udf
-        from pyspark.sql.types import StringType
+        from pyspark.sql.types import (
+            StringType,
+            StructField,
+            StructType,
+        )
 
         if isinstance(values, dict):
             values_col = struct(*[v.alias(k) for k, v in values.items()])
         else:
             values_col = values
 
         gateway = self._llm_gateway
+        result_schema = StructType(
+            [
+                StructField("status", StringType(), True),
+                StructField("response", StringType(), True),
+                StructField("error_code", StringType(), True),
+                StructField("error_message", StringType(), True),
+            ]
+        )
 
-        def _generate(values_row: Any) -> str:
+        def _generate(values_row: Any) -> Dict[str, Optional[str]]:
             if values_row is None:
-                return ""
+                return {
+                    "status": _STATUS_ERROR,
+                    "response": None,
+                    "error_code": None,
+                    "error_message": "values column was null for this row",
+                }
             subs = (
                 values_row.asDict()
                 if hasattr(values_row, "asDict")
                 else dict(values_row)
             )
             prompt = template.format(**subs)
-            return _invoke_llm_gateway(gateway, prompt, model_id)
+            return _invoke_llm_gateway_as_struct(gateway, prompt, model_id)
 
-        return udf(_generate, StringType())(values_col)
+        return udf(_generate, result_schema)(values_col)
 
 
 def _build_underlying_gateway() -> "LLMGateway":
@@ -97,22 +124,33 @@ def _build_underlying_gateway() -> "LLMGateway":
     return cfg.to_object()
 
 
-def _invoke_llm_gateway(
+def _call_llm_gateway(
     gateway: "LLMGateway",
     prompt: str,
     model_id: Optional[str],
-) -> str:
-    from datacustomcode.llm_gateway.errors import LLMGatewayCallError
+) -> "GenerateTextResponse":
+    """Build the request and dispatch it to the underlying gateway."""
     from datacustomcode.llm_gateway.types.generate_text_request_builder import (
         GenerateTextRequestBuilder,
     )
 
-    builder = (
+    request = (
         GenerateTextRequestBuilder()
         .set_prompt(prompt)
         .set_model(model_id or _DEFAULT_LLM_MODEL_ID)
+        .build()
     )
-    response = gateway.generate_text(builder.build())
+    return gateway.generate_text(request)
+
+
+def _invoke_llm_gateway(
+    gateway: "LLMGateway",
+    prompt: str,
+    model_id: Optional[str],
+) -> str:
+    from datacustomcode.llm_gateway.errors import LLMGatewayCallError
+
+    response = _call_llm_gateway(gateway, prompt, model_id)
     if response.is_error:
         raise LLMGatewayCallError(
             f"LLM Gateway call failed: status_code={response.status_code}, "
@@ -123,3 +161,24 @@ def _invoke_llm_gateway(
             error_message=str(response.data) if response.data else None,
         )
     return response.text
+
+
+def _invoke_llm_gateway_as_struct(
+    gateway: "LLMGateway",
+    prompt: str,
+    model_id: Optional[str],
+) -> Dict[str, Optional[str]]:
+    response = _call_llm_gateway(gateway, prompt, model_id)
+    if response.is_error:
+        return {
+            "status": _STATUS_ERROR,
+            "response": None,
+            "error_code": response.error_code or None,
+            "error_message": str(response.data) if response.data else None,
+        }
+    return {
+        "status": _STATUS_SUCCESS,
+        "response": response.text,
+        "error_code": None,
+        "error_message": None,
+    }

src/datacustomcode/templates/script/payload/entrypoint.py

@@ -16,7 +16,9 @@ def main():
     You can use your AI models configured in Salesforce to generate column
     values. See README.md for how to test locally before deploying to Data Cloud.
 
-    Example:
+    Example (the per-row helper returns a struct
+    ``{status, response, error_code, error_message}`` — pick the field you
+    want with ``[...]``):
 
         >>> from datacustomcode.client import llm_gateway_generate_text_col
             df_generated = df.withColumn(
@@ -25,7 +27,7 @@ def main():
             ...         "In one sentence, greet {name} from {city}.",
             ...         {"name": col("name__c"), "city": col("homecity__c")},
             ...         model_id="sfdc_ai__DefaultGPT4Omni",
-            ...     ),
+            ...     )["response"],
             ... )
 
     You can also invoke the LLM with a literal plain text prompt — no

tests/test_spark_llm_gateway.py

@@ -6,8 +6,11 @@
 
 from datacustomcode.llm_gateway import DefaultSparkLLMGateway, LLMGatewayCallError
 from datacustomcode.llm_gateway.spark_default import (
+    _STATUS_ERROR,
+    _STATUS_SUCCESS,
     _build_underlying_gateway,
     _invoke_llm_gateway,
+    _invoke_llm_gateway_as_struct,
 )
 from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
 
@@ -134,7 +137,12 @@ def test_dict_values_built_into_struct_and_wrapped_in_udf(
         row.asDict.return_value = {"name": "Ada", "city": "London"}
         out = udf_fn(row)
 
-        assert out == "row-out"
+        assert out == {
+            "status": _STATUS_SUCCESS,
+            "response": "row-out",
+            "error_code": None,
+            "error_message": None,
+        }
         sent = mock_inner.generate_text.call_args.args[0]
         assert sent.prompt == "Greet Ada from London."
         assert sent.model_name == "test-model"
@@ -158,7 +166,7 @@ def test_column_values_passed_through_without_struct(self, mock_struct, mock_udf
 
     @patch("pyspark.sql.functions.udf")
     @patch("pyspark.sql.functions.struct")
-    def test_udf_returns_empty_for_null_row(self, mock_struct, mock_udf):
+    def test_udf_returns_error_struct_for_null_row(self, mock_struct, mock_udf):
         mock_struct.return_value = MagicMock()
         mock_udf.return_value = MagicMock()
         mock_inner = MagicMock()
@@ -167,9 +175,37 @@ def test_udf_returns_empty_for_null_row(self, mock_struct, mock_udf):
         gateway.llm_gateway_generate_text_col("template", {"placeholder": MagicMock()})
 
         udf_fn = mock_udf.call_args.args[0]
-        assert udf_fn(None) == ""
+        out = udf_fn(None)
+        assert out["status"] == _STATUS_ERROR
+        assert out["response"] is None
+        assert "null" in out["error_message"].lower()
         mock_inner.generate_text.assert_not_called()
 
+    @patch("pyspark.sql.functions.udf")
+    @patch("pyspark.sql.functions.struct")
+    def test_udf_returns_error_struct_on_http_error(self, mock_struct, mock_udf):
+        """Per-row HTTP errors are returned as ``status="ERROR"`` structs so
+        one bad row does not abort the Spark job."""
+        mock_struct.return_value = MagicMock()
+        mock_udf.return_value = MagicMock()
+        mock_inner = MagicMock()
+        mock_inner.generate_text.return_value = _error_response(
+            status_code=503, error_code="UNAVAILABLE"
+        )
+        gateway = DefaultSparkLLMGateway(llm_gateway=mock_inner)
+
+        gateway.llm_gateway_generate_text_col("Greet {name}", {"name": MagicMock()})
+
+        udf_fn = mock_udf.call_args.args[0]
+        row = MagicMock()
+        row.asDict.return_value = {"name": "Ada"}
+        out = udf_fn(row)
+
+        assert out["status"] == _STATUS_ERROR
+        assert out["response"] is None
+        assert out["error_code"] == "UNAVAILABLE"
+        assert out["error_message"] is not None
+
 
 class TestInvokeLLMGateway:
 
@@ -207,6 +243,46 @@ def test_raises_llm_gateway_call_error_on_error_response(self):
         assert "UNAVAILABLE" in str(excinfo.value)
 
 
+class TestInvokeLLMGatewayAsStruct:
+    """Non-raising variant of ``_invoke_llm_gateway`` used by the per-row UDF.
+    Both SUCCESS and ERROR cases land in the same struct shape so callers can
+    select fields uniformly."""
+
+    def test_success_returns_success_struct(self):
+        mock_inner = MagicMock()
+        mock_inner.generate_text.return_value = _success_response("howdy")
+
+        out = _invoke_llm_gateway_as_struct(mock_inner, "prompt", "model")
+
+        assert out == {
+            "status": _STATUS_SUCCESS,
+            "response": "howdy",
+            "error_code": None,
+            "error_message": None,
+        }
+
+    def test_error_returns_error_struct_without_raising(self):
+        mock_inner = MagicMock()
+        mock_inner.generate_text.return_value = _error_response(
+            status_code=503, error_code="UNAVAILABLE"
+        )
+
+        out = _invoke_llm_gateway_as_struct(mock_inner, "prompt", "model")
+
+        assert out["status"] == _STATUS_ERROR
+        assert out["response"] is None
+        assert out["error_code"] == "UNAVAILABLE"
+        assert out["error_message"] is not None
+
+    def test_uses_default_model_when_none(self):
+        mock_inner = MagicMock()
+        mock_inner.generate_text.return_value = _success_response("ok")
+
+        _invoke_llm_gateway_as_struct(mock_inner, "prompt", None)
+        sent = mock_inner.generate_text.call_args.args[0]
+        assert sent.model_name == "sfdc_ai__DefaultGPT4Omni"
+
+
 class TestDefaultSparkLLMGatewayGenerateTextErrorHandling:
     """The scalar generate_text path raises when the underlying gateway errors."""