has_more_rows -> is_direct_results

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -758,7 +758,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
 
-        has_more_rows = (
+        is_direct_results = (
             (not direct_results)
             or (not direct_results.resultSet)
             or direct_results.resultSet.hasMoreRows
@@ -796,7 +796,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             result_format=t_result_set_metadata_resp.resultFormat,
         )
 
-        return execute_response, has_more_rows
+        return execute_response, is_direct_results
 
     def get_execution_result(
         self, command_id: CommandId, cursor: "Cursor"
@@ -838,7 +838,7 @@ def get_execution_result(
 
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
         is_staging_operation = t_result_set_metadata_resp.isStagingOperation
-        has_more_rows = resp.hasMoreRows
+        is_direct_results = resp.hasMoreRows
 
         status = self.get_query_state(command_id)
 
@@ -863,7 +863,7 @@ def get_execution_result(
             t_row_set=resp.results,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            has_more_rows=has_more_rows,
+            is_direct_results=is_direct_results,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -976,7 +976,7 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, has_more_rows = self._handle_execute_response(
+            execute_response, is_direct_results = self._handle_execute_response(
                 resp, cursor
             )
 
@@ -994,7 +994,7 @@ def execute_command(
                 t_row_set=t_row_set,
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
-                has_more_rows=has_more_rows,
+                is_direct_results=is_direct_results,
             )
 
     def get_catalogs(
@@ -1016,7 +1016,9 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
+        execute_response, is_direct_results = self._handle_execute_response(
+            resp, cursor
+        )
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1032,7 +1034,7 @@ def get_catalogs(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            has_more_rows=has_more_rows,
+            is_direct_results=is_direct_results,
         )
 
     def get_schemas(
@@ -1058,7 +1060,9 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
+        execute_response, is_direct_results = self._handle_execute_response(
+            resp, cursor
+        )
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1074,7 +1078,7 @@ def get_schemas(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            has_more_rows=has_more_rows,
+            is_direct_results=is_direct_results,
         )
 
     def get_tables(
@@ -1104,7 +1108,9 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
+        execute_response, is_direct_results = self._handle_execute_response(
+            resp, cursor
+        )
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1120,7 +1126,7 @@ def get_tables(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            has_more_rows=has_more_rows,
+            is_direct_results=is_direct_results,
         )
 
     def get_columns(
@@ -1150,7 +1156,9 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
+        execute_response, is_direct_results = self._handle_execute_response(
+            resp, cursor
+        )
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1166,7 +1174,7 @@ def get_columns(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            has_more_rows=has_more_rows,
+            is_direct_results=is_direct_results,
         )
 
     def _handle_execute_response(self, resp, cursor):

src/databricks/sql/result_set.py

@@ -41,7 +41,7 @@ def __init__(
         command_id: CommandId,
         status: CommandState,
         has_been_closed_server_side: bool = False,
-        has_more_rows: bool = False,
+        is_direct_results: bool = False,
         results_queue=None,
         description=None,
         is_staging_operation: bool = False,
@@ -57,7 +57,7 @@ def __init__(
             command_id: The command ID
             status: The command status
             has_been_closed_server_side: Whether the command has been closed on the server
-            has_more_rows: Whether the command has more rows
+            is_direct_results: Whether the command has more rows
             results_queue: The results queue
             description: column description of the results
             is_staging_operation: Whether the command is a staging operation
@@ -72,7 +72,7 @@ def __init__(
         self.command_id = command_id
         self.status = status
         self.has_been_closed_server_side = has_been_closed_server_side
-        self.has_more_rows = has_more_rows
+        self.is_direct_results = is_direct_results
         self.results = results_queue
         self._is_staging_operation = is_staging_operation
 
@@ -160,7 +160,7 @@ def __init__(
         t_row_set=None,
         max_download_threads: int = 10,
         ssl_options=None,
-        has_more_rows: bool = True,
+        is_direct_results: bool = True,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -175,7 +175,7 @@ def __init__(
             t_row_set: The TRowSet containing result data (if available)
             max_download_threads: Maximum number of download threads for cloud fetch
             ssl_options: SSL options for cloud fetch
-            has_more_rows: Whether there are more rows to fetch
+            is_direct_results: Whether there are more rows to fetch
         """
         # Initialize ThriftResultSet-specific attributes
         self._arrow_schema_bytes = execute_response.arrow_schema_bytes
@@ -207,7 +207,7 @@ def __init__(
             command_id=execute_response.command_id,
             status=execute_response.status,
             has_been_closed_server_side=execute_response.has_been_closed_server_side,
-            has_more_rows=has_more_rows,
+            is_direct_results=is_direct_results,
             results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
@@ -218,7 +218,7 @@ def __init__(
             self._fill_results_buffer()
 
     def _fill_results_buffer(self):
-        results, has_more_rows = self.backend.fetch_results(
+        results, is_direct_results = self.backend.fetch_results(
             command_id=self.command_id,
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
@@ -229,7 +229,7 @@ def _fill_results_buffer(self):
             use_cloud_fetch=self._use_cloud_fetch,
         )
         self.results = results
-        self.has_more_rows = has_more_rows
+        self.is_direct_results = is_direct_results
 
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]
@@ -313,7 +313,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         while (
             n_remaining_rows > 0
             and not self.has_been_closed_server_side
-            and self.has_more_rows
+            and self.is_direct_results
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
@@ -338,7 +338,7 @@ def fetchmany_columnar(self, size: int):
         while (
             n_remaining_rows > 0
             and not self.has_been_closed_server_side
-            and self.has_more_rows
+            and self.is_direct_results
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
@@ -353,7 +353,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
 
-        while not self.has_been_closed_server_side and self.has_more_rows:
+        while not self.has_been_closed_server_side and self.is_direct_results:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             if isinstance(results, ColumnTable) and isinstance(
@@ -379,7 +379,7 @@ def fetchall_columnar(self):
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
 
-        while not self.has_been_closed_server_side and self.has_more_rows:
+        while not self.has_been_closed_server_side and self.is_direct_results:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             results = self.merge_columnar(results, partial_results)

tests/unit/test_client.py

@@ -48,7 +48,7 @@ def new(cls):
             is_staging_operation=False,
             command_id=None,
             has_been_closed_server_side=True,
-            has_more_rows=True,
+            is_direct_results=True,
             lz4_compressed=True,
             arrow_schema_bytes=b"schema",
         )

tests/unit/test_fetches_bench.py

@@ -36,7 +36,7 @@ def make_dummy_result_set_from_initial_results(arrow_table):
             execute_response=ExecuteResponse(
                 status=None,
                 has_been_closed_server_side=True,
-                has_more_rows=False,
+                is_direct_results=False,
                 description=Mock(),
                 command_id=None,
                 arrow_queue=arrow_queue,

tests/unit/test_thrift_backend.py

@@ -1004,10 +1004,10 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
     def test_handle_execute_response_reads_has_more_rows_in_direct_results(
         self, tcli_service_class, build_queue
     ):
-        for has_more_rows, resp_type in itertools.product(
+        for is_direct_results, resp_type in itertools.product(
             [True, False], self.execute_response_types
         ):
-            with self.subTest(has_more_rows=has_more_rows, resp_type=resp_type):
+            with self.subTest(is_direct_results=is_direct_results, resp_type=resp_type):
                 tcli_service_instance = tcli_service_class.return_value
                 results_mock = Mock()
                 results_mock.startRowOffset = 0
@@ -1019,7 +1019,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                     resultSetMetadata=self.metadata_resp,
                     resultSet=ttypes.TFetchResultsResp(
                         status=self.okay_status,
-                        hasMoreRows=has_more_rows,
+                        hasMoreRows=is_direct_results,
                         results=results_mock,
                     ),
                     closeOperation=Mock(),
@@ -1040,7 +1040,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                     has_more_rows_result,
                 ) = thrift_backend._handle_execute_response(execute_resp, Mock())
 
-                self.assertEqual(has_more_rows, has_more_rows_result)
+                self.assertEqual(is_direct_results, has_more_rows_result)
 
     @patch(
         "databricks.sql.utils.ResultSetQueueFactory.build_queue", return_value=Mock()
@@ -1049,10 +1049,10 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
     def test_handle_execute_response_reads_has_more_rows_in_result_response(
         self, tcli_service_class, build_queue
     ):
-        for has_more_rows, resp_type in itertools.product(
+        for is_direct_results, resp_type in itertools.product(
             [True, False], self.execute_response_types
         ):
-            with self.subTest(has_more_rows=has_more_rows, resp_type=resp_type):
+            with self.subTest(is_direct_results=is_direct_results, resp_type=resp_type):
                 tcli_service_instance = tcli_service_class.return_value
                 results_mock = MagicMock()
                 results_mock.startRowOffset = 0
@@ -1065,7 +1065,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
 
                 fetch_results_resp = ttypes.TFetchResultsResp(
                     status=self.okay_status,
-                    hasMoreRows=has_more_rows,
+                    hasMoreRows=is_direct_results,
                     results=results_mock,
                     resultSetMetadata=ttypes.TGetResultSetMetadataResp(
                         resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET
@@ -1098,7 +1098,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
                     description=Mock(),
                 )
 
-                self.assertEqual(has_more_rows, has_more_rows_resp)
+                self.assertEqual(is_direct_results, has_more_rows_resp)
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_arrow_batches_row_count_are_respected(self, tcli_service_class):