Skip to content

Commit 335fc0c

Browse files
committed
Implemented ColumnQueue to test the fetchall without pyarrow
Removed token removed token
1 parent b438c38 commit 335fc0c

File tree

12 files changed

+227
-7
lines changed

12 files changed

+227
-7
lines changed

.idea/.gitignore

Lines changed: 8 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/codeStyles/Project.xml

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/codeStyles/codeStyleConfig.xml

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/databricks-sql-python.iml

Lines changed: 11 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/misc.xml

Lines changed: 9 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/modules.xml

Lines changed: 8 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/vcs.xml

Lines changed: 6 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

check.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
import os
2+
import sys
3+
# import logging
4+
#
5+
# logging.basicConfig(level=logging.DEBUG)
6+
7+
#
8+
# # Get the parent directory of the current file
9+
# target_folder_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "databricks-sql-python", "src"))
10+
#
11+
# # Add the parent directory to sys.path
12+
# sys.path.append(target_folder_path)
13+
14+
from src.databricks import sql
15+
16+
# from dotenv import load_dotenv
17+
18+
# export DATABRICKS_TOKEN=whatever
19+
20+
21+
# Load environment variables from .env file
22+
# load_dotenv()
23+
24+
host = "e2-dogfood.staging.cloud.databricks.com"
25+
http_path = "/sql/1.0/warehouses/dd43ee29fedd958d"
26+
27+
access_token = ""
28+
connection = sql.connect(
29+
server_hostname=host,
30+
http_path=http_path,
31+
access_token=access_token)
32+
33+
34+
cursor = connection.cursor()
35+
cursor.execute('SELECT :param `p`, * FROM RANGE(10)', {"param": "foo"})
36+
# cursor.execute('SELECT 1')
37+
result = cursor.fetchall()
38+
for row in result:
39+
print(row)
40+
41+
cursor.close()
42+
connection.close()

src/databricks/sql/client.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -777,6 +777,9 @@ def execute(
777777
use_cloud_fetch=self.connection.use_cloud_fetch,
778778
parameters=prepared_params,
779779
)
780+
781+
print("Line 781")
782+
print(execute_response)
780783
self.active_result_set = ResultSet(
781784
self.connection,
782785
execute_response,
@@ -1129,6 +1132,20 @@ def _fill_results_buffer(self):
11291132
self.results = results
11301133
self.has_more_rows = has_more_rows
11311134

1135+
def _convert_columnar_table(self, table):
1136+
column_names = [c[0] for c in self.description]
1137+
ResultRow = Row(*column_names)
1138+
1139+
result = []
1140+
for row_index in range(len(table[0])):
1141+
curr_row = []
1142+
for col_index in range(len(table)-1, -1, -1):
1143+
curr_row.append(table[col_index][row_index])
1144+
result.append(ResultRow(*curr_row))
1145+
1146+
return result
1147+
1148+
11321149
def _convert_arrow_table(self, table):
11331150
column_names = [c[0] for c in self.description]
11341151
ResultRow = Row(*column_names)
@@ -1209,6 +1226,11 @@ def fetchall_arrow(self) -> pyarrow.Table:
12091226

12101227
return results
12111228

1229+
def fetchall_columnar(self):
1230+
results = self.results.remaining_rows()
1231+
self._next_row_index += len(results[0])
1232+
return results
1233+
12121234
def fetchone(self) -> Optional[Row]:
12131235
"""
12141236
Fetch the next row of a query result set, returning a single sequence,
@@ -1224,6 +1246,9 @@ def fetchall(self) -> List[Row]:
12241246
"""
12251247
Fetch all (remaining) rows of a query result, returning them as a list of rows.
12261248
"""
1249+
1250+
return self._convert_columnar_table(self.fetchall_columnar())
1251+
12271252
return self._convert_arrow_table(self.fetchall_arrow())
12281253

12291254
def fetchmany(self, size: int) -> List[Row]:

src/databricks/sql/thrift_api/TCLIService/ttypes.py

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)