@@ -370,6 +370,25 @@ def test_mixed_fetches_across_pages_keep_order_and_refill_offsets(self):
370370 self .assertEqual (rs .rownumber , 10 )
371371 self .assertEqual (offsets , [0 , 3 , 6 , 9 ])
372372
373+ def test_close_drops_fetchone_row_buffer (self ):
374+ # Rows parked in the fetchone() buffer must not survive close(). The warm
375+ # refill pulls every row into the buffer, so the queue itself is empty and
376+ # any row returned after close() could only come from a stale buffer.
377+ from databricks .sql .result_set import _FETCHONE_BUFFER_MAX_ROWS
378+
379+ for fetch_after_close in (
380+ lambda rs : rs .fetchall (),
381+ lambda rs : rs .fetchmany_arrow (10 ).to_pylist (),
382+ ):
383+ rs = self .make_dummy_result_set_from_initial_results ([[1 ], [2 ], [3 ], [4 ]])
384+ rs ._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
385+ self .assertSequenceEqual (rs .fetchone (), [1 ])
386+ self .assertEqual (rs ._row_buffer_count (), 3 )
387+ rs .close ()
388+ self .assertEqual (rs ._row_buffer_count (), 0 )
389+ self .assertEqual (fetch_after_close (rs ), [])
390+ self .assertIsNone (rs .fetchone ())
391+
373392 def test_fetchone_refill_does_not_fetch_extra_pages (self ):
374393 # A refill must not make extra round-trips to fill ``arraysize`` when
375394 # server pages are smaller (e.g. byte-limited): it serves what the current
@@ -408,6 +427,10 @@ def test_cloudfetch_fetchone_refill_capped_to_current_chunk(self):
408427 description = [("col0" , "integer" , None , None , None , None , None )],
409428 )
410429 rs .arraysize = 100
430+ # Skip the adaptive refill ramp so the chunk cap, not the ramp, bounds it.
431+ from databricks .sql .result_set import _FETCHONE_BUFFER_MAX_ROWS
432+
433+ rs ._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
411434
412435 self .assertSequenceEqual (rs .fetchone (), [8 ])
413436 queue .next_n_rows .assert_called_once_with (2 )
@@ -482,13 +505,21 @@ class ColumnarFetchoneBufferTests(unittest.TestCase):
482505 rows in batches (not one row at a time) while keeping order and rownumber."""
483506
484507 @staticmethod
485- def make_columnar_result_set (n_rows , arraysize ):
508+ def make_columnar_result_set (n_rows , arraysize , warm = True ):
509+ """Build a result set over ``n_rows`` columnar rows.
510+
511+ With ``warm`` (the default) the adaptive fetchone() refill ramp is skipped,
512+ so every refill is ``min(arraysize, _FETCHONE_BUFFER_MAX_ROWS)`` rows and the
513+ batch-boundary assertions below stay simple. Pass ``warm=False`` to test
514+ the ramp itself.
515+ """
516+ from databricks .sql .result_set import _FETCHONE_BUFFER_MAX_ROWS
486517 from databricks .sql .utils import ColumnQueue , ColumnTable
487518
488519 table = ColumnTable ([list (range (n_rows ))], ["col0" ])
489520 mock_thrift_backend = Mock (spec = ThriftDatabricksClient )
490521 mock_thrift_backend .fetch_results .return_value = (ColumnQueue (table ), False , 0 )
491- return ThriftResultSet (
522+ rs = ThriftResultSet (
492523 connection = Mock (),
493524 execute_response = ExecuteResponse (
494525 command_id = None ,
@@ -501,6 +532,9 @@ def make_columnar_result_set(n_rows, arraysize):
501532 thrift_client = mock_thrift_backend ,
502533 arraysize = arraysize ,
503534 )
535+ if warm :
536+ rs ._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
537+ return rs
504538
505539 @staticmethod
506540 def _count_conversions ():
@@ -543,6 +577,40 @@ def test_fetchone_refill_is_capped_for_large_arraysize(self):
543577 self .assertEqual ([r [0 ] for r in rs .fetchall ()], list (range (1 , n_rows )))
544578 self .assertEqual (rs .rownumber , n_rows )
545579
580+ def test_first_fetchone_converts_a_single_row (self ):
581+ # A lone fetchone() (SQLAlchemy .first()/.scalar(), EXISTS probes) must not
582+ # pull and convert a whole batch: the first refill is a single row.
583+ rs = self .make_columnar_result_set (5000 , arraysize = 100000 , warm = False )
584+ with self ._count_conversions () as m :
585+ self .assertEqual (rs .fetchone ()[0 ], 0 )
586+ self .assertEqual (m .call_count , 1 )
587+ self .assertEqual (m .call_args [0 ][1 ].num_rows , 1 )
588+ self .assertEqual (rs ._row_buffer_count (), 0 )
589+ self .assertEqual (rs .rownumber , 1 )
590+
591+ def test_fetchone_refill_size_ramps_up_to_cap (self ):
592+ from databricks .sql .result_set import _FETCHONE_BUFFER_MAX_ROWS
593+
594+ n_rows = 1 + 10 + 100 + _FETCHONE_BUFFER_MAX_ROWS * 2
595+ rs = self .make_columnar_result_set (n_rows , arraysize = 100000 , warm = False )
596+ with self ._count_conversions () as m :
597+ rows = [r [0 ] for r in rs ]
598+ self .assertEqual (rows , list (range (n_rows )))
599+ sizes = [c [0 ][1 ].num_rows for c in m .call_args_list ]
600+ self .assertEqual (
601+ sizes [:5 ],
602+ [1 , 10 , 100 , _FETCHONE_BUFFER_MAX_ROWS , _FETCHONE_BUFFER_MAX_ROWS ],
603+ )
604+ self .assertEqual (rs .rownumber , n_rows )
605+
606+ def test_fetchone_refill_ramp_is_capped_by_arraysize (self ):
607+ rs = self .make_columnar_result_set (10 , arraysize = 4 , warm = False )
608+ with self ._count_conversions () as m :
609+ rows = [r [0 ] for r in rs ]
610+ self .assertEqual (rows , list (range (10 )))
611+ sizes = [c [0 ][1 ].num_rows for c in m .call_args_list ]
612+ self .assertEqual (sizes [:4 ], [1 , 4 , 4 , 1 ])
613+
546614 def test_mixed_fetches_keep_order_and_rownumber (self ):
547615 rs = self .make_columnar_result_set (10 , arraysize = 4 )
548616 self .assertEqual (rs .fetchone ()[0 ], 0 )
0 commit comments