test: move test_collect_interrupted to test_dataframe.py

Author: kosiew

python/tests/test_dataframe.py

@@ -17,6 +17,9 @@
 import datetime
 import os
 import re
+import threading
+import time
+import ctypes
 from typing import Any
 
 import pyarrow as pa
@@ -1914,7 +1917,7 @@ def test_fill_null_date32_column(null_df):
     dates = result.column(4).to_pylist()
     assert dates[0] == datetime.date(2000, 1, 1)  # Original value
     assert dates[1] == epoch_date  # Filled value
-    assert dates[2] == datetime.date(2022, 1, 1)  # Original value
+    assert dates[2] == datetime.date(2022, 1, 1)     # Original value
     assert dates[3] == epoch_date  # Filled value
 
     # Other date column should be unchanged
@@ -2061,3 +2064,106 @@ def test_fill_null_all_null_column(ctx):
     # Check that all nulls were filled
     result = filled_df.collect()[0]
     assert result.column(1).to_pylist() == ["filled", "filled", "filled"]
+
+
+def test_collect_interrupted():
+    """Test that a long-running query can be interrupted with Ctrl-C.
+    
+    This test simulates a Ctrl-C keyboard interrupt by raising a KeyboardInterrupt
+    exception in the main thread during a long-running query execution.
+    """
+    # Create a context and a DataFrame with a query that will run for a while
+    ctx = SessionContext()
+    
+    # Create a recursive computation that will run for some time
+    batches = []
+    for i in range(10):
+        batch = pa.RecordBatch.from_arrays(
+            [
+                pa.array(list(range(i * 1000, (i + 1) * 1000))),
+                pa.array([f"value_{j}" for j in range(i * 1000, (i + 1) * 1000)]),
+            ],
+            names=["a", "b"],
+        )
+        batches.append(batch)
+    
+    # Register tables
+    ctx.register_record_batches("t1", [batches])
+    ctx.register_record_batches("t2", [batches])
+    
+    # Create a large join operation that will take time to process
+    df = ctx.sql("""
+        WITH t1_expanded AS (
+            SELECT 
+                a, 
+                b, 
+                CAST(a AS DOUBLE) / 1.5 AS c,
+                CAST(a AS DOUBLE) * CAST(a AS DOUBLE) AS d
+            FROM t1
+            CROSS JOIN (SELECT 1 AS dummy FROM t1 LIMIT 5)
+        ),
+        t2_expanded AS (
+            SELECT 
+                a,
+                b,
+                CAST(a AS DOUBLE) * 2.5 AS e,
+                CAST(a AS DOUBLE) * CAST(a AS DOUBLE) * CAST(a AS DOUBLE) AS f
+            FROM t2
+            CROSS JOIN (SELECT 1 AS dummy FROM t2 LIMIT 5)
+        )
+        SELECT 
+            t1.a, t1.b, t1.c, t1.d, 
+            t2.a AS a2, t2.b AS b2, t2.e, t2.f
+        FROM t1_expanded t1
+        JOIN t2_expanded t2 ON t1.a % 100 = t2.a % 100
+        WHERE t1.a > 100 AND t2.a > 100
+    """)
+    
+    # Flag to track if the query was interrupted
+    interrupted = False
+    interrupt_error = None
+    main_thread = threading.main_thread()
+    
+    # This function will be run in a separate thread and will raise 
+    # KeyboardInterrupt in the main thread
+    def trigger_interrupt():
+        """Wait a moment then raise KeyboardInterrupt in the main thread"""
+        time.sleep(0.5)  # Give the query time to start
+        
+        # Check if thread ID is available
+        thread_id = main_thread.ident
+        if thread_id is None:
+            msg = "Cannot get main thread ID"
+            raise RuntimeError(msg)
+            
+        # Use ctypes to raise exception in main thread
+        exception = ctypes.py_object(KeyboardInterrupt)
+        res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
+            ctypes.c_long(thread_id), exception)
+        if res != 1:
+            # If res is 0, the thread ID was invalid
+            # If res > 1, we modified multiple threads
+            ctypes.pythonapi.PyThreadState_SetAsyncExc(
+                ctypes.c_long(thread_id), ctypes.py_object(0))
+            msg = "Failed to raise KeyboardInterrupt in main thread"
+            raise RuntimeError(msg)
+    
+    # Start a thread to trigger the interrupt
+    interrupt_thread = threading.Thread(target=trigger_interrupt)
+    interrupt_thread.daemon = True
+    interrupt_thread.start()
+    
+    # Execute the query and expect it to be interrupted
+    try:
+        df.collect()
+    except KeyboardInterrupt:
+        interrupted = True
+    except Exception as e:
+        interrupt_error = e
+    
+    # Assert that the query was interrupted properly
+    assert interrupted, "Query was not interrupted by KeyboardInterrupt"
+    assert interrupt_error is None, f"Unexpected error occurred: {interrupt_error}"
+    
+    # Make sure the interrupt thread has finished
+    interrupt_thread.join(timeout=1.0)