align remaining_rows with prev impl

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

Author: varun-edachali-dbx

src/databricks/sql/utils.py

@@ -255,27 +255,19 @@ def next_n_rows(self, num_rows: int) -> "pyarrow.Table":
             # Return empty pyarrow table to cause retry of fetch
             return self._create_empty_table()
         logger.debug("CloudFetchQueue: trying to get {} next rows".format(num_rows))
-
         results = self.table.slice(0, 0)
-
         while num_rows > 0 and self.table:
             # Get remaining of num_rows or the rest of the current table, whichever is smaller
             length = min(num_rows, self.table.num_rows - self.table_row_index)
             table_slice = self.table.slice(self.table_row_index, length)
-
-            # Concatenate results if we have any
-            if results.num_rows > 0:
-                results = pyarrow.concat_tables([results, table_slice])
-            else:
-                results = table_slice
+            results = pyarrow.concat_tables([results, table_slice])
 
             self.table_row_index += table_slice.num_rows
 
             # Replace current table with the next table if we are at the end of the current table
             if self.table_row_index == self.table.num_rows:
                 self.table = self._create_next_table()
                 self.table_row_index = 0
-
             num_rows -= table_slice.num_rows
 
         logger.debug("CloudFetchQueue: collected {} next rows".format(results.num_rows))
@@ -288,32 +280,23 @@ def remaining_rows(self) -> "pyarrow.Table":
         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
+        results = self.table.slice(0, 0)
         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
+            results = pyarrow.concat_tables([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_empty_table(self) -> "pyarrow.Table":
-        """Create a 0-row table with just the schema bytes."""
-        if not self.schema_bytes:
-            return pyarrow.Table.from_pydict({})
-        return create_arrow_table_from_arrow_file(self.schema_bytes, self.description)
-
     def _create_table_at_offset(self, offset: int) -> Union["pyarrow.Table", None]:
         """Create next table by retrieving the logical next downloaded file."""
         # Create next table by retrieving the logical next downloaded file, or return None to signal end of queue
@@ -323,6 +306,9 @@ def _create_table_at_offset(self, offset: int) -> Union["pyarrow.Table", None]:
 
         downloaded_file = self.download_manager.get_next_downloaded_file(offset)
         if not downloaded_file:
+            logger.debug(
+                "CloudFetchQueue: Cannot find downloaded file for row {}".format(offset)
+            )
             # None signals no more Arrow tables can be built from the remaining handlers if any remain
             return None
 
@@ -345,6 +331,12 @@ 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."""
+        if not self.schema_bytes:
+            return pyarrow.Table.from_pydict({})
+        return create_arrow_table_from_arrow_file(self.schema_bytes, self.description)
+
 
 class ThriftCloudFetchQueue(CloudFetchQueue):
     """Queue implementation for EXTERNAL_LINKS disposition with ARROW format for Thrift backend."""