Skip to content

Commit bb015e6

Browse files
mandate ResultData in SeaResultSet constructor
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 21e3078 commit bb015e6

File tree

6 files changed

+56
-49
lines changed

6 files changed

+56
-49
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -616,10 +616,10 @@ def get_execution_result(
616616
connection=cursor.connection,
617617
execute_response=execute_response,
618618
sea_client=self,
619-
buffer_size_bytes=cursor.buffer_size_bytes,
620-
arraysize=cursor.arraysize,
621619
result_data=response.result,
622620
manifest=response.manifest,
621+
buffer_size_bytes=cursor.buffer_size_bytes,
622+
arraysize=cursor.arraysize,
623623
)
624624

625625
# == Metadata Operations ==

src/databricks/sql/backend/sea/utils/filters.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,9 +77,9 @@ def _filter_sea_result_set(
7777
connection=result_set.connection,
7878
execute_response=execute_response,
7979
sea_client=cast(SeaDatabricksClient, result_set.backend),
80+
result_data=result_data,
8081
buffer_size_bytes=result_set.buffer_size_bytes,
8182
arraysize=result_set.arraysize,
82-
result_data=result_data,
8383
)
8484

8585
return filtered_result_set

src/databricks/sql/result_set.py

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from __future__ import annotations
2+
13
from abc import ABC, abstractmethod
24
from typing import List, Optional, TYPE_CHECKING
35

@@ -450,13 +452,13 @@ class SeaResultSet(ResultSet):
450452

451453
def __init__(
452454
self,
453-
connection: "Connection",
454-
execute_response: "ExecuteResponse",
455-
sea_client: "SeaDatabricksClient",
455+
connection: Connection,
456+
execute_response: ExecuteResponse,
457+
sea_client: SeaDatabricksClient,
458+
result_data: ResultData,
459+
manifest: Optional[ResultManifest] = None,
456460
buffer_size_bytes: int = 104857600,
457461
arraysize: int = 10000,
458-
result_data: Optional["ResultData"] = None,
459-
manifest: Optional["ResultManifest"] = None,
460462
):
461463
"""
462464
Initialize a SeaResultSet with the response from a SEA query execution.
@@ -467,21 +469,19 @@ def __init__(
467469
sea_client: The SeaDatabricksClient instance for direct access
468470
buffer_size_bytes: Buffer size for fetching results
469471
arraysize: Default number of rows to fetch
470-
result_data: Result data from SEA response (optional)
471-
manifest: Manifest from SEA response (optional)
472+
result_data: Result data from SEA response
473+
manifest: Manifest from SEA response
472474
"""
473475

474-
results_queue = None
475-
if result_data:
476-
results_queue = SeaResultSetQueueFactory.build_queue(
477-
result_data,
478-
manifest,
479-
str(execute_response.command_id.to_sea_statement_id()),
480-
description=execute_response.description,
481-
max_download_threads=sea_client.max_download_threads,
482-
sea_client=sea_client,
483-
lz4_compressed=execute_response.lz4_compressed,
484-
)
476+
results_queue = SeaResultSetQueueFactory.build_queue(
477+
result_data,
478+
manifest,
479+
str(execute_response.command_id.to_sea_statement_id()),
480+
description=execute_response.description,
481+
max_download_threads=sea_client.max_download_threads,
482+
sea_client=sea_client,
483+
lz4_compressed=execute_response.lz4_compressed,
484+
)
485485

486486
# Call parent constructor with common attributes
487487
super().__init__(
@@ -503,6 +503,8 @@ def _convert_json_to_arrow(self, rows: List) -> "pyarrow.Table":
503503
"""
504504
Convert raw data rows to Arrow table.
505505
"""
506+
if not rows:
507+
return pyarrow.Table.from_pydict({})
506508

507509
columns = []
508510
num_cols = len(rows[0])

src/databricks/sql/utils.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@
2121
except ImportError:
2222
pyarrow = None
2323

24-
from databricks.sql import OperationalError, exc
24+
from databricks.sql import OperationalError
25+
from databricks.sql.exc import ProgrammingError
2526
from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
2627
from databricks.sql.thrift_api.TCLIService.ttypes import (
2728
TRowSet,
@@ -148,9 +149,7 @@ def build_queue(
148149
raise NotImplementedError(
149150
"EXTERNAL_LINKS disposition is not implemented for SEA backend"
150151
)
151-
else:
152-
# Empty result set
153-
return JsonQueue([])
152+
raise ProgrammingError("No result data or external links found")
154153

155154

156155
class JsonQueue(ResultSetQueue):

tests/unit/test_sea_queue.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ def test_build_queue_with_inline_data(self, mock_sea_client, mock_description):
135135
def test_build_queue_with_empty_data(self, mock_sea_client, mock_description):
136136
"""Test building a queue with empty data."""
137137
# Create a ResultData object with no data
138-
result_data = ResultData(data=None, external_links=None, row_count=0)
138+
result_data = ResultData(data=[], external_links=None, row_count=0)
139139

140140
# Build the queue
141141
queue = SeaResultSetQueueFactory.build_queue(

tests/unit/test_sea_result_set.py

Lines changed: 29 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -74,10 +74,10 @@ def result_set_with_data(
7474
connection=mock_connection,
7575
execute_response=execute_response,
7676
sea_client=mock_sea_client,
77-
buffer_size_bytes=1000,
78-
arraysize=100,
7977
result_data=result_data,
8078
manifest=None,
79+
buffer_size_bytes=1000,
80+
arraysize=100,
8181
)
8282
result_set.results = JsonQueue(sample_data)
8383

@@ -96,6 +96,7 @@ def test_init_with_execute_response(
9696
connection=mock_connection,
9797
execute_response=execute_response,
9898
sea_client=mock_sea_client,
99+
result_data=ResultData(data=[]),
99100
buffer_size_bytes=1000,
100101
arraysize=100,
101102
)
@@ -115,6 +116,7 @@ def test_close(self, mock_connection, mock_sea_client, execute_response):
115116
connection=mock_connection,
116117
execute_response=execute_response,
117118
sea_client=mock_sea_client,
119+
result_data=ResultData(data=[]),
118120
buffer_size_bytes=1000,
119121
arraysize=100,
120122
)
@@ -135,6 +137,7 @@ def test_close_when_already_closed_server_side(
135137
connection=mock_connection,
136138
execute_response=execute_response,
137139
sea_client=mock_sea_client,
140+
result_data=ResultData(data=[]),
138141
buffer_size_bytes=1000,
139142
arraysize=100,
140143
)
@@ -157,6 +160,7 @@ def test_close_when_connection_closed(
157160
connection=mock_connection,
158161
execute_response=execute_response,
159162
sea_client=mock_sea_client,
163+
result_data=ResultData(data=[]),
160164
buffer_size_bytes=1000,
161165
arraysize=100,
162166
)
@@ -301,39 +305,40 @@ def test_fetchmany_arrow_not_implemented(
301305
self, mock_connection, mock_sea_client, execute_response, sample_data
302306
):
303307
"""Test that fetchmany_arrow raises NotImplementedError for non-JSON data."""
304-
# Create a result set without JSON data
305-
result_set = SeaResultSet(
306-
connection=mock_connection,
307-
execute_response=execute_response,
308-
sea_client=mock_sea_client,
309-
buffer_size_bytes=1000,
310-
arraysize=100,
311-
)
312308

313309
# Test that NotImplementedError is raised
314310
with pytest.raises(
315-
NotImplementedError, match="fetchmany_arrow only supported for JSON data"
311+
NotImplementedError,
312+
match="EXTERNAL_LINKS disposition is not implemented for SEA backend",
316313
):
317-
result_set.fetchmany_arrow(10)
314+
# Create a result set without JSON data
315+
result_set = SeaResultSet(
316+
connection=mock_connection,
317+
execute_response=execute_response,
318+
sea_client=mock_sea_client,
319+
result_data=ResultData(data=None, external_links=[]),
320+
buffer_size_bytes=1000,
321+
arraysize=100,
322+
)
318323

319324
def test_fetchall_arrow_not_implemented(
320325
self, mock_connection, mock_sea_client, execute_response, sample_data
321326
):
322327
"""Test that fetchall_arrow raises NotImplementedError for non-JSON data."""
323-
# Create a result set without JSON data
324-
result_set = SeaResultSet(
325-
connection=mock_connection,
326-
execute_response=execute_response,
327-
sea_client=mock_sea_client,
328-
buffer_size_bytes=1000,
329-
arraysize=100,
330-
)
331-
332328
# Test that NotImplementedError is raised
333329
with pytest.raises(
334-
NotImplementedError, match="fetchall_arrow only supported for JSON data"
330+
NotImplementedError,
331+
match="EXTERNAL_LINKS disposition is not implemented for SEA backend",
335332
):
336-
result_set.fetchall_arrow()
333+
# Create a result set without JSON data
334+
result_set = SeaResultSet(
335+
connection=mock_connection,
336+
execute_response=execute_response,
337+
sea_client=mock_sea_client,
338+
result_data=ResultData(data=None, external_links=[]),
339+
buffer_size_bytes=1000,
340+
arraysize=100,
341+
)
337342

338343
def test_is_staging_operation(
339344
self, mock_connection, mock_sea_client, execute_response
@@ -347,6 +352,7 @@ def test_is_staging_operation(
347352
connection=mock_connection,
348353
execute_response=execute_response,
349354
sea_client=mock_sea_client,
355+
result_data=ResultData(data=[]),
350356
buffer_size_bytes=1000,
351357
arraysize=100,
352358
)

0 commit comments

Comments
 (0)