diff --git a/vortex-python/src/iter/mod.rs b/vortex-python/src/iter/mod.rs index 3b8e30cd91e..be338a90e20 100644 --- a/vortex-python/src/iter/mod.rs +++ b/vortex-python/src/iter/mod.rs @@ -27,6 +27,7 @@ use vortex::array::iter::ArrayIteratorAdapter; use vortex::array::iter::ArrayIteratorExt; use vortex::dtype::DType; use vortex::error::VortexResult; +use vortex::error::vortex_err; use vortex_arrow::ArrowSessionExt; use crate::arrays::ArrowExport; @@ -99,7 +100,8 @@ impl PyArrayIterator { /// Returns the next chunk from the iterator. fn __next__(&self, py: Python) -> PyVortexResult> { - py.detach(|| { + py.check_signals()?; + let next = py.detach(|| -> VortexResult> { Ok(self .iter .lock() @@ -107,13 +109,29 @@ impl PyArrayIterator { .and_then(|iter| iter.next()) .transpose()? .map(PyArrayRef::from)) - }) + })?; + py.check_signals()?; + Ok(next) } /// Read all chunks into a single :class:`vortex.Array`. If there are multiple chunks, /// this will be a :class:`vortex.ChunkedArray`, otherwise it will be a single array. fn read_all(&self, py: Python) -> PyVortexResult { - let array = py.detach(|| self.read_remaining())?; + let Some(mut iter) = self.iter.lock().take() else { + // Otherwise, we continue to return an empty array. + return Ok(PyArrayRef::from(Canonical::empty(&self.dtype).into_array())); + }; + let dtype = iter.dtype().clone(); + let mut chunks = Vec::new(); + loop { + py.check_signals()?; + match py.detach(|| iter.next()) { + Some(chunk) => chunks.push(chunk?), + None => break, + } + } + py.check_signals()?; + let array = ArrayIteratorAdapter::new(dtype, chunks.into_iter().map(Ok)).read_all()?; Ok(PyArrayRef::from(array)) } @@ -169,6 +187,7 @@ impl PyArrayIterator { let record_batch_reader: Box = Box::new(RecordBatchIterator::new( iter.map(move |chunk| { + Python::attach(|py| py.check_signals()).map_err(|e| vortex_err!("{}", e))?; let target = target.clone(); session().arrow().execute_arrow( chunk?, diff --git a/vortex-python/src/iter/python.rs b/vortex-python/src/iter/python.rs index 10d061ff0b8..4a1d4d20e8c 100644 --- a/vortex-python/src/iter/python.rs +++ b/vortex-python/src/iter/python.rs @@ -36,6 +36,9 @@ impl Iterator for PythonArrayIterator { fn next(&mut self) -> Option { // Check for any signals on this chunk. Python::attach(|py| { + if let Err(pyerr) = py.check_signals() { + return Some(Err(vortex_err!("{}", pyerr))); + } let mut iter = self.iter.clone_ref(py).into_bound(py); iter.next().map(|array| { array diff --git a/vortex-python/test/test_iter_signals.py b/vortex-python/test/test_iter_signals.py new file mode 100644 index 00000000000..557ef8d81c6 --- /dev/null +++ b/vortex-python/test/test_iter_signals.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Tests that ArrayIterator scan paths honour pending Python signals. + +Regression tests for https://github.com/vortex-data/vortex/issues/1459: +long scans run with the GIL released, so the Rust code must call +``py.check_signals()`` at chunk boundaries. Otherwise a Ctrl+C +(``SIGINT``) appears to hang until the whole scan finishes. + +Strategy: drive each scan entry point with a counting (and, where useful, +slow) Python iterator while a background thread delivers a real ``SIGINT`` +mid-scan. The scan must surface it promptly as ``KeyboardInterrupt`` +without consuming every chunk first. +""" + +import os +import signal +import threading +import time + +import pyarrow as pa +import pytest + +import vortex as vx + +N_CHUNKS = 10 + + +def _raise_sigint() -> None: + if hasattr(signal, "raise_signal"): + signal.raise_signal(signal.SIGINT) + else: # pragma: no cover - very old Pythons + os.kill(os.getpid(), signal.SIGINT) + + +def _counting_chunks(dtype, n_chunks: int, consumed: list[int]): + chunk = vx.array(pa.array([1, 2, 3])) + + def gen(): + for _ in range(n_chunks): + consumed.append(1) + yield chunk + + assert chunk.dtype == dtype + return gen() + + +@pytest.fixture +def int_dtype(): + return vx.array(pa.array([1, 2, 3])).dtype + + +def _install_default_sigint(): + old = signal.getsignal(signal.SIGINT) + signal.signal(signal.SIGINT, signal.default_int_handler) + return old + + +def test_next_raises_on_pending_sigint(int_dtype) -> None: + consumed: list[int] = [] + # Make the first chunk slow so a SIGINT delivered from a background + # thread arrives while the scan is in flight. + chunk = vx.array(pa.array([1, 2, 3])) + assert chunk.dtype == int_dtype + + def slow_gen(): + consumed.append(1) + time.sleep(0.2) + yield chunk + + it = vx.ArrayIterator.from_iter(int_dtype, slow_gen()) + old = _install_default_sigint() + try: + + def _interrupt() -> None: + time.sleep(0.05) + _raise_sigint() + + t = threading.Thread(target=_interrupt, daemon=True) + t.start() + start = time.monotonic() + with pytest.raises(Exception) as exc_info: # noqa: BLE001, PT011 + next(it) + message = str(exc_info.value) + assert isinstance(exc_info.value, KeyboardInterrupt) or "interrupt" in message.lower() or "keyboardinterrupt" in message.lower() + elapsed = time.monotonic() - start + assert elapsed < 1.0 + t.join(timeout=5) + finally: + signal.signal(signal.SIGINT, old) + + +def test_read_all_aborts_early_on_pending_sigint(int_dtype) -> None: + consumed: list[int] = [] + chunk = vx.array(pa.array([1, 2, 3])) + assert chunk.dtype == int_dtype + + def slow_gen(): + for _ in range(N_CHUNKS): + consumed.append(1) + time.sleep(0.05) + yield chunk + + it = vx.ArrayIterator.from_iter(int_dtype, slow_gen()) + old = _install_default_sigint() + try: + + def _interrupt() -> None: + time.sleep(0.02) + _raise_sigint() + + t = threading.Thread(target=_interrupt, daemon=True) + start = time.monotonic() + t.start() + with pytest.raises(Exception) as exc_info: # noqa: BLE001, PT011 + it.read_all() + message = str(exc_info.value).lower() + assert isinstance(exc_info.value, KeyboardInterrupt) or "interrupt" in message or "keyboardinterrupt" in message + assert len(consumed) < N_CHUNKS + t.join(timeout=5) + finally: + signal.signal(signal.SIGINT, old) + + +def test_read_all_interrupted_mid_scan(int_dtype) -> None: + """A SIGINT arriving mid-scan aborts `read_all` promptly.""" + consumed: list[int] = [] + chunk = vx.array(pa.array([1, 2, 3])) + assert chunk.dtype == int_dtype + + def slow_gen(): + for _ in range(200): + consumed.append(1) + time.sleep(0.01) + yield chunk + + it = vx.ArrayIterator.from_iter(int_dtype, slow_gen()) + old = _install_default_sigint() + try: + + def _interrupt() -> None: + time.sleep(0.1) + _raise_sigint() + + t = threading.Thread(target=_interrupt, daemon=True) + start = time.monotonic() + t.start() + with pytest.raises(Exception) as exc_info: # noqa: BLE001, PT011 + it.read_all() + message = str(exc_info.value) + assert isinstance(exc_info.value, KeyboardInterrupt) or "interrupt" in message.lower() or "keyboardinterrupt" in message.lower() + elapsed = time.monotonic() - start + assert elapsed < 1.5 + assert len(consumed) < 200 + t.join(timeout=5) + finally: + signal.signal(signal.SIGINT, old) + + +def test_to_arrow_reader_aborts_on_pending_sigint() -> None: + """The `to_arrow` record-batch reader checks signals per chunk.""" + chunk = vx.array([{"x": 1}, {"x": 2}]) + consumed: list[int] = [] + + def gen(): + for _ in range(N_CHUNKS): + consumed.append(1) + time.sleep(0.02) + yield chunk + + it = vx.ArrayIterator.from_iter(chunk.dtype, gen()) + old = _install_default_sigint() + try: + + def _interrupt() -> None: + time.sleep(0.05) + _raise_sigint() + + t = threading.Thread(target=_interrupt, daemon=True) + t.start() + reader = it.to_arrow() + + with pytest.raises(Exception) as exc_info: # noqa: BLE001, PT011 + reader.read_all() + message = str(exc_info.value).lower() + assert isinstance(exc_info.value, KeyboardInterrupt) or "interrupt" in message + assert len(consumed) < N_CHUNKS + t.join(timeout=5) + finally: + signal.signal(signal.SIGINT, old)