Skip to content

Commit b99d0c4

Browse files
Revert "remove complex types code"
This reverts commit 138359d.
1 parent 138359d commit b99d0c4

File tree

1 file changed

+43
-0
lines changed

1 file changed

+43
-0
lines changed

src/databricks/sql/result_set.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -552,6 +552,43 @@ def fetchall_json(self):
552552

553553
return results
554554

555+
def _convert_complex_types_to_string(
556+
self, rows: "pyarrow.Table"
557+
) -> "pyarrow.Table":
558+
"""
559+
Convert complex types (array, struct, map) to string representation.
560+
561+
Args:
562+
rows: Input PyArrow table
563+
564+
Returns:
565+
PyArrow table with complex types converted to strings
566+
"""
567+
568+
if not pyarrow:
569+
return rows
570+
571+
def convert_complex_column_to_string(col: "pyarrow.Array") -> "pyarrow.Array":
572+
python_values = col.to_pylist()
573+
json_strings = [
574+
(None if val is None else json.dumps(val)) for val in python_values
575+
]
576+
return pyarrow.array(json_strings, type=pyarrow.string())
577+
578+
converted_columns = []
579+
for col in rows.columns:
580+
converted_col = col
581+
if (
582+
pyarrow.types.is_list(col.type)
583+
or pyarrow.types.is_large_list(col.type)
584+
or pyarrow.types.is_struct(col.type)
585+
or pyarrow.types.is_map(col.type)
586+
):
587+
converted_col = convert_complex_column_to_string(col)
588+
converted_columns.append(converted_col)
589+
590+
return pyarrow.Table.from_arrays(converted_columns, names=rows.column_names)
591+
555592
def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
556593
"""
557594
Fetch the next set of rows as an Arrow table.
@@ -572,6 +609,9 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
572609
results = self.results.next_n_rows(size)
573610
self._next_row_index += results.num_rows
574611

612+
if not self.backend._use_arrow_native_complex_types:
613+
results = self._convert_complex_types_to_string(results)
614+
575615
return results
576616

577617
def fetchall_arrow(self) -> "pyarrow.Table":
@@ -581,6 +621,9 @@ def fetchall_arrow(self) -> "pyarrow.Table":
581621
results = self.results.remaining_rows()
582622
self._next_row_index += results.num_rows
583623

624+
if not self.backend._use_arrow_native_complex_types:
625+
results = self._convert_complex_types_to_string(results)
626+
584627
return results
585628

586629
def fetchone(self) -> Optional[Row]:

0 commit comments

Comments
 (0)