is_direct_results -> has_more_rows

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -790,7 +790,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
 
-        is_direct_results = (
+        has_more_rows = (
             (not direct_results)
             or (not direct_results.resultSet)
             or direct_results.resultSet.hasMoreRows
@@ -831,7 +831,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             result_format=t_result_set_metadata_resp.resultFormat,
         )
 
-        return execute_response, is_direct_results
+        return execute_response, has_more_rows
 
     def get_execution_result(
         self, command_id: CommandId, cursor: Cursor
@@ -876,7 +876,7 @@ def get_execution_result(
 
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
         is_staging_operation = t_result_set_metadata_resp.isStagingOperation
-        is_direct_results = resp.hasMoreRows
+        has_more_rows = resp.hasMoreRows
 
         status = CommandState.from_thrift_state(resp.status) or CommandState.RUNNING
 
@@ -902,7 +902,7 @@ def get_execution_result(
             t_row_set=resp.results,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            has_more_rows=has_more_rows,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -1021,7 +1021,7 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, is_direct_results = self._handle_execute_response(
+            execute_response, has_more_rows = self._handle_execute_response(
                 resp, cursor
             )
 
@@ -1039,7 +1039,7 @@ def execute_command(
                 t_row_set=t_row_set,
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
-                is_direct_results=is_direct_results,
+                has_more_rows=has_more_rows,
                 session_id_hex=self._session_id_hex,
             )
 
@@ -1062,9 +1062,7 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1080,7 +1078,7 @@ def get_catalogs(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            has_more_rows=has_more_rows,
             session_id_hex=self._session_id_hex,
         )
 
@@ -1109,9 +1107,7 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1127,7 +1123,7 @@ def get_schemas(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            has_more_rows=has_more_rows,
             session_id_hex=self._session_id_hex,
         )
 
@@ -1160,9 +1156,7 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1178,7 +1172,7 @@ def get_tables(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            has_more_rows=has_more_rows,
             session_id_hex=self._session_id_hex,
         )
 
@@ -1211,9 +1205,7 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1229,7 +1221,7 @@ def get_columns(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            has_more_rows=has_more_rows,
             session_id_hex=self._session_id_hex,
         )
 

src/databricks/sql/result_set.py

@@ -43,7 +43,7 @@ def __init__(
         command_id: CommandId,
         status: CommandState,
         has_been_closed_server_side: bool = False,
-        is_direct_results: bool = False,
+        has_more_rows: bool = False,
         results_queue=None,
         description: List[Tuple] = [],
         is_staging_operation: bool = False,
@@ -61,7 +61,7 @@ def __init__(
             :param command_id: The command ID
             :param status: The command status
             :param has_been_closed_server_side: Whether the command has been closed on the server
-            :param is_direct_results: Whether the command has more rows
+            :param has_more_rows: Whether the command has more rows
             :param results_queue: The results queue
             :param description: column description of the results
             :param is_staging_operation: Whether the command is a staging operation
@@ -76,7 +76,7 @@ def __init__(
         self.command_id = command_id
         self.status = status
         self.has_been_closed_server_side = has_been_closed_server_side
-        self.is_direct_results = is_direct_results
+        self.has_more_rows = has_more_rows
         self.results = results_queue
         self._is_staging_operation = is_staging_operation
         self.lz4_compressed = lz4_compressed
@@ -200,7 +200,7 @@ def __init__(
         t_row_set=None,
         max_download_threads: int = 10,
         ssl_options=None,
-        is_direct_results: bool = True,
+        has_more_rows: bool = True,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -215,13 +215,13 @@ def __init__(
             :param t_row_set: The TRowSet containing result data (if available)
             :param max_download_threads: Maximum number of download threads for cloud fetch
             :param ssl_options: SSL options for cloud fetch
-            :param is_direct_results: Whether there are more rows to fetch
+            :param has_more_rows: Whether there are more rows to fetch
         """
         self.num_downloaded_chunks = 0
 
         # Initialize ThriftResultSet-specific attributes
         self._use_cloud_fetch = use_cloud_fetch
-        self.is_direct_results = is_direct_results
+        self.has_more_rows = has_more_rows
 
         # Build the results queue if t_row_set is provided
         results_queue = None
@@ -253,7 +253,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,
-            is_direct_results=is_direct_results,
+            has_more_rows=has_more_rows,
             results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
@@ -266,7 +266,7 @@ def __init__(
             self._fill_results_buffer()
 
     def _fill_results_buffer(self):
-        results, is_direct_results, result_links_count = self.backend.fetch_results(
+        results, has_more_rows, result_links_count = self.backend.fetch_results(
             command_id=self.command_id,
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
@@ -278,7 +278,7 @@ def _fill_results_buffer(self):
             chunk_id=self.num_downloaded_chunks,
         )
         self.results = results
-        self.is_direct_results = is_direct_results
+        self.has_more_rows = has_more_rows
         self.num_downloaded_chunks += result_links_count
 
     def _convert_columnar_table(self, table):
@@ -326,7 +326,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         while (
             n_remaining_rows > 0
             and not self.has_been_closed_server_side
-            and self.is_direct_results
+            and self.has_more_rows
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
@@ -351,7 +351,7 @@ def fetchmany_columnar(self, size: int):
         while (
             n_remaining_rows > 0
             and not self.has_been_closed_server_side
-            and self.is_direct_results
+            and self.has_more_rows
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
@@ -366,7 +366,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
         partial_result_chunks = [results]
-        while not self.has_been_closed_server_side and self.is_direct_results:
+        while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             if isinstance(results, ColumnTable) and isinstance(
@@ -392,7 +392,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.is_direct_results:
+        while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             results = self.merge_columnar(results, partial_results)

tests/unit/test_client.py

@@ -46,7 +46,7 @@ def new(cls):
             is_staging_operation=False,
             command_id=None,
             has_been_closed_server_side=True,
-            is_direct_results=True,
+            has_more_rows=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,
-                is_direct_results=False,
+                has_more_rows=False,
                 description=Mock(),
                 command_id=None,
                 arrow_schema_bytes=arrow_table.schema,

tests/unit/test_sea_result_set.py

@@ -46,7 +46,7 @@ def execute_response(self):
         mock_response.command_id = CommandId.from_sea_statement_id("test-statement-123")
         mock_response.status = CommandState.SUCCEEDED
         mock_response.has_been_closed_server_side = False
-        mock_response.is_direct_results = False
+        mock_response.has_more_rows = False
         mock_response.results_queue = None
         mock_response.description = [
             ("col1", "string", None, None, None, None, None),

tests/unit/test_thrift_backend.py

@@ -1016,10 +1016,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 is_direct_results, resp_type in itertools.product(
+        for has_more_rows, resp_type in itertools.product(
             [True, False], self.execute_response_types
         ):
-            with self.subTest(is_direct_results=is_direct_results, resp_type=resp_type):
+            with self.subTest(has_more_rows=has_more_rows, resp_type=resp_type):
                 tcli_service_instance = tcli_service_class.return_value
                 results_mock = Mock()
                 results_mock.startRowOffset = 0
@@ -1031,7 +1031,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=is_direct_results,
+                        hasMoreRows=has_more_rows,
                         results=results_mock,
                     ),
                     closeOperation=Mock(),
@@ -1052,7 +1052,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(is_direct_results, has_more_rows_result)
+                self.assertEqual(has_more_rows, has_more_rows_result)
 
     @patch(
         "databricks.sql.utils.ThriftResultSetQueueFactory.build_queue",
@@ -1062,10 +1062,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 is_direct_results, resp_type in itertools.product(
+        for has_more_rows, resp_type in itertools.product(
             [True, False], self.execute_response_types
         ):
-            with self.subTest(is_direct_results=is_direct_results, resp_type=resp_type):
+            with self.subTest(has_more_rows=has_more_rows, resp_type=resp_type):
                 tcli_service_instance = tcli_service_class.return_value
                 results_mock = MagicMock()
                 results_mock.startRowOffset = 0
@@ -1078,7 +1078,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
 
                 fetch_results_resp = ttypes.TFetchResultsResp(
                     status=self.okay_status,
-                    hasMoreRows=is_direct_results,
+                    hasMoreRows=has_more_rows,
                     results=results_mock,
                     resultSetMetadata=ttypes.TGetResultSetMetadataResp(
                         resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET
@@ -1112,7 +1112,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
                     chunk_id=0,
                 )
 
-                self.assertEqual(is_direct_results, has_more_rows_resp)
+                self.assertEqual(has_more_rows, 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):