Skip to content

Commit fa734cc

Browse files
re-introduce preliminary complex types support
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 75c330a commit fa734cc

File tree

1 file changed

+42
-0
lines changed

1 file changed

+42
-0
lines changed

src/databricks/sql/backend/sea/result_set.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import json
34
from typing import Any, List, Optional, TYPE_CHECKING
45

56
import logging
@@ -82,6 +83,41 @@ def __init__(
8283
arrow_schema_bytes=execute_response.arrow_schema_bytes,
8384
)
8485

86+
def _convert_complex_types_to_string(
87+
self, rows: "pyarrow.Table"
88+
) -> "pyarrow.Table":
89+
"""
90+
Convert complex types (array, struct, map) to string representation.
91+
Args:
92+
rows: Input PyArrow table
93+
Returns:
94+
PyArrow table with complex types converted to strings
95+
"""
96+
97+
if not pyarrow:
98+
return rows
99+
100+
def convert_complex_column_to_string(col: "pyarrow.Array") -> "pyarrow.Array":
101+
python_values = col.to_pylist()
102+
json_strings = [
103+
(None if val is None else json.dumps(val)) for val in python_values
104+
]
105+
return pyarrow.array(json_strings, type=pyarrow.string())
106+
107+
converted_columns = []
108+
for col in rows.columns:
109+
converted_col = col
110+
if (
111+
pyarrow.types.is_list(col.type)
112+
or pyarrow.types.is_large_list(col.type)
113+
or pyarrow.types.is_struct(col.type)
114+
or pyarrow.types.is_map(col.type)
115+
):
116+
converted_col = convert_complex_column_to_string(col)
117+
converted_columns.append(converted_col)
118+
119+
return pyarrow.Table.from_arrays(converted_columns, names=rows.column_names)
120+
85121
def _convert_json_types(self, row: List[str]) -> List[Any]:
86122
"""
87123
Convert string values in the row to appropriate Python types based on column metadata.
@@ -200,6 +236,9 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
200236
if isinstance(self.results, JsonQueue):
201237
results = self._convert_json_to_arrow_table(results)
202238

239+
if not self.backend._use_arrow_native_complex_types:
240+
results = self._convert_complex_types_to_string(results)
241+
203242
self._next_row_index += results.num_rows
204243

205244
return results
@@ -213,6 +252,9 @@ def fetchall_arrow(self) -> "pyarrow.Table":
213252
if isinstance(self.results, JsonQueue):
214253
results = self._convert_json_to_arrow_table(results)
215254

255+
if not self.backend._use_arrow_native_complex_types:
256+
results = self._convert_complex_types_to_string(results)
257+
216258
self._next_row_index += results.num_rows
217259

218260
return results

0 commit comments

Comments
 (0)