diff --git a/crates/sail-common-datafusion/src/error.rs b/crates/sail-common-datafusion/src/error.rs index a9c1a81740..f9df981a2a 100644 --- a/crates/sail-common-datafusion/src/error.rs +++ b/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>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub failure_kind: Option, } /// A trait to extract Python error cause from a generic error. @@ -166,6 +193,14 @@ impl CommonErrorCause { }; } + if let Some(failure) = error.downcast_ref::() { + 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::(error, &mut HashSet::new()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_python_failure_kind_round_trips() -> Result<(), Box> { + 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> { + 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(()) + } +} diff --git a/crates/sail-data-source/src/formats/python/error.rs b/crates/sail-data-source/src/formats/python/error.rs index ecff918436..4cc9b0d234 100644 --- a/crates/sail-data-source/src/formats/python/error.rs +++ b/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 { + use pyo3::prelude::PyAnyMethods; + use pyo3::types::{PyTuple, PyTupleMethods, PyType}; + + pyo3::Python::attach(|py| { + let type_type = py.get_type::(); + let getattribute = type_type.getattr("__getattribute__").ok()?; + let exception_type = error.get_type(py); + let mro = getattribute + .call1((&exception_type, "__mro__")) + .ok()? + .cast_into::() + .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::() 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 = Result; @@ -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 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 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", + ) + } +} diff --git a/crates/sail-python-udf/src/error.rs b/crates/sail-python-udf/src/error.rs index 8bbf3a049a..8e6e6b2c33 100644 --- a/crates/sail-python-udf/src/error.rs +++ b/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")); + } +} diff --git a/crates/sail-spark-connect/src/error.rs b/crates/sail-spark-connect/src/error.rs index 450a8ea86a..cab5547e36 100644 --- a/crates/sail-spark-connect/src/error.rs +++ b/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 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) diff --git a/docs/guide/sources/python/index.md b/docs/guide/sources/python/index.md index 23d344f82c..6fa8fe9380 100644 --- a/docs/guide/sources/python/index.md +++ b/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 diff --git a/python/pysail/tests/spark/datasource/test_python.py b/python/pysail/tests/spark/datasource/test_python.py index ead765716e..5a152a4072 100644 --- a/python/pysail/tests/spark/datasource/test_python.py +++ b/python/pysail/tests/spark/datasource/test_python.py @@ -5,6 +5,7 @@ including both Arrow RecordBatch and tuple-based paths. """ +import contextlib import json from collections.abc import Iterator from pathlib import Path @@ -580,6 +581,81 @@ def read(self, partition): # noqa: ARG002 df.collect() +@pytest.mark.parametrize( + ("failure_kind", "expected_type", "expected_message"), + [ + ("terminal", "AnalysisException", "Python data source reported a terminal failure"), + ("transient", "SparkRuntimeException", "Python data source reported a transient failure"), + ], +) +def test_python_declared_failure_kind_is_structured_and_message_free( + failure_kind: str, expected_type: str, expected_message: str +): + """A declared failure crosses Spark Connect by finite class, not Python text.""" + import pyarrow as pa + from pyspark import errors + from pyspark.sql import SparkSession + from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition + + import pysail.spark + + class DeclaredDataSourceError(Exception): + __sail_data_source_failure_kind__ = failure_kind + + def __getattribute__(self, name: str): + if name == "__sail_data_source_failure_kind__": + return "masked-on-instance" + return super().__getattribute__(name) + + private_detail = "detail that must not cross the boundary" + + def direct_read(_reader, _partition): + raise DeclaredDataSourceError(private_detail) + + def generator_read(_reader, _partition): + raise DeclaredDataSourceError(private_detail) + yield (0,) # pragma: no cover - makes this a generator like streaming readers + + class DeclaredFailureReader(DataSourceReader): + read = direct_read if failure_kind == "terminal" else generator_read + + def partitions(self): + return [InputPartition(0)] + + class DeclaredFailureDataSource(DataSource): + @classmethod + def name(cls) -> str: + return f"declared_failure_{failure_kind}" + + def schema(self): + return pa.schema([("id", pa.int32())]) + + def reader(self, schema): # noqa: ARG002 + return DeclaredFailureReader() + + server = pysail.spark.SparkConnectServer("127.0.0.1", 0) + server.start() + host, port = server.listening_address + spark = SparkSession.builder.remote(f"sc://{host}:{port}").create() + try: + spark.conf.set("spark.sql.session.localRelationSizeLimit", "3g") + spark.dataSource.register(DeclaredFailureDataSource) + frame = spark.read.format(DeclaredFailureDataSource.name()).load() + exception_type = getattr(errors, expected_type) + + with pytest.raises(exception_type) as caught: + frame.collect() + assert expected_message in str(caught.value) + assert private_detail not in str(caught.value) + assert caught.value.__cause__ is None + assert private_detail not in repr(caught.value.__context__) + finally: + with contextlib.suppress(Exception): + spark.stop() + with contextlib.suppress(Exception): + server.stop() + + def test_python_session_isolation(remote: str): """Test that datasources registered in one session are not visible in another.