Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitbb015e6

Browse files
mandate ResultData in SeaResultSet constructor
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent21e3078 commitbb015e6

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
returnfiltered_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__importannotations
2+
13
fromabcimportABC,abstractmethod
24
fromtypingimportList,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-
ifresult_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+
ifnotrows:
507+
returnpyarrow.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
exceptImportError:
2222
pyarrow=None
2323

24-
fromdatabricks.sqlimportOperationalError,exc
24+
fromdatabricks.sqlimportOperationalError
25+
fromdatabricks.sql.excimportProgrammingError
2526
fromdatabricks.sql.cloudfetch.download_managerimportResultFileDownloadManager
2627
fromdatabricks.sql.thrift_api.TCLIService.ttypesimport (
2728
TRowSet,
@@ -148,9 +149,7 @@ def build_queue(
148149
raiseNotImplementedError(
149150
"EXTERNAL_LINKS disposition is not implemented for SEA backend"
150151
)
151-
else:
152-
# Empty result set
153-
returnJsonQueue([])
152+
raiseProgrammingError("No result data or external links found")
154153

155154

156155
classJsonQueue(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
deftest_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
withpytest.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
deftest_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
withpytest.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
deftest_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)
[8]ページ先頭
©2009-2025 Movatter.jp