Skip to content

Commit dd7abb8

Browse files
introduce preliminary complex type conversion
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent b3273c7 commit dd7abb8

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
@@ -600,6 +600,43 @@ def fetchall_json(self) -> List:
600600

601601
return results
602602

603+
def _convert_complex_types_to_string(
604+
self, rows: "pyarrow.Table"
605+
) -> "pyarrow.Table":
606+
"""
607+
Convert complex types (array, struct, map) to string representation.
608+
609+
Args:
610+
rows: Input PyArrow table
611+
612+
Returns:
613+
PyArrow table with complex types converted to strings
614+
"""
615+
616+
if not pyarrow:
617+
return rows
618+
619+
def convert_complex_column_to_string(col: "pyarrow.Array") -> "pyarrow.Array":
620+
python_values = col.to_pylist()
621+
json_strings = [
622+
(None if val is None else json.dumps(val)) for val in python_values
623+
]
624+
return pyarrow.array(json_strings, type=pyarrow.string())
625+
626+
converted_columns = []
627+
for col in rows.columns:
628+
converted_col = col
629+
if (
630+
pyarrow.types.is_list(col.type)
631+
or pyarrow.types.is_large_list(col.type)
632+
or pyarrow.types.is_struct(col.type)
633+
or pyarrow.types.is_map(col.type)
634+
):
635+
converted_col = convert_complex_column_to_string(col)
636+
converted_columns.append(converted_col)
637+
638+
return pyarrow.Table.from_arrays(converted_columns, names=rows.column_names)
639+
603640
def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
604641
"""
605642
Fetch the next set of rows as an Arrow table.
@@ -625,6 +662,9 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
625662
results = self._convert_json_to_arrow(rows)
626663
self._next_row_index += results.num_rows
627664

665+
if not self.backend._use_arrow_native_complex_types:
666+
results = self._convert_complex_types_to_string(results)
667+
628668
return results
629669

630670
def fetchall_arrow(self) -> "pyarrow.Table":
@@ -639,6 +679,9 @@ def fetchall_arrow(self) -> "pyarrow.Table":
639679
results = self._convert_json_to_arrow(rows)
640680
self._next_row_index += results.num_rows
641681

682+
if not self.backend._use_arrow_native_complex_types:
683+
results = self._convert_complex_types_to_string(results)
684+
642685
return results
643686

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

0 commit comments

Comments
 (0)