diff --git a/CHANGELOG.md b/CHANGELOG.md index b9f997b842c..8c2459a22b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,7 @@ - Fixed a bug in `session.get_session_stage` that referenced a non-existing stage after switching database or schema. - Fixed a bug where calling `DataFrame.to_snowpark_pandas_dataframe` without explicitly initializing the Snowpark pandas plugin caused an error. - Fixed a bug where using the `explode` function in dynamic table creation caused a SQL compilation error due to improper boolean type casting on the `outer` parameter. +- Fixed a bug where using `to_pandas_batches` with async jobs caused an error due to improper handling of waiting for asynchronous query completion. ### Snowpark Local Testing Updates diff --git a/src/snowflake/snowpark/_internal/server_connection.py b/src/snowflake/snowpark/_internal/server_connection.py index 22a394cda0f..498a1fe406d 100644 --- a/src/snowflake/snowpark/_internal/server_connection.py +++ b/src/snowflake/snowpark/_internal/server_connection.py @@ -465,7 +465,7 @@ def _to_data_or_iter( qid = results_cursor.sfqid if to_iter: new_cursor = results_cursor.connection.cursor() - new_cursor.execute(f"SELECT * FROM TABLE(RESULT_SCAN('{qid}'))") + new_cursor.get_results_from_sfqid(qid) results_cursor = new_cursor if to_pandas: diff --git a/tests/integ/scala/test_async_job_suite.py b/tests/integ/scala/test_async_job_suite.py index 347086971e0..be2ebdf4bf3 100644 --- a/tests/integ/scala/test_async_job_suite.py +++ b/tests/integ/scala/test_async_job_suite.py @@ -480,3 +480,16 @@ def test_async_job_result_wait_no_result(session): t1 = time() assert t1 - t0 >= 3.0 assert result is None + + +@pytest.mark.parametrize( + "action", + [ + lambda df: df.to_local_iterator(block=False), + lambda df: df.to_pandas_batches(block=False), + ], +) +def test_iter_cursor_wait_for_result(session, action): + df = session.sql("call system$wait(5)") + async_job = action(df) + assert async_job.result() is not None