From acf0817047d44469736a368dc28baf07418fdcd3 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sun, 8 Nov 2026 12:26:11 +0900 Subject: [PATCH] [SPARK-61157][CONNECT][PYTHON] Capture the call sites of the origins of Column operations without importing a missing IPython module on every call Each Column operation captured its call site with `_capture_call_site` of pyspark.errors.utils, which runs `import IPython` and `import ipykernel` on every call. Where ipykernel is not installed, each failing import scans sys.path again, so 90k Column operations took 7.3 s against 0.32 s with debugging off. The call site is now computed in origin.rs with the same result: the frames outside the pyspark package, and within an IPython shell without the frames of IPython and ipykernel, as "line N in cell [K]". IPython and ipykernel are read from sys.modules, and imported at most once per process. The vendored module is unchanged. Co-authored-by: Isaac --- crates/pyspark-rs/src/origin.rs | 104 +++++++++++++++++++++++++++- python/tests/test_dropin_offline.py | 28 ++++++++ 2 files changed, 129 insertions(+), 3 deletions(-) diff --git a/crates/pyspark-rs/src/origin.rs b/crates/pyspark-rs/src/origin.rs index c73b78af..786ae809 100644 --- a/crates/pyspark-rs/src/origin.rs +++ b/crates/pyspark-rs/src/origin.rs @@ -1,18 +1,28 @@ //! As PySpark's `_with_origin`: the origin in user code of the expressions that an API function //! builds, which the server reports in the `DataFrameQueryContext` of an error in them. +use std::sync::atomic::{AtomicBool, Ordering}; + use pyo3::prelude::*; use pyo3::sync::PyOnceLock; +use pyo3::types::PyDict; use spark_connect::origin::{set_current, Origin, OriginGuard}; static UTILS: PyOnceLock> = PyOnceLock::new(); +static SYS: PyOnceLock> = PyOnceLock::new(); +/// The directory of the `pyspark` package, as `os.path.dirname(pyspark.__file__)`. +static PYSPARK_ROOT: PyOnceLock = PyOnceLock::new(); +/// Whether the imports of IPython and ipykernel were tried, once per process. A failed import +/// is not retried, and a module imported later is still seen in `sys.modules`. +static IPYTHON_IMPORTS_TRIED: AtomicBool = AtomicBool::new(false); /// While the returned guard lives, the expressions that this thread builds have the origin of /// `fragment`, the name of the API function, at its call site in user code, when DataFrame /// debugging is enabled. As PySpark's `_with_origin`, with the functions of /// `pyspark.errors.utils` that it calls: `is_debugging_enabled`, which reads -/// `spark.python.sql.dataFrameDebugging.enabled` once, and `_capture_call_site`. An origin that -/// cannot be read is not set, as an expression is the same without one. +/// `spark.python.sql.dataFrameDebugging.enabled` once, and `_capture_call_site`, which +/// `call_site` gives without importing a module per call. An origin that cannot be read is +/// not set, as an expression is the same without one. pub(crate) fn with_origin(fragment: &str) -> Option { Python::attach(|py| -> PyResult> { let utils = UTILS @@ -21,7 +31,7 @@ pub(crate) fn with_origin(fragment: &str) -> Option { if !utils.call_method0("is_debugging_enabled")?.is_truthy()? { return Ok(None); } - let call_site: String = utils.call_method1("_capture_call_site", (1,))?.extract()?; + let call_site = call_site(py, 1)?; Ok(Some(set_current(Origin { fragment: fragment.to_string(), call_site, @@ -29,3 +39,91 @@ pub(crate) fn with_origin(fragment: &str) -> Option { }) .unwrap_or(None) } + +/// The call site of the first `depth` frames of user code, as `_capture_call_site(depth)` of +/// `pyspark.errors.utils` gives it: the frames outside the pyspark package, and within an +/// IPython shell (when IPython and ipykernel are loaded) without the frames of IPython and +/// ipykernel, each "line N in cell [K]", else "file:line", joined by newlines. +fn call_site(py: Python<'_>, depth: usize) -> PyResult { + let sys = SYS + .get_or_try_init(py, || py.import("sys").map(Bound::unbind))? + .bind(py); + let pyspark_root = PYSPARK_ROOT.get_or_try_init(py, || { + let file = py.import("pyspark")?.getattr("__file__")?; + py.import("os")? + .getattr("path")? + .call_method1("dirname", (file,))? + .extract::() + })?; + + // The frames as `inspect` walks them from the innermost: `sys._getframe(0)` is the Python + // frame that called the API function, as the native code of it adds no frame. + let mut selected = Vec::with_capacity(depth); + let mut frame = sys.call_method1("_getframe", (0,))?; + while selected.len() < depth && !frame.is_none() { + let file: String = frame.getattr("f_code")?.getattr("co_filename")?.extract()?; + if !file.contains(pyspark_root.as_str()) { + selected.push((file, frame.getattr("f_lineno")?)); + } + frame = frame.getattr("f_back")?; + } + + let Some((ipython, ipykernel)) = loaded_ipython(py, sys)? else { + return join_frames(&selected); + }; + // The frames of IPython and ipykernel are dropped after the first `depth` frames, as + // `_capture_call_site` does. + let os_path = py.import("os")?.getattr("path")?; + let ipy_root: String = os_path + .call_method1("dirname", (ipython.getattr("__file__")?,))? + .extract()?; + let kernel_root: String = os_path + .call_method1("dirname", (ipykernel.getattr("__file__")?,))? + .extract()?; + let frames: Vec<_> = selected + .into_iter() + .filter(|(file, _)| { + !file.contains(ipy_root.as_str()) && !file.contains(kernel_root.as_str()) + }) + .collect(); + let shell = ipython.call_method0("get_ipython")?; + if !shell.is_truthy()? { + return join_frames(&frames); + } + // Identifying the cell is useful when the error is generated from IPython Notebook. + let count = shell.getattr("execution_count")?; + let cells: Vec = frames + .iter() + .map(|(_, lineno)| Ok(format!("line {lineno} in cell [{count}]"))) + .collect::>()?; + Ok(cells.join("\n")) +} + +/// The "file:line" of each frame, joined by newlines. +fn join_frames(frames: &[(String, Bound<'_, PyAny>)]) -> PyResult { + let lines: Vec = frames + .iter() + .map(|(file, lineno)| Ok(format!("{file}:{lineno}"))) + .collect::>()?; + Ok(lines.join("\n")) +} + +/// The modules IPython and ipykernel, when both are loaded. They are read from `sys.modules`, +/// and imported once per process when one is missing, as `_capture_call_site` imports them. +fn loaded_ipython<'py>( + py: Python<'py>, + sys: &Bound<'py, PyModule>, +) -> PyResult, Bound<'py, PyAny>)>> { + let modules = sys.getattr("modules")?; + let modules = modules.cast::()?; + let find = |name: &str| modules.get_item(name); + let (mut ipython, mut ipykernel) = (find("IPython")?, find("ipykernel")?); + if (ipython.is_none() || ipykernel.is_none()) + && !IPYTHON_IMPORTS_TRIED.swap(true, Ordering::Relaxed) + { + // `import IPython`, then `import ipykernel`, as `_capture_call_site` does. + let _ = py.import("IPython").and_then(|_| py.import("ipykernel")); + (ipython, ipykernel) = (find("IPython")?, find("ipykernel")?); + } + Ok(ipython.zip(ipykernel)) +} diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 5417a022..ce1a391c 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -11653,6 +11653,34 @@ def unsupported(name): assert error.value.errno in (errno.ENOENT, errno.EISDIR), path +def test_column_origins_capture_their_call_site_without_importing_per_call(mock_connect, monkeypatch): + # As PySpark's `_with_origin`: a Column operation carries the call site of its user code, + # which `_capture_call_site` gives, but the capture does not import IPython and ipykernel + # on every call, where ipykernel is missing. + import builtins + + import pyspark.errors.utils as errors_utils + from pyspark.sql import functions as F + + spark, service = mock_connect + monkeypatch.setattr(errors_utils, "_enable_debugging_cache", True) + column = F.col("id") + 1; expected = errors_utils._capture_call_site(1) + imports = [] + real_import = builtins.__import__ + + def counting_import(name, *args, **kwargs): + if name in ("IPython", "ipykernel"): + imports.append(name) + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", counting_import) + for _ in range(200): + F.col("id") + 1 + assert len(imports) <= 1 + spark.range(1).select(column).collect() + assert expected.encode() in service.execute_requests[-1] + + def test_add_artifacts_puts_a_pyfile_zip_on_the_client_path_before_opening_it(tmp_path, monkeypatch): # As PySpark's _parse_artifacts: a zip, egg or jar of pyfile is put on sys.path before its # file is opened, and a Python file is not.