Skip to content

Commit

Permalink
Minor fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
jprakash-db committed Sep 20, 2024
1 parent eeaee96 commit 6a8646d
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 17 deletions.
18 changes: 11 additions & 7 deletions src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1227,13 +1227,17 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":

return results

def merge_columnar(self, result1, result2):
"""
Function to merge / combining the columnar results into a single result
:param result1:
:param result2:
:return:
"""
def merge_columnar(self, result1, result2):
"""
Function to merge / combining the columnar results into a single result
:param result1:
:param result2:
:return:
"""

if len(result1) != len(result2):
raise ValueError("The number of columns in both results must be the same")

merged_result = [result1[i] + result2[i] for i in range(len(result1))]
return merged_result

Expand Down
24 changes: 14 additions & 10 deletions src/databricks/sql/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -589,15 +589,17 @@ def convert_decimals_in_arrow_table(table, description) -> "pyarrow.Table":


def convert_to_assigned_datatypes_in_column_table(column_table, description):

converted_column_table = []
for i, col in enumerate(column_table):
if description[i][1] == "decimal":
column_table[i] = tuple(v if v is None else Decimal(v) for v in col)
converted_column_table.append(tuple(v if v is None else Decimal(v) for v in col))
elif description[i][1] == "date":
column_table[i] = tuple(
converted_column_table[i].append(tuple(
v if v is None else datetime.date.fromisoformat(v) for v in col
)
))
elif description[i][1] == "timestamp":
column_table[i] = tuple(
converted_column_table[i].append(tuple(
(
v
if v is None
Expand All @@ -606,9 +608,11 @@ def convert_to_assigned_datatypes_in_column_table(column_table, description):
)
)
for v in col
)
))
else:
converted_column_table.append(col)

return column_table
return converted_column_table


def convert_column_based_set_to_arrow_table(columns, description):
Expand All @@ -624,7 +628,7 @@ def convert_column_based_set_to_arrow_table(columns, description):

def convert_column_based_set_to_column_table(columns, description):
column_names = [c[0] for c in description]
column_table = [_covert_column_to_list(c) for c in columns]
column_table = [_convert_column_to_list(c) for c in columns]

return column_table, column_names

Expand Down Expand Up @@ -653,8 +657,8 @@ def _convert_column_to_arrow_array(t_col):
raise OperationalError("Empty TColumn instance {}".format(t_col))


def _covert_column_to_list(t_col):
supported_field_types = (
def _convert_column_to_list(t_col):
SUPPORTED_FIELD_TYPES = (
"boolVal",
"byteVal",
"i16Val",
Expand All @@ -665,7 +669,7 @@ def _covert_column_to_list(t_col):
"binaryVal",
)

for field in supported_field_types:
for field in SUPPORTED_FIELD_TYPES:
wrapper = getattr(t_col, field)
if wrapper:
return _create_python_tuple(wrapper)
Expand Down

0 comments on commit 6a8646d

Please sign in to comment.