Skip to content

Commit cc514ad

Browse files
return None for immediate out of bounds
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 8f2953e commit cc514ad

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

src/databricks/sql/backend/sea/queue.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -173,12 +173,10 @@ def _convert_to_thrift_link(self, link: ExternalLink) -> TSparkArrowResultLink:
173173
httpHeaders=link.http_headers or {},
174174
)
175175

176-
def _get_chunk_link(self, chunk_index: int) -> ExternalLink:
176+
def _get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]:
177177
"""Progress to the next chunk link."""
178178
if chunk_index >= self._total_chunk_count:
179-
raise ValueError(
180-
f"Chunk index {chunk_index} is out of bounds, total chunks: {self._total_chunk_count}"
181-
)
179+
return None
182180

183181
try:
184182
return self._sea_client.get_chunk_link(self._statement_id, chunk_index)
@@ -208,4 +206,6 @@ def _create_next_table(self) -> Union["pyarrow.Table", None]:
208206
"""Create next table by retrieving the logical next downloaded file."""
209207
self._current_chunk_index += 1
210208
next_chunk_link = self._get_chunk_link(self._current_chunk_index)
209+
if not next_chunk_link:
210+
return None
211211
return self._create_table_from_link(next_chunk_link)

0 commit comments

Comments
 (0)