fix einstein predict response object

Author: atulikumwenayo

src/datacustomcode/einstein_predictions/spark_default.py

@@ -39,6 +39,9 @@
 _STATUS_SUCCESS = "SUCCESS"
 _STATUS_ERROR = "ERROR"
 
+# HTTP status considered a successful prediction call.
+_HTTP_OK = 200
+
 
 class DefaultSparkEinsteinPredictions(SparkEinsteinPredictions):
 
@@ -104,6 +107,8 @@ def einstein_predict_col(
 
         def _predict(values_row: Any) -> Dict[str, Optional[str]]:
             if values_row is None:
+                # An entirely null features struct is not the normal per-feature null
+                # case; surface it directly rather than masking it (local debuggability).
                 return {
                     "status": _STATUS_ERROR,
                     "response": None,
@@ -181,6 +186,21 @@ def _call_predictions(
     return predictions.predict(request)
 
 
+def _null_feature_name(features: Dict[str, Any]) -> Optional[str]:
+    """Return the name of the first null feature value, or ``None``."""
+    for name, value in features.items():
+        if value is None:
+            return name
+    return None
+
+
+def _null_feature_message(name: str) -> str:
+    return (
+        f"Feature '{name}' has null value. Use coalesce() or when() to handle "
+        f"nulls before calling einstein_predict."
+    )
+
+
 def _invoke_predictions(
     predictions: "EinsteinPredictions",
     model_api_name: str,
@@ -190,18 +210,42 @@ def _invoke_predictions(
 ) -> 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)
+    null_feature = _null_feature_name(features)
+    if null_feature is not None:
+        message = _null_feature_message(null_feature)
+        raise EinsteinPredictionsCallError(
+            f"Einstein Predictions call failed: {message}",
+            status=None,
+            error_code=None,
+            error_message=message,
+        )
+
+    try:
+        response = _call_predictions(
+            predictions, model_api_name, prediction_type, features, settings
+        )
+    except EinsteinPredictionsCallError:
+        raise
+    except Exception as exc:
+        # Transport/build failures: surface the real error (no masking) so local runs stay
+        # debuggable. error_code stays None since there is no HTTP status.
+        raise EinsteinPredictionsCallError(
+            f"Einstein Predictions call failed: {exc}",
+            status=None,
+            error_code=None,
+            error_message=str(exc),
+        )
+
+    if response.status_code != _HTTP_OK:
+        error_message = (
+            json.dumps(response.data) if response.data is not None else None
+        )
         raise EinsteinPredictionsCallError(
             f"Einstein Predictions call failed: "
-            f"status_code={response.status_code}, "
-            f"error_code={error_code!r}, message={response.data!r}",
+            f"status_code={response.status_code}, message={error_message!r}",
             status=response.status_code,
-            error_code=error_code,
-            error_message=str(response.data) if response.data else None,
+            error_code=str(response.status_code),
+            error_message=error_message,
         )
     return response.data or {}
 
@@ -213,27 +257,46 @@ def _invoke_predictions_as_struct(
     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:
+    # (a) Customer-actionable data condition — surface the actionable message directly.
+    null_feature = _null_feature_name(features)
+    if null_feature is not None:
         return {
             "status": _STATUS_ERROR,
             "response": None,
-            "error_code": _extract_error_code(response),
-            "error_message": str(response.data) if response.data else None,
+            "error_code": None,
+            "error_message": _null_feature_message(null_feature),
         }
-    return {
-        "status": _STATUS_SUCCESS,
-        "response": json.dumps(response.data) if response.data is not None else None,
-        "error_code": None,
-        "error_message": None,
-    }
 
+    # (b) Transport/build failures — surface the real error (no masking) so local runs stay
+    # debuggable. error_code stays None since there is no HTTP status.
+    try:
+        response = _call_predictions(
+            predictions, model_api_name, prediction_type, features, settings
+        )
+    except Exception as exc:
+        return {
+            "status": _STATUS_ERROR,
+            "response": None,
+            "error_code": None,
+            "error_message": str(exc),
+        }
 
-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
+    if response.status_code == _HTTP_OK:
+        return {
+            "status": _STATUS_SUCCESS,
+            "response": (
+                json.dumps(response.data) if response.data is not None else None
+            ),
+            "error_code": None,
+            "error_message": None,
+        }
+
+    # (c) Non-200 SFAP HTTP error: error_code = status code, error_message = data JSON.
+    return {
+        "status": _STATUS_ERROR,
+        "response": None,
+        "error_code": str(response.status_code),
+        "error_message": (
+            json.dumps(response.data) if response.data is not None else None
+        ),
+    }

tests/test_spark_einstein_predictions.py

@@ -206,14 +206,16 @@ def test_udf_returns_error_struct_for_null_row(self, mock_struct, mock_udf):
         out = udf_fn(None)
         assert out["status"] == _STATUS_ERROR
         assert out["response"] is None
+        assert out["error_code"] is None
         assert "null" in out["error_message"].lower()
         mock_inner.predict.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 errors are returned as ``status="ERROR"`` structs so one bad
-        row does not abort the Spark job."""
+        row does not abort the Spark job. ``error_code`` is the HTTP status code
+        and ``error_message`` is the SFAP body as JSON."""
         mock_struct.return_value = MagicMock()
         mock_udf.return_value = MagicMock()
         mock_inner = MagicMock()
@@ -233,8 +235,99 @@ def test_udf_returns_error_struct_on_http_error(self, mock_struct, mock_udf):
 
         assert out["status"] == _STATUS_ERROR
         assert out["response"] is None
-        assert out["error_code"] == "UNAVAILABLE"
-        assert out["error_message"] is not None
+        assert out["error_code"] == "503"
+        assert json.loads(out["error_message"]) == {"errorCode": "UNAVAILABLE"}
+
+    @patch("pyspark.sql.functions.udf")
+    @patch("pyspark.sql.functions.struct")
+    def test_udf_returns_specific_error_for_null_feature(self, mock_struct, mock_udf):
+        """A null feature value is a customer-actionable data condition: it is
+        surfaced with error_code None and an actionable message, never coerced
+        to the string "None"."""
+        mock_struct.return_value = MagicMock()
+        mock_udf.return_value = MagicMock()
+        mock_inner = MagicMock()
+        predictions = DefaultSparkEinsteinPredictions(einstein_predictions=mock_inner)
+
+        predictions.einstein_predict_col(
+            "model1", PredictionType.REGRESSION, {"beds": MagicMock()}
+        )
+
+        udf_fn = mock_udf.call_args.args[0]
+        row = MagicMock()
+        row.asDict.return_value = {"beds": None}
+        out = udf_fn(row)
+
+        assert out["status"] == _STATUS_ERROR
+        assert out["response"] is None
+        assert out["error_code"] is None
+        assert "beds" in out["error_message"]
+        assert "coalesce" in out["error_message"]
+        mock_inner.predict.assert_not_called()
+
+    @patch("pyspark.sql.functions.udf")
+    @patch("pyspark.sql.functions.struct")
+    def test_udf_returns_generic_error_on_transport_failure(
+        self, mock_struct, mock_udf
+    ):
+        """Transport/build exceptions are logged and surfaced with error_code None
+        and the exception text as error_message so local runs stay debuggable."""
+        mock_struct.return_value = MagicMock()
+        mock_udf.return_value = MagicMock()
+        mock_inner = MagicMock()
+        mock_inner.predict.side_effect = RuntimeError("connection refused to 10.0.0.1")
+        predictions = DefaultSparkEinsteinPredictions(einstein_predictions=mock_inner)
+
+        predictions.einstein_predict_col(
+            "model1", PredictionType.REGRESSION, {"beds": MagicMock()}
+        )
+
+        udf_fn = mock_udf.call_args.args[0]
+        row = MagicMock()
+        row.asDict.return_value = {"beds": 3.0}
+        out = udf_fn(row)
+
+        assert out["status"] == _STATUS_ERROR
+        assert out["response"] is None
+        assert out["error_code"] is None
+        assert out["error_message"] == "connection refused to 10.0.0.1"
+
+    @patch("pyspark.sql.functions.udf")
+    @patch("pyspark.sql.functions.struct")
+    def test_udf_passes_through_prediction_failure_as_success(
+        self, mock_struct, mock_udf
+    ):
+        """A 200 response carrying a PredictionFailure stays SUCCESS; the failure
+        is passed through in ``response`` for the script to handle."""
+        mock_struct.return_value = MagicMock()
+        mock_udf.return_value = MagicMock()
+        failure_body = {
+            "results": [
+                {
+                    "type": "PredictionFailure",
+                    "error": {
+                        "message": "no match",
+                        "predictionErrorCode": "PREDICTION_ERROR_CODE_NO_MATCH",
+                    },
+                }
+            ]
+        }
+        mock_inner = MagicMock()
+        mock_inner.predict.return_value = _success_response(failure_body)
+        predictions = DefaultSparkEinsteinPredictions(einstein_predictions=mock_inner)
+
+        predictions.einstein_predict_col(
+            "model1", PredictionType.BINARY_CLASSIFICATION, {"beds": MagicMock()}
+        )
+
+        udf_fn = mock_udf.call_args.args[0]
+        row = MagicMock()
+        row.asDict.return_value = {"beds": 3.0}
+        out = udf_fn(row)
+
+        assert out["status"] == _STATUS_SUCCESS
+        assert json.loads(out["response"]) == failure_body
+        assert out["error_code"] is None
 
 
 class TestInvokePredictions:
@@ -260,9 +353,37 @@ def test_raises_call_error_on_error_response(self):
             )
 
         assert excinfo.value.status == 503
-        assert excinfo.value.error_code == "UNAVAILABLE"
+        assert excinfo.value.error_code == "503"
+        assert excinfo.value.error_message == json.dumps({"errorCode": "UNAVAILABLE"})
         assert "503" in str(excinfo.value)
-        assert "UNAVAILABLE" in str(excinfo.value)
+
+    def test_raises_specific_error_on_null_feature(self):
+        mock_inner = MagicMock()
+
+        with pytest.raises(EinsteinPredictionsCallError) as excinfo:
+            _invoke_predictions(
+                mock_inner, "model", PredictionType.REGRESSION, {"x": None}, None
+            )
+
+        assert excinfo.value.status is None
+        assert excinfo.value.error_code is None
+        assert "x" in str(excinfo.value.error_message)
+        assert "coalesce" in str(excinfo.value.error_message)
+        mock_inner.predict.assert_not_called()
+
+    def test_raises_generic_error_on_transport_failure(self):
+        mock_inner = MagicMock()
+        mock_inner.predict.side_effect = RuntimeError("connection refused to 10.0.0.1")
+
+        with pytest.raises(EinsteinPredictionsCallError) as excinfo:
+            _invoke_predictions(
+                mock_inner, "model", PredictionType.REGRESSION, {"x": 1.0}, None
+            )
+
+        assert excinfo.value.status is None
+        assert excinfo.value.error_code is None
+        assert excinfo.value.error_message == "connection refused to 10.0.0.1"
+        assert "connection refused" in str(excinfo.value)
 
 
 class TestInvokePredictionsAsStruct:
@@ -295,8 +416,35 @@ def test_error_returns_error_struct_without_raising(self):
 
         assert out["status"] == _STATUS_ERROR
         assert out["response"] is None
-        assert out["error_code"] == "UNAVAILABLE"
-        assert out["error_message"] is not None
+        assert out["error_code"] == "503"
+        assert json.loads(out["error_message"]) == {"errorCode": "UNAVAILABLE"}
+
+    def test_null_feature_returns_specific_error_struct(self):
+        mock_inner = MagicMock()
+
+        out = _invoke_predictions_as_struct(
+            mock_inner, "model", PredictionType.REGRESSION, {"x": None}, None
+        )
+
+        assert out["status"] == _STATUS_ERROR
+        assert out["response"] is None
+        assert out["error_code"] is None
+        assert "x" in out["error_message"]
+        assert "None" != out["error_message"]
+        mock_inner.predict.assert_not_called()
+
+    def test_transport_failure_returns_generic_error_struct(self):
+        mock_inner = MagicMock()
+        mock_inner.predict.side_effect = RuntimeError("connection refused to 10.0.0.1")
+
+        out = _invoke_predictions_as_struct(
+            mock_inner, "model", PredictionType.REGRESSION, {"x": 1.0}, None
+        )
+
+        assert out["status"] == _STATUS_ERROR
+        assert out["response"] is None
+        assert out["error_code"] is None
+        assert out["error_message"] == "connection refused to 10.0.0.1"
 
 
 class TestDefaultSparkEinsteinPredictionsErrorHandling:
@@ -316,4 +464,4 @@ def test_raises_on_error_response(self):
             )
 
         assert excinfo.value.status == 429
-        assert excinfo.value.error_code == "RATE_LIMITED"
+        assert excinfo.value.error_code == "429"