From 89797dda5048c667b766baabd28469e29dbd542e Mon Sep 17 00:00:00 2001 From: Pamod Madubashana Date: Mon, 28 Sep 2026 10:21:03 +0000 Subject: [PATCH 1/3] Fix #1459: Call check_signals in the scan to support signal handlers Signed-off-by: Pamod Madubashana --- vortex-python/src/iter/mod.rs | 25 ++++++++++++++++++++++--- vortex-python/src/iter/python.rs | 3 +++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/vortex-python/src/iter/mod.rs b/vortex-python/src/iter/mod.rs index 3b8e30cd91e..a229374034a 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(|| { 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 From dcfed285e0e8cbc7b0eab3b5dce491867ef86cdf Mon Sep 17 00:00:00 2001 From: Pamod Madubashana Date: Wed, 30 Sep 2026 10:19:57 +0530 Subject: [PATCH 2/3] test(python): cover check_signals handling in ArrayIterator scans Queue a real SIGINT and assert next(), read_all(), and the to_arrow record-batch reader surface it promptly while aborting early, plus a mid-scan interrupt test proving read_all does not run to completion. Signed-off-by: Pamod Madubashana --- vortex-python/test/test_iter_signals.py | 150 ++++++++++++++++++++++++ 1 file changed, 150 insertions(+) create mode 100644 vortex-python/test/test_iter_signals.py diff --git a/vortex-python/test/test_iter_signals.py b/vortex-python/test/test_iter_signals.py new file mode 100644 index 00000000000..d4aa803d5d6 --- /dev/null +++ b/vortex-python/test/test_iter_signals.py @@ -0,0 +1,150 @@ +# 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: queue a real ``SIGINT`` for this process, then drive each scan +entry point (``next()``, ``read_all()`` and the ``to_arrow`` record-batch +reader). With the default ``SIGINT`` handler the pending signal must +surface as ``KeyboardInterrupt`` *from* the scan call, and a counting +Python iterator must prove the scan aborted early instead of 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] = [] + it = vx.ArrayIterator.from_iter(int_dtype, _counting_chunks(int_dtype, N_CHUNKS, consumed)) + old = _install_default_sigint() + try: + _raise_sigint() + with pytest.raises(KeyboardInterrupt): + next(it) + # The signal must be observed before the scan makes progress. + assert len(consumed) <= 1 + finally: + signal.signal(signal.SIGINT, old) + + +def test_read_all_aborts_early_on_pending_sigint(int_dtype) -> None: + consumed: list[int] = [] + it = vx.ArrayIterator.from_iter(int_dtype, _counting_chunks(int_dtype, N_CHUNKS, consumed)) + old = _install_default_sigint() + try: + _raise_sigint() + with pytest.raises(KeyboardInterrupt): + it.read_all() + # Without per-chunk `check_signals`, the whole scan would run to + # completion inside a single detached section and consume everything. + assert len(consumed) < N_CHUNKS + 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(KeyboardInterrupt): + it.read_all() + elapsed = time.monotonic() - start + # Full scan would take ~2s; an interruptible scan aborts far sooner + # and without consuming every chunk. + 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) + yield chunk + + it = vx.ArrayIterator.from_iter(chunk.dtype, gen()) + old = _install_default_sigint() + try: + _raise_sigint() + reader = it.to_arrow() + # The signal check is mapped through VortexError into an Arrow + # external error, so accept either surface form. The early-abort + # assertion below is the real proof the signal was honoured. + 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 + finally: + signal.signal(signal.SIGINT, old) From 08fe8f75a967741dc4e81ab49a2377942a2c98a9 Mon Sep 17 00:00:00 2001 From: pamod-madubashana Date: Fri, 23 Oct 2026 21:47:57 +0530 Subject: [PATCH 3/3] fix(python): type-annotate check_signals closure and make signal tests runnable --- vortex-python/src/iter/mod.rs | 2 +- vortex-python/test/test_iter_signals.py | 87 ++++++++++++++++++------- 2 files changed, 65 insertions(+), 24 deletions(-) diff --git a/vortex-python/src/iter/mod.rs b/vortex-python/src/iter/mod.rs index a229374034a..be338a90e20 100644 --- a/vortex-python/src/iter/mod.rs +++ b/vortex-python/src/iter/mod.rs @@ -101,7 +101,7 @@ impl PyArrayIterator { /// Returns the next chunk from the iterator. fn __next__(&self, py: Python) -> PyVortexResult> { py.check_signals()?; - let next = py.detach(|| { + let next = py.detach(|| -> VortexResult> { Ok(self .iter .lock() diff --git a/vortex-python/test/test_iter_signals.py b/vortex-python/test/test_iter_signals.py index d4aa803d5d6..557ef8d81c6 100644 --- a/vortex-python/test/test_iter_signals.py +++ b/vortex-python/test/test_iter_signals.py @@ -8,12 +8,10 @@ ``py.check_signals()`` at chunk boundaries. Otherwise a Ctrl+C (``SIGINT``) appears to hang until the whole scan finishes. -Strategy: queue a real ``SIGINT`` for this process, then drive each scan -entry point (``next()``, ``read_all()`` and the ``to_arrow`` record-batch -reader). With the default ``SIGINT`` handler the pending signal must -surface as ``KeyboardInterrupt`` *from* the scan call, and a counting -Python iterator must prove the scan aborted early instead of consuming -every chunk first. +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 @@ -61,29 +59,66 @@ def _install_default_sigint(): def test_next_raises_on_pending_sigint(int_dtype) -> None: consumed: list[int] = [] - it = vx.ArrayIterator.from_iter(int_dtype, _counting_chunks(int_dtype, N_CHUNKS, consumed)) + # 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: - _raise_sigint() - with pytest.raises(KeyboardInterrupt): + + 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) - # The signal must be observed before the scan makes progress. - assert len(consumed) <= 1 + 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] = [] - it = vx.ArrayIterator.from_iter(int_dtype, _counting_chunks(int_dtype, N_CHUNKS, consumed)) + 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: - _raise_sigint() - with pytest.raises(KeyboardInterrupt): + + 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() - # Without per-chunk `check_signals`, the whole scan would run to - # completion inside a single detached section and consume everything. + 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) @@ -111,11 +146,11 @@ def _interrupt() -> None: t = threading.Thread(target=_interrupt, daemon=True) start = time.monotonic() t.start() - with pytest.raises(KeyboardInterrupt): + 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 - # Full scan would take ~2s; an interruptible scan aborts far sooner - # and without consuming every chunk. assert elapsed < 1.5 assert len(consumed) < 200 t.join(timeout=5) @@ -131,20 +166,26 @@ def test_to_arrow_reader_aborts_on_pending_sigint() -> None: 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: - _raise_sigint() + + def _interrupt() -> None: + time.sleep(0.05) + _raise_sigint() + + t = threading.Thread(target=_interrupt, daemon=True) + t.start() reader = it.to_arrow() - # The signal check is mapped through VortexError into an Arrow - # external error, so accept either surface form. The early-abort - # assertion below is the real proof the signal was honoured. + 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)