make get_chunk_link return mandatory ExternalLink

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

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

@@ -178,9 +178,11 @@ def _trigger_next_batch_download(self) -> bool:
 
         return True
 
-    def get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]:
+    def get_chunk_link(self, chunk_index: int) -> ExternalLink:
         if chunk_index >= self.total_chunk_count:
-            return None
+            raise ValueError(
+                f"Chunk index {chunk_index} is out of range for total chunk count {self.total_chunk_count}"
+            )
 
         with self._link_data_update:
             while chunk_index not in self.chunk_index_to_link:
@@ -194,7 +196,7 @@ def get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]:
                     )
                 self._link_data_update.wait()
 
-            return self.chunk_index_to_link.get(chunk_index, None)
+            return self.chunk_index_to_link[chunk_index]
 
     @staticmethod
     def _convert_to_thrift_link(link: ExternalLink) -> TSparkArrowResultLink:
@@ -296,8 +298,6 @@ def _create_next_table(self) -> Union["pyarrow.Table", None]:
             return None
 
         chunk_link = self.link_fetcher.get_chunk_link(self._current_chunk_index)
-        if not chunk_link:
-            return None
 
         row_offset = chunk_link.row_offset
         arrow_table = self._create_table_at_offset(row_offset)

tests/unit/test_sea_queue.py

@@ -717,10 +717,14 @@ def _worker():
         # The thread should have finished and captured link1
         assert result_container.get("link") == link1
 
-    def test_get_chunk_link_out_of_range_returns_none(self, sample_links):
+    def test_get_chunk_link_out_of_range_raises_value_error(self, sample_links):
         """Requesting a chunk index >= total_chunk_count should immediately return None."""
         link0, _ = sample_links
 
         fetcher, _backend, _dm = self._create_fetcher([link0], total_chunk_count=1)
 
-        assert fetcher.get_chunk_link(10) is None
+        with pytest.raises(
+            ValueError,
+            match="Chunk index 10 is out of range for total chunk count 1",
+        ):
+            fetcher.get_chunk_link(10)