reduce repetition in cloud fetch queue

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/cloud_fetch_queue.py

@@ -162,6 +162,33 @@ def __init__(
         # Initialize download manager - will be set by subclasses
         self.download_manager: Optional["ResultFileDownloadManager"] = None
 
+    def remaining_rows(self) -> "pyarrow.Table":
+        """
+        Get all remaining rows of the cloud fetch Arrow dataframes.
+
+        Returns:
+            pyarrow.Table
+        """
+        if not self.table:
+            # Return empty pyarrow table to cause retry of fetch
+            return self._create_empty_table()
+
+        results = pyarrow.Table.from_pydict({})  # Empty table
+        while self.table:
+            table_slice = self.table.slice(
+                self.table_row_index, self.table.num_rows - self.table_row_index
+            )
+            if results.num_rows > 0:
+                results = pyarrow.concat_tables([results, table_slice])
+            else:
+                results = table_slice
+
+            self.table_row_index += table_slice.num_rows
+            self.table = self._create_next_table()
+            self.table_row_index = 0
+
+        return results
+
     def next_n_rows(self, num_rows: int) -> "pyarrow.Table":
         """Get up to the next n rows of the cloud fetch Arrow dataframes."""
         if not self.table:
@@ -216,20 +243,15 @@ def next_n_rows(self, num_rows: int) -> "pyarrow.Table":
         logger.info("SeaCloudFetchQueue: Retrieved {} rows".format(results.num_rows))
         return results
 
-    @abstractmethod
-    def remaining_rows(self) -> "pyarrow.Table":
-        """Get all remaining rows of the cloud fetch Arrow dataframes."""
-        pass
+    def _create_empty_table(self) -> "pyarrow.Table":
+        """Create a 0-row table with just the schema bytes."""
+        return create_arrow_table_from_arrow_file(self.schema_bytes, self.description)
 
     @abstractmethod
     def _create_next_table(self) -> Union["pyarrow.Table", None]:
         """Create next table by retrieving the logical next downloaded file."""
         pass
 
-    def _create_empty_table(self) -> "pyarrow.Table":
-        """Create a 0-row table with just the schema bytes."""
-        return create_arrow_table_from_arrow_file(self.schema_bytes, self.description)
-
 
 class SeaCloudFetchQueue(CloudFetchQueue):
     """Queue implementation for EXTERNAL_LINKS disposition with ARROW format for SEA backend."""
@@ -283,7 +305,7 @@ def __init__(
         )
 
         if initial_links:
-            initial_links = [] 
+            initial_links = []
             # logger.debug("SeaCloudFetchQueue: Initial links provided:")
             # for link in initial_links:
             #     logger.debug(
@@ -386,33 +408,6 @@ def _fetch_chunk_link(self, chunk_index: int) -> Optional["ExternalLink"]:
 
         return link
 
-    def remaining_rows(self) -> "pyarrow.Table":
-        """
-        Get all remaining rows of the cloud fetch Arrow dataframes.
-
-        Returns:
-            pyarrow.Table
-        """
-        if not self.table:
-            # Return empty pyarrow table to cause retry of fetch
-            return self._create_empty_table()
-
-        results = pyarrow.Table.from_pydict({})  # Empty table
-        while self.table:
-            table_slice = self.table.slice(
-                self.table_row_index, self.table.num_rows - self.table_row_index
-            )
-            if results.num_rows > 0:
-                results = pyarrow.concat_tables([results, table_slice])
-            else:
-                results = table_slice
-
-            self.table_row_index += table_slice.num_rows
-            self.table = self._create_next_table()
-            self.table_row_index = 0
-
-        return results
-
     def _create_next_table(self) -> Union["pyarrow.Table", None]:
         """Create next table by retrieving the logical next downloaded file."""
         # if we're still processing the current table, just return it
@@ -596,33 +591,6 @@ def __init__(
         # Initialize table and position
         self.table = self._create_next_table()
 
-    def remaining_rows(self) -> "pyarrow.Table":
-        """
-        Get all remaining rows of the cloud fetch Arrow dataframes.
-
-        Returns:
-            pyarrow.Table
-        """
-        if not self.table:
-            # Return empty pyarrow table to cause retry of fetch
-            return self._create_empty_table()
-
-        results = pyarrow.Table.from_pydict({})  # Empty table
-        while self.table:
-            table_slice = self.table.slice(
-                self.table_row_index, self.table.num_rows - self.table_row_index
-            )
-            if results.num_rows > 0:
-                results = pyarrow.concat_tables([results, table_slice])
-            else:
-                results = table_slice
-
-            self.table_row_index += table_slice.num_rows
-            self.table = self._create_next_table()
-            self.table_row_index = 0
-
-        return results
-
     def _create_next_table(self) -> Union["pyarrow.Table", None]:
         """Create next table by retrieving the logical next downloaded file."""
         logger.debug(