Commit4cb0d70

authored and

committed

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

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

1 parentbb69dc9 commit4cb0d70Copy full SHA for 4cb0d70

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
`@@ -1433,9 +1433,22 @@ def fetchall_arrow(self) -> "pyarrow.Table":`
`1433`	`1433`	`whilenotself.has_been_closed_server_sideandself.has_more_rows:`
`1434`	`1434`	`self._fill_results_buffer()`
`1435`	`1435`	`partial_results=self.results.remaining_rows()`
`1436`		`-results=pyarrow.concat_tables([results,partial_results])`
	`1436`	`+ifisinstance(results,ColumnTable)andisinstance(`
	`1437`	`+partial_results,ColumnTable`
	`1438`	`+ ):`
	`1439`	`+results=self.merge_columnar(results,partial_results)`
	`1440`	`+else:`
	`1441`	`+results=pyarrow.concat_tables([results,partial_results])`
`1437`	`1442`	`self._next_row_index+=partial_results.num_rows`
`1438`	`1443`
	`1444`	`+# If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table`
	`1445`	`+# Valid only for metadata commands result set`
	`1446`	`+ifisinstance(results,ColumnTable)andpyarrow:`
	`1447`	`+data= {`
	`1448`	`+name:col`
	`1449`	`+forname,colinzip(results.column_names,results.column_table)`
	`1450`	`+ }`
	`1451`	`+returnpyarrow.Table.from_pydict(data)`
`1439`	`1452`	`returnresults`
`1440`	`1453`
`1441`	`1454`	`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)