feat: preserve typed Python data source failures

crates/sail-common-datafusion/src/error.rs

@@ -14,10 +14,37 @@ pub struct RemoteError {
     pub cause: CommonErrorCause,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub enum PythonFailureKind {
+    Terminal,
+    Transient,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
+pub enum PythonDataSourceFailure {
+    #[error("Python data source reported a terminal failure")]
+    Terminal,
+    #[error("Python data source reported a transient failure")]
+    Transient,
+}
+
+impl PythonDataSourceFailure {
+    pub fn kind(self) -> PythonFailureKind {
+        match self {
+            Self::Terminal => PythonFailureKind::Terminal,
+            Self::Transient => PythonFailureKind::Transient,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
 pub struct PythonErrorCause {
     pub summary: String,
     pub traceback: Option<Vec<String>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub failure_kind: Option<PythonFailureKind>,
 }
 
 /// A trait to extract Python error cause from a generic error.
@@ -166,6 +193,14 @@ impl CommonErrorCause {
             };
         }
 
+        if let Some(failure) = error.downcast_ref::<PythonDataSourceFailure>() {
+            return Self::Python(PythonErrorCause {
+                summary: failure.to_string(),
+                traceback: None,
+                failure_kind: Some(failure.kind()),
+            });
+        }
+
         if let Some(cause) = Py::extract(error) {
             return Self::Python(cause);
         }
@@ -187,3 +222,37 @@ impl CommonErrorCause {
         Self::build::<Py>(error, &mut HashSet::new())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_python_failure_kind_round_trips() -> Result<(), Box<dyn std::error::Error>> {
+        let cause = CommonErrorCause::Python(PythonErrorCause {
+            summary: "Python data source reported a terminal failure".to_string(),
+            traceback: None,
+            failure_kind: Some(PythonFailureKind::Terminal),
+        });
+
+        let encoded = serde_json::to_string(&cause)?;
+        assert!(encoded.contains(r#""failureKind":"terminal""#));
+        let decoded: CommonErrorCause = serde_json::from_str(&encoded)?;
+        let CommonErrorCause::Python(decoded) = decoded else {
+            return Err(std::io::Error::other("expected Python error cause").into());
+        };
+        assert_eq!(decoded.failure_kind, Some(PythonFailureKind::Terminal));
+        Ok(())
+    }
+
+    #[test]
+    fn test_legacy_python_cause_defaults_failure_kind() -> Result<(), Box<dyn std::error::Error>> {
+        let encoded = r#"{"python":{"summary":"legacy","traceback":null}}"#;
+        let decoded: CommonErrorCause = serde_json::from_str(encoded)?;
+        let CommonErrorCause::Python(decoded) = decoded else {
+            return Err(std::io::Error::other("expected Python error cause").into());
+        };
+        assert_eq!(decoded.failure_kind, None);
+        Ok(())
+    }
+}

crates/sail-data-source/src/formats/python/error.rs

@@ -3,8 +3,43 @@
 //! Provides structured error types with context for debugging Python datasource issues.
 
 use datafusion_common::DataFusionError;
+use sail_common_datafusion::error::PythonDataSourceFailure;
 use thiserror::Error;
 
+const FAILURE_KIND_ATTRIBUTE: &str = "__sail_data_source_failure_kind__";
+
+fn declared_failure_kind(error: &pyo3::PyErr) -> Option<PythonDataSourceFailure> {
+    use pyo3::prelude::PyAnyMethods;
+    use pyo3::types::{PyTuple, PyTupleMethods, PyType};
+
+    pyo3::Python::attach(|py| {
+        let type_type = py.get_type::<PyType>();
+        let getattribute = type_type.getattr("__getattribute__").ok()?;
+        let exception_type = error.get_type(py);
+        let mro = getattribute
+            .call1((&exception_type, "__mro__"))
+            .ok()?
+            .cast_into::<PyTuple>()
+            .ok()?;
+
+        for base in mro.iter() {
+            let namespace = getattribute.call1((&base, "__dict__")).ok()?;
+            let Ok(value) = namespace.get_item(FAILURE_KIND_ATTRIBUTE) else {
+                continue;
+            };
+            let Ok(value) = value.extract::<String>() else {
+                return None;
+            };
+            return match value.as_str() {
+                "terminal" => Some(PythonDataSourceFailure::Terminal),
+                "transient" => Some(PythonDataSourceFailure::Transient),
+                _ => None,
+            };
+        }
+        None
+    })
+}
+
 /// Result type alias for Python data source operations.
 #[expect(dead_code)]
 pub type PythonDataSourceResult<T> = Result<T, PythonDataSourceError>;
@@ -30,6 +65,9 @@ pub enum PythonDataSourceError {
     /// Resource exhaustion (e.g., partition too large)
     #[error("Resource exhausted: {0}")]
     ResourceExhausted(String),
+    /// Application-declared failure with private Python details discarded.
+    #[error("{0}")]
+    DeclaredFailure(#[from] PythonDataSourceFailure),
 }
 
 impl PythonDataSourceError {
@@ -92,7 +130,10 @@ impl PythonDataSourceContext {
 
     /// Wrap a Python error with context information, preserving traceback.
     pub fn wrap_py_error(&self, e: pyo3::PyErr) -> PythonDataSourceError {
-        self.wrap_error(format_py_error_with_traceback(e))
+        match declared_failure_kind(&e) {
+            Some(failure) => failure.into(),
+            None => self.wrap_error(format_py_error_with_traceback(e)),
+        }
     }
 }
 
@@ -125,7 +166,10 @@ pub fn format_py_error_with_traceback(e: pyo3::PyErr) -> String {
 
 impl From<pyo3::PyErr> for PythonDataSourceError {
     fn from(e: pyo3::PyErr) -> Self {
-        Self::python(format_py_error_with_traceback(e))
+        match declared_failure_kind(&e) {
+            Some(failure) => failure.into(),
+            None => Self::python(format_py_error_with_traceback(e)),
+        }
     }
 }
 
@@ -134,9 +178,12 @@ impl From<pyo3::PyErr> for PythonDataSourceError {
 /// This is a shared helper to avoid duplicating this conversion pattern
 /// across multiple modules (stream.rs, executor.rs, arrow_utils.rs, etc.).
 pub fn py_err(e: pyo3::PyErr) -> DataFusionError {
-    DataFusionError::External(Box::new(std::io::Error::other(
-        format_py_error_with_traceback(e),
-    )))
+    match declared_failure_kind(&e) {
+        Some(failure) => PythonDataSourceError::from(failure).into(),
+        None => DataFusionError::External(Box::new(std::io::Error::other(
+            format_py_error_with_traceback(e),
+        ))),
+    }
 }
 
 /// Import cloudpickle from PySpark.
@@ -154,3 +201,64 @@ pub fn import_cloudpickle(
         ))
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use pyo3::ffi::c_str;
+    use pyo3::prelude::*;
+    use pyo3::types::{PyDict, PyDictMethods};
+
+    use super::*;
+
+    #[expect(clippy::unwrap_used)]
+    fn declared_error(kind: &str) -> PyErr {
+        Python::initialize();
+        Python::attach(|py| {
+            let namespace = PyDict::new(py);
+            namespace.set_item("failure_kind", kind).unwrap();
+            py.run(
+                c_str!(
+                    "class DeclaredError(RuntimeError):\n    __sail_data_source_failure_kind__ = failure_kind\n"
+                ),
+                Some(&namespace),
+                None,
+            )
+            .unwrap();
+            let exception_type = namespace.get_item("DeclaredError").unwrap().unwrap();
+            let value = exception_type.call1(("private Python detail",)).unwrap();
+            PyErr::from_value(value)
+        })
+    }
+
+    fn assert_declared_failure_is_constant(kind: &str, message: &str) -> Result<(), String> {
+        match py_err(declared_error(kind)) {
+            DataFusionError::External(error) => {
+                assert_eq!(error.to_string(), message);
+                assert!(!error.to_string().contains("private Python detail"));
+                let Some(source) = error.source() else {
+                    return Err("classified marker source was not preserved".to_string());
+                };
+                assert_eq!(source.to_string(), message);
+                assert!(source.source().is_none());
+                Ok(())
+            }
+            other => Err(format!("expected external error, got {other:?}")),
+        }
+    }
+
+    #[test]
+    fn test_terminal_declared_failure_preserves_only_finite_marker() -> Result<(), String> {
+        assert_declared_failure_is_constant(
+            "terminal",
+            "Python data source reported a terminal failure",
+        )
+    }
+
+    #[test]
+    fn test_transient_declared_failure_preserves_only_finite_marker() -> Result<(), String> {
+        assert_declared_failure_is_constant(
+            "transient",
+            "Python data source reported a transient failure",
+        )
+    }
+}

crates/sail-python-udf/src/error.rs

@@ -59,9 +59,35 @@ impl PythonErrorCauseExtractor for PyErrExtractor {
             Some(PythonErrorCause {
                 summary: e.to_string(),
                 traceback: traceback.ok(),
+                failure_kind: None,
             })
         } else {
             None
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use pyo3::exceptions::PyRuntimeError;
+
+    use super::*;
+
+    #[test]
+    #[expect(clippy::unwrap_used)]
+    fn test_generic_python_error_ignores_data_source_attribute() {
+        Python::initialize();
+        let error = Python::attach(|py| {
+            let error = PyRuntimeError::new_err("generic Python detail");
+            error
+                .value(py)
+                .setattr("__sail_data_source_failure_kind__", "terminal")
+                .unwrap();
+            error
+        });
+
+        let cause = PyErrExtractor::extract(&error).unwrap();
+        assert_eq!(cause.failure_kind, None);
+        assert!(cause.summary.contains("generic Python detail"));
+    }
+}

crates/sail-spark-connect/src/error.rs

@@ -8,7 +8,7 @@ use datafusion::common::DataFusionError;
 use prost::{DecodeError, UnknownEnumValue};
 use sail_cache::error::CacheError;
 use sail_common::error::CommonError;
-use sail_common_datafusion::error::{CommonErrorCause, PythonErrorCause};
+use sail_common_datafusion::error::{CommonErrorCause, PythonErrorCause, PythonFailureKind};
 use sail_execution::error::ExecutionError;
 use sail_plan::error::PlanError;
 use sail_python_udf::error::PyErrExtractor;
@@ -369,21 +369,31 @@ impl From<CommonErrorCause> for SparkThrowable {
                 SparkThrowable::ArithmeticException(x)
             }
             CommonErrorCause::ArrowParse(x) => SparkThrowable::ParseException(x),
-            CommonErrorCause::Python(PythonErrorCause { summary, traceback }) => {
-                // The message must end with a newline character
-                // since the PySpark unit tests expect it.
-                let message = if let Some(traceback) = traceback {
-                    // Each line string already ends with a newline character.
-                    traceback.join("")
-                } else {
-                    format!("{summary}\n")
-                };
-                if message.contains("net.razorvine.pickle.PickleException") {
-                    SparkThrowable::SparkException(message)
-                } else {
-                    SparkThrowable::PythonException(message)
+            CommonErrorCause::Python(PythonErrorCause {
+                summary,
+                traceback,
+                failure_kind,
+            }) => match failure_kind {
+                Some(PythonFailureKind::Terminal) => SparkThrowable::AnalysisException(summary),
+                Some(PythonFailureKind::Transient) => {
+                    SparkThrowable::SparkRuntimeException(summary)
                 }
-            }
+                None => {
+                    // The message must end with a newline character
+                    // since the PySpark unit tests expect it.
+                    let message = if let Some(traceback) = traceback {
+                        // Each line string already ends with a newline character.
+                        traceback.join("")
+                    } else {
+                        format!("{summary}\n")
+                    };
+                    if message.contains("net.razorvine.pickle.PickleException") {
+                        SparkThrowable::SparkException(message)
+                    } else {
+                        SparkThrowable::PythonException(message)
+                    }
+                }
+            },
             CommonErrorCause::ArrowCast(x) => cast_error_to_throwable(x),
             CommonErrorCause::Schema(x)
             | CommonErrorCause::Plan(x)

docs/guide/sources/python/index.md

@@ -12,6 +12,34 @@ You can define a Python class that inherits from the `pyspark.sql.datasource.Dat
 
 Currently, Sail supports Python data sources for batch reading and writing.
 
+## Classified Reader Failures
+
+By default, Sail propagates a Python data source exception with its Python
+traceback. A data source can instead declare a finite failure category when its
+caller must distinguish a deterministic failure from a retryable one without
+inspecting exception text. Define `__sail_data_source_failure_kind__` on the
+exception class with one of these exact values:
+
+- `"terminal"` becomes a Spark `AnalysisException`.
+- `"transient"` becomes a Spark `SparkRuntimeException`.
+
+For a declared failure, Sail replaces the Python message and traceback with a
+constant category message before the error crosses Spark Connect. Unknown
+values retain the default exception behavior.
+
+```python
+class RetryableReadError(TimeoutError):
+    __sail_data_source_failure_kind__ = "transient"
+
+
+class ContractReadError(ValueError):
+    __sail_data_source_failure_kind__ = "terminal"
+```
+
+Use this protocol only for data-source-controlled exception classes. Do not put
+record values, credentials, endpoints, or other runtime details in the category
+attribute.
+
 ## Examples
 
 <!--@include: ../../_common/spark-session.md-->