File tree Expand file tree Collapse file tree 3 files changed +15
-5
lines changed Expand file tree Collapse file tree 3 files changed +15
-5
lines changed Original file line number Diff line number Diff line change @@ -73,11 +73,15 @@ def _filter_sea_result_set(
73
73
from databricks .sql .result_set import SeaResultSet
74
74
75
75
# Create a new SeaResultSet with the filtered data
76
+ manifest = result_set .manifest
77
+ manifest .total_row_count = len (filtered_rows )
78
+
76
79
filtered_result_set = SeaResultSet (
77
80
connection = result_set .connection ,
78
81
execute_response = execute_response ,
79
82
sea_client = cast (SeaDatabricksClient , result_set .backend ),
80
83
result_data = result_data ,
84
+ manifest = manifest ,
81
85
buffer_size_bytes = result_set .buffer_size_bytes ,
82
86
arraysize = result_set .arraysize ,
83
87
)
Original file line number Diff line number Diff line change @@ -369,7 +369,7 @@ def from_sea_statement_id(cls, statement_id: str):
369
369
370
370
return cls (BackendType .SEA , statement_id )
371
371
372
- def to_thrift_handle (self ):
372
+ def to_thrift_handle (self ) -> Optional [ ttypes . TOperationHandle ] :
373
373
"""
374
374
Convert this CommandId to a Thrift TOperationHandle.
375
375
@@ -390,7 +390,7 @@ def to_thrift_handle(self):
390
390
modifiedRowCount = self .modified_row_count ,
391
391
)
392
392
393
- def to_sea_statement_id (self ):
393
+ def to_sea_statement_id (self ) -> Optional [ str ] :
394
394
"""
395
395
Get the SEA statement ID string.
396
396
@@ -401,7 +401,7 @@ def to_sea_statement_id(self):
401
401
if self .backend_type != BackendType .SEA :
402
402
return None
403
403
404
- return self .guid
404
+ return str ( self .guid )
405
405
406
406
def to_hex_guid (self ) -> str :
407
407
"""
Original file line number Diff line number Diff line change @@ -473,10 +473,16 @@ def __init__(
473
473
manifest: Manifest from SEA response
474
474
"""
475
475
476
+ self .manifest = manifest
477
+
478
+ statement_id = execute_response .command_id .to_sea_statement_id ()
479
+ if statement_id is None :
480
+ raise ValueError ("Command ID is not a SEA statement ID" )
481
+
476
482
results_queue = SeaResultSetQueueFactory .build_queue (
477
483
result_data ,
478
- manifest ,
479
- str ( execute_response . command_id . to_sea_statement_id ()) ,
484
+ self . manifest ,
485
+ statement_id ,
480
486
description = execute_response .description ,
481
487
max_download_threads = sea_client .max_download_threads ,
482
488
sea_client = sea_client ,
You can’t perform that action at this time.
0 commit comments