Skip to content

Commit e2f1e1c

Browse files
committed
fix ruff errors
1 parent 84f31bb commit e2f1e1c

File tree

1 file changed

+27
-26
lines changed

1 file changed

+27
-26
lines changed

python/tests/test_dataframe.py

Lines changed: 27 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,12 @@
1414
# KIND, either express or implied. See the License for the
1515
# specific language governing permissions and limitations
1616
# under the License.
17+
import ctypes
1718
import datetime
1819
import os
1920
import re
2021
import threading
2122
import time
22-
import ctypes
2323
from typing import Any
2424

2525
import pyarrow as pa
@@ -1299,7 +1299,6 @@ def test_collect_partitioned():
12991299
assert [[batch]] == ctx.create_dataframe([[batch]]).collect_partitioned()
13001300

13011301

1302-
13031302
def test_union(ctx):
13041303
batch = pa.RecordBatch.from_arrays(
13051304
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
@@ -1917,7 +1916,7 @@ def test_fill_null_date32_column(null_df):
19171916
dates = result.column(4).to_pylist()
19181917
assert dates[0] == datetime.date(2000, 1, 1) # Original value
19191918
assert dates[1] == epoch_date # Filled value
1920-
assert dates[2] == datetime.date(2022, 1, 1) # Original value
1919+
assert dates[2] == datetime.date(2022, 1, 1) # Original value
19211920
assert dates[3] == epoch_date # Filled value
19221921

19231922
# Other date column should be unchanged
@@ -2068,13 +2067,13 @@ def test_fill_null_all_null_column(ctx):
20682067

20692068
def test_collect_interrupted():
20702069
"""Test that a long-running query can be interrupted with Ctrl-C.
2071-
2070+
20722071
This test simulates a Ctrl-C keyboard interrupt by raising a KeyboardInterrupt
20732072
exception in the main thread during a long-running query execution.
20742073
"""
20752074
# Create a context and a DataFrame with a query that will run for a while
20762075
ctx = SessionContext()
2077-
2076+
20782077
# Create a recursive computation that will run for some time
20792078
batches = []
20802079
for i in range(10):
@@ -2086,49 +2085,49 @@ def test_collect_interrupted():
20862085
names=["a", "b"],
20872086
)
20882087
batches.append(batch)
2089-
2088+
20902089
# Register tables
20912090
ctx.register_record_batches("t1", [batches])
20922091
ctx.register_record_batches("t2", [batches])
2093-
2092+
20942093
# Create a large join operation that will take time to process
20952094
df = ctx.sql("""
20962095
WITH t1_expanded AS (
2097-
SELECT
2098-
a,
2099-
b,
2096+
SELECT
2097+
a,
2098+
b,
21002099
CAST(a AS DOUBLE) / 1.5 AS c,
21012100
CAST(a AS DOUBLE) * CAST(a AS DOUBLE) AS d
21022101
FROM t1
21032102
CROSS JOIN (SELECT 1 AS dummy FROM t1 LIMIT 5)
21042103
),
21052104
t2_expanded AS (
2106-
SELECT
2105+
SELECT
21072106
a,
21082107
b,
21092108
CAST(a AS DOUBLE) * 2.5 AS e,
21102109
CAST(a AS DOUBLE) * CAST(a AS DOUBLE) * CAST(a AS DOUBLE) AS f
21112110
FROM t2
21122111
CROSS JOIN (SELECT 1 AS dummy FROM t2 LIMIT 5)
21132112
)
2114-
SELECT
2115-
t1.a, t1.b, t1.c, t1.d,
2113+
SELECT
2114+
t1.a, t1.b, t1.c, t1.d,
21162115
t2.a AS a2, t2.b AS b2, t2.e, t2.f
21172116
FROM t1_expanded t1
21182117
JOIN t2_expanded t2 ON t1.a % 100 = t2.a % 100
21192118
WHERE t1.a > 100 AND t2.a > 100
21202119
""")
2121-
2120+
21222121
# Flag to track if the query was interrupted
21232122
interrupted = False
21242123
interrupt_error = None
21252124
main_thread = threading.main_thread()
2126-
2125+
21272126
# Shared flag to indicate query execution has started
21282127
query_started = threading.Event()
21292128
max_wait_time = 5.0 # Maximum wait time in seconds
2130-
2131-
# This function will be run in a separate thread and will raise
2129+
2130+
# This function will be run in a separate thread and will raise
21322131
# KeyboardInterrupt in the main thread
21332132
def trigger_interrupt():
21342133
"""Poll for query start, then raise KeyboardInterrupt in the main thread"""
@@ -2139,31 +2138,33 @@ def trigger_interrupt():
21392138
if time.time() - start_time > max_wait_time:
21402139
msg = f"Query did not start within {max_wait_time} seconds"
21412140
raise RuntimeError(msg)
2142-
2141+
21432142
# Check if thread ID is available
21442143
thread_id = main_thread.ident
21452144
if thread_id is None:
21462145
msg = "Cannot get main thread ID"
21472146
raise RuntimeError(msg)
2148-
2147+
21492148
# Use ctypes to raise exception in main thread
21502149
exception = ctypes.py_object(KeyboardInterrupt)
21512150
res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
2152-
ctypes.c_long(thread_id), exception)
2151+
ctypes.c_long(thread_id), exception
2152+
)
21532153
if res != 1:
21542154
# If res is 0, the thread ID was invalid
21552155
# If res > 1, we modified multiple threads
21562156
ctypes.pythonapi.PyThreadState_SetAsyncExc(
2157-
ctypes.c_long(thread_id), ctypes.py_object(0))
2157+
ctypes.c_long(thread_id), ctypes.py_object(0)
2158+
)
21582159
msg = "Failed to raise KeyboardInterrupt in main thread"
21592160
raise RuntimeError(msg)
2160-
2161+
21612162
# Start a thread to trigger the interrupt
21622163
interrupt_thread = threading.Thread(target=trigger_interrupt)
2163-
# we mark as daemon so the test process can exit even if this thread doesnt finish
2164+
# we mark as daemon so the test process can exit even if this thread doesn't finish
21642165
interrupt_thread.daemon = True
21652166
interrupt_thread.start()
2166-
2167+
21672168
# Execute the query and expect it to be interrupted
21682169
try:
21692170
# Signal that we're about to start the query
@@ -2173,10 +2174,10 @@ def trigger_interrupt():
21732174
interrupted = True
21742175
except Exception as e:
21752176
interrupt_error = e
2176-
2177+
21772178
# Assert that the query was interrupted properly
21782179
if not interrupted:
21792180
pytest.fail(f"Query was not interrupted; got error: {interrupt_error}")
2180-
2181+
21812182
# Make sure the interrupt thread has finished
21822183
interrupt_thread.join(timeout=1.0)

0 commit comments

Comments
 (0)