Commit61cc398

authored and

committed

[PECOBLR-361] convert column table to arrow if arrow present (#551)

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

1 parentf7d3865 commit61cc398Copy full SHA for 61cc398

File tree

+21

-1

lines changed

+21

-1

lines changed

Lines changed: 14 additions & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -1415,9 +1415,22 @@ def fetchall_arrow(self) -> "pyarrow.Table":`
`1415`	`1415`	`whilenotself.has_been_closed_server_sideandself.has_more_rows:`
`1416`	`1416`	`self._fill_results_buffer()`
`1417`	`1417`	`partial_results=self.results.remaining_rows()`
`1418`		`-results=pyarrow.concat_tables([results,partial_results])`
	`1418`	`+ifisinstance(results,ColumnTable)andisinstance(`
	`1419`	`+partial_results,ColumnTable`
	`1420`	`+ ):`
	`1421`	`+results=self.merge_columnar(results,partial_results)`
	`1422`	`+else:`
	`1423`	`+results=pyarrow.concat_tables([results,partial_results])`
`1419`	`1424`	`self._next_row_index+=partial_results.num_rows`
`1420`	`1425`
	`1426`	`+# If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table`
	`1427`	`+# Valid only for metadata commands result set`
	`1428`	`+ifisinstance(results,ColumnTable)andpyarrow:`
	`1429`	`+data= {`
	`1430`	`+name:col`
	`1431`	`+forname,colinzip(results.column_names,results.column_table)`
	`1432`	`+ }`
	`1433`	`+returnpyarrow.Table.from_pydict(data)`
`1421`	`1434`	`returnresults`
`1422`	`1435`
`1423`	`1436`	`deffetchall_columnar(self):`

Lines changed: 7 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -801,6 +801,13 @@ def test_decimal_not_returned_as_strings_arrow(self):`
`801`	`801`	`decimal_type=arrow_df.field(0).type`
`802`	`802`	`assertpyarrow.types.is_decimal(decimal_type)`
`803`	`803`
	`804`	`+@skipUnless(pysql_supports_arrow(),"arrow test needs arrow support")`
	`805`	`+deftest_catalogs_returns_arrow_table(self):`
	`806`	`+withself.cursor()ascursor:`
	`807`	`+cursor.catalogs()`
	`808`	`+results=cursor.fetchall_arrow()`
	`809`	`+assertisinstance(results,pyarrow.Table)`
	`810`	`+`
`804`	`811`	`deftest_close_connection_closes_cursors(self):`
`805`	`812`
`806`	`813`	`fromdatabricks.sql.thrift_api.TCLIServiceimportttypes`

Comments

(0)