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

Commitfd81c5a

Browse files
authored
Concat tables to be backward compatible (#647)
* fixed* Minor fix* more types
1 parent701f7f6 commitfd81c5a

File tree

3 files changed

+77
-34
lines changed

3 files changed

+77
-34
lines changed

‎src/databricks/sql/result_set.py‎

Lines changed: 15 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
fromdatabricks.sql.utilsimport (
2121
ColumnTable,
2222
ColumnQueue,
23+
concat_table_chunks,
2324
)
2425
fromdatabricks.sql.backend.typesimportCommandId,CommandState,ExecuteResponse
2526
fromdatabricks.sql.telemetry.models.eventimportStatementType
@@ -296,23 +297,6 @@ def _convert_columnar_table(self, table):
296297

297298
returnresult
298299

299-
defmerge_columnar(self,result1,result2)->"ColumnTable":
300-
"""
301-
Function to merge / combining the columnar results into a single result
302-
:param result1:
303-
:param result2:
304-
:return:
305-
"""
306-
307-
ifresult1.column_names!=result2.column_names:
308-
raiseValueError("The columns in the results don't match")
309-
310-
merged_result= [
311-
result1.column_table[i]+result2.column_table[i]
312-
foriinrange(result1.num_columns)
313-
]
314-
returnColumnTable(merged_result,result1.column_names)
315-
316300
deffetchmany_arrow(self,size:int)->"pyarrow.Table":
317301
"""
318302
Fetch the next set of rows of a query result, returning a PyArrow table.
@@ -337,7 +321,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
337321
n_remaining_rows-=partial_results.num_rows
338322
self._next_row_index+=partial_results.num_rows
339323

340-
returnpyarrow.concat_tables(partial_result_chunks,use_threads=True)
324+
returnconcat_table_chunks(partial_result_chunks)
341325

342326
deffetchmany_columnar(self,size:int):
343327
"""
@@ -350,19 +334,19 @@ def fetchmany_columnar(self, size: int):
350334
results=self.results.next_n_rows(size)
351335
n_remaining_rows=size-results.num_rows
352336
self._next_row_index+=results.num_rows
353-
337+
partial_result_chunks= [results]
354338
while (
355339
n_remaining_rows>0
356340
andnotself.has_been_closed_server_side
357341
andself.has_more_rows
358342
):
359343
self._fill_results_buffer()
360344
partial_results=self.results.next_n_rows(n_remaining_rows)
361-
results=self.merge_columnar(results,partial_results)
345+
partial_result_chunks.append(partial_results)
362346
n_remaining_rows-=partial_results.num_rows
363347
self._next_row_index+=partial_results.num_rows
364348

365-
returnresults
349+
returnconcat_table_chunks(partial_result_chunks)
366350

367351
deffetchall_arrow(self)->"pyarrow.Table":
368352
"""Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
@@ -372,36 +356,34 @@ def fetchall_arrow(self) -> "pyarrow.Table":
372356
whilenotself.has_been_closed_server_sideandself.has_more_rows:
373357
self._fill_results_buffer()
374358
partial_results=self.results.remaining_rows()
375-
ifisinstance(results,ColumnTable)andisinstance(
376-
partial_results,ColumnTable
377-
):
378-
results=self.merge_columnar(results,partial_results)
379-
else:
380-
partial_result_chunks.append(partial_results)
359+
partial_result_chunks.append(partial_results)
381360
self._next_row_index+=partial_results.num_rows
382361

362+
result_table=concat_table_chunks(partial_result_chunks)
383363
# If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table
384364
# Valid only for metadata commands result set
385-
ifisinstance(results,ColumnTable)andpyarrow:
365+
ifisinstance(result_table,ColumnTable)andpyarrow:
386366
data= {
387367
name:col
388-
forname,colinzip(results.column_names,results.column_table)
368+
forname,colinzip(
369+
result_table.column_names,result_table.column_table
370+
)
389371
}
390372
returnpyarrow.Table.from_pydict(data)
391-
returnpyarrow.concat_tables(partial_result_chunks,use_threads=True)
373+
returnresult_table
392374

393375
deffetchall_columnar(self):
394376
"""Fetch all (remaining) rows of a query result, returning them as a Columnar table."""
395377
results=self.results.remaining_rows()
396378
self._next_row_index+=results.num_rows
397-
379+
partial_result_chunks= [results]
398380
whilenotself.has_been_closed_server_sideandself.has_more_rows:
399381
self._fill_results_buffer()
400382
partial_results=self.results.remaining_rows()
401-
results=self.merge_columnar(results,partial_results)
383+
partial_result_chunks.append(partial_results)
402384
self._next_row_index+=partial_results.num_rows
403385

404-
returnresults
386+
returnconcat_table_chunks(partial_result_chunks)
405387

406388
deffetchone(self)->Optional[Row]:
407389
"""

‎src/databricks/sql/utils.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -853,3 +853,25 @@ def _create_python_tuple(t_col_value_wrapper):
853853
result[i]=None
854854

855855
returntuple(result)
856+
857+
858+
defconcat_table_chunks(
859+
table_chunks:List[Union["pyarrow.Table",ColumnTable]]
860+
)->Union["pyarrow.Table",ColumnTable]:
861+
iflen(table_chunks)==0:
862+
returntable_chunks
863+
864+
ifisinstance(table_chunks[0],ColumnTable):
865+
## Check if all have the same column names
866+
ifnotall(
867+
table.column_names==table_chunks[0].column_namesfortableintable_chunks
868+
):
869+
raiseValueError("The columns in the results don't match")
870+
871+
result_table:List[List[Any]]= [[]for_inrange(table_chunks[0].num_columns)]
872+
foriinrange(0,len(table_chunks)):
873+
forjinrange(table_chunks[i].num_columns):
874+
result_table[j].extend(table_chunks[i].column_table[j])
875+
returnColumnTable(result_table,table_chunks[0].column_names)
876+
else:
877+
returnpyarrow.concat_tables(table_chunks,use_threads=True)

‎tests/unit/test_util.py‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,17 @@
11
importdecimal
22
importdatetime
33
fromdatetimeimporttimezone,timedelta
4+
importpytest
5+
fromdatabricks.sql.utilsimport (
6+
convert_to_assigned_datatypes_in_column_table,
7+
ColumnTable,
8+
concat_table_chunks,
9+
)
410

5-
fromdatabricks.sql.utilsimportconvert_to_assigned_datatypes_in_column_table
11+
try:
12+
importpyarrow
13+
exceptImportError:
14+
pyarrow=None
615

716

817
classTestUtils:
@@ -122,3 +131,33 @@ def test_convert_to_assigned_datatypes_in_column_table(self):
122131
forindex,entryinenumerate(converted_column_table):
123132
assertentry[0]==expected_convertion[index][0]
124133
assertisinstance(entry[0],expected_convertion[index][1])
134+
135+
deftest_concat_table_chunks_column_table(self):
136+
column_table1=ColumnTable([[1,2], [5,6]], ["col1","col2"])
137+
column_table2=ColumnTable([[3,4], [7,8]], ["col1","col2"])
138+
139+
result_table=concat_table_chunks([column_table1,column_table2])
140+
141+
assertresult_table.column_table== [[1,2,3,4], [5,6,7,8]]
142+
assertresult_table.column_names== ["col1","col2"]
143+
144+
@pytest.mark.skipif(pyarrowisNone,reason="PyArrow is not installed")
145+
deftest_concat_table_chunks_arrow_table(self):
146+
arrow_table1=pyarrow.Table.from_pydict({"col1": [1,2],"col2": [5,6]})
147+
arrow_table2=pyarrow.Table.from_pydict({"col1": [3,4],"col2": [7,8]})
148+
149+
result_table=concat_table_chunks([arrow_table1,arrow_table2])
150+
assertresult_table.column_names== ["col1","col2"]
151+
assertresult_table.column("col1").to_pylist()== [1,2,3,4]
152+
assertresult_table.column("col2").to_pylist()== [5,6,7,8]
153+
154+
deftest_concat_table_chunks_empty(self):
155+
result_table=concat_table_chunks([])
156+
assertresult_table== []
157+
158+
deftest_concat_table_chunks__incorrect_column_names_error(self):
159+
column_table1=ColumnTable([[1,2], [5,6]], ["col1","col2"])
160+
column_table2=ColumnTable([[3,4], [7,8]], ["col1","col3"])
161+
162+
withpytest.raises(ValueError):
163+
concat_table_chunks([column_table1,column_table2])

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp