1616 from databricks .sql .client import Connection
1717from databricks .sql .backend .databricks_client import DatabricksClient
1818from databricks .sql .types import Row
19- from databricks .sql .exc import RequestError , CursorAlreadyClosedError
19+ from databricks .sql .exc import InternalError , RequestError , CursorAlreadyClosedError
2020from databricks .sql .utils import (
21+ CloudFetchQueue ,
2122 ColumnTable ,
2223 ColumnQueue ,
2324 concat_table_chunks ,
@@ -196,6 +197,14 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
196197 raise row .error
197198 return taken
198199
200+ @staticmethod
201+ def _raw_row_count (raw : Any ) -> Optional [int ]:
202+ """Row count of a raw chunk (Arrow table, ColumnTable or list); None if unknown."""
203+ if raw is None :
204+ return 0
205+ n = len (raw ) if isinstance (raw , list ) else getattr (raw , "num_rows" , None )
206+ return n if isinstance (n , int ) else None
207+
199208 @staticmethod
200209 def _convert_rows_individually (
201210 raw : Any , convert : Callable [[Any ], List [Row ]]
@@ -206,9 +215,7 @@ def _convert_rows_individually(
206215 raised only when that row is handed out. Returns None if ``raw`` has no rows
207216 to split, in which case the caller re-raises the batch error.
208217 """
209- if raw is None :
210- return None
211- n = len (raw ) if isinstance (raw , list ) else getattr (raw , "num_rows" , 0 )
218+ n = ResultSet ._raw_row_count (raw )
212219 if not n :
213220 return None
214221 rows : List [Any ] = []
@@ -275,6 +282,16 @@ def _fetchone_buffered(
275282 rows = fallback
276283 if not rows :
277284 return None
285+ raw_count = self ._raw_row_count (raw )
286+ if raw_count is not None and raw_count != len (rows ):
287+ # fetch_raw advanced _next_row_index (and the Thrift fetch offset)
288+ # by raw_count; the buffer must line up 1:1 with the raw chunk or
289+ # rownumber and fetchmany_arrow slicing would silently drift.
290+ raise InternalError (
291+ "fetchone(): converted {} row(s) from a chunk of {} raw row(s)" .format (
292+ len (rows ), raw_count
293+ )
294+ )
278295 self ._row_buffer_raw = raw
279296 self ._row_buffer_rows = rows
280297 self ._row_buffer_pos = 0
@@ -552,19 +569,31 @@ def _fetch_available_from_queue(self, size: int):
552569 until ``size`` rows are collected: it only refills the queue while it is
553570 empty, so fetchone() makes no more round-trips than it needs to return one
554571 row (server pages can be smaller than ``size`` when they hit the byte limit).
572+ On CloudFetch, ``next_n_rows`` keeps walking downloaded files until it has
573+ ``size`` rows, so the request is also capped at what the current chunk still
574+ holds; otherwise a refill near a chunk boundary would block on the next
575+ download before returning a row that is already available.
555576 Only called while the fetchone() row buffer is empty.
556577 """
557- results = self .results .next_n_rows (size )
578+ results = self .results .next_n_rows (self . _cap_to_current_chunk ( size ) )
558579 while (
559580 results .num_rows == 0
560581 and not self .has_been_closed_server_side
561582 and self .has_more_rows
562583 ):
563584 self ._fill_results_buffer ()
564- results = self .results .next_n_rows (size )
585+ results = self .results .next_n_rows (self . _cap_to_current_chunk ( size ) )
565586 self ._next_row_index += results .num_rows
566587 return results
567588
589+ def _cap_to_current_chunk (self , size : int ) -> int :
590+ """Cap ``size`` at the rows left in the current CloudFetch chunk, if any."""
591+ if isinstance (self .results , CloudFetchQueue ) and self .results .table :
592+ available = self .results .table .num_rows - self .results .table_row_index
593+ if available > 0 :
594+ return min (size , available )
595+ return size
596+
568597 def fetchall_arrow (self ) -> "pyarrow.Table" :
569598 """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
570599 buffered = self ._row_buffer_take_raw ()
0 commit comments