convert complex types to string if not _use_arrow_native_complex_types

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -11,6 +11,8 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING
 
+from databricks.sql.types import SSLOptions
+
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
 
@@ -25,6 +27,13 @@
 
 
 class DatabricksClient(ABC):
+    def __init__(self, ssl_options: SSLOptions, **kwargs):
+        self._use_arrow_native_complex_types = kwargs.get(
+            "_use_arrow_native_complex_types", True
+        )
+        self._max_download_threads = kwargs.get("max_download_threads", 10)
+        self._ssl_options = ssl_options
+
     # == Connection and Session Management ==
     @abstractmethod
     def open_session(

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

@@ -124,7 +124,7 @@ def __init__(
             http_path,
         )
 
-        self._max_download_threads = kwargs.get("max_download_threads", 10)
+        super().__init__(ssl_options, **kwargs)
 
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
@@ -136,7 +136,7 @@ def __init__(
             http_path=http_path,
             http_headers=http_headers,
             auth_provider=auth_provider,
-            ssl_options=ssl_options,
+            ssl_options=self._ssl_options,
             **kwargs,
         )
 

src/databricks/sql/backend/thrift_backend.py

@@ -147,6 +147,8 @@ def __init__(
             http_path,
         )
 
+        super().__init__(ssl_options, **kwargs)
+
         port = port or 443
         if kwargs.get("_connection_uri"):
             uri = kwargs.get("_connection_uri")
@@ -160,19 +162,13 @@ def __init__(
             raise ValueError("No valid connection settings.")
 
         self._initialize_retry_args(kwargs)
-        self._use_arrow_native_complex_types = kwargs.get(
-            "_use_arrow_native_complex_types", True
-        )
+
         self._use_arrow_native_decimals = kwargs.get("_use_arrow_native_decimals", True)
         self._use_arrow_native_timestamps = kwargs.get(
             "_use_arrow_native_timestamps", True
         )
 
         # Cloud fetch
-        self._max_download_threads = kwargs.get("max_download_threads", 10)
-
-        self._ssl_options = ssl_options
-
         self._auth_provider = auth_provider
 
         # Connector version 3 retry approach

src/databricks/sql/result_set.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+import json
 from typing import List, Optional, Any, Union, Tuple, TYPE_CHECKING
 
 import logging
@@ -551,6 +552,43 @@ def fetchall_json(self):
 
         return results
 
+    def _convert_complex_types_to_string(
+        self, rows: "pyarrow.Table"
+    ) -> "pyarrow.Table":
+        """
+        Convert complex types (array, struct, map) to string representation.
+
+        Args:
+            rows: Input PyArrow table
+
+        Returns:
+            PyArrow table with complex types converted to strings
+        """
+
+        if not pyarrow:
+            return rows
+
+        def convert_complex_column_to_string(col: "pyarrow.Array") -> "pyarrow.Array":
+            python_values = col.to_pylist()
+            json_strings = [
+                (None if val is None else json.dumps(val)) for val in python_values
+            ]
+            return pyarrow.array(json_strings, type=pyarrow.string())
+
+        converted_columns = []
+        for col in rows.columns:
+            converted_col = col
+            if (
+                pyarrow.types.is_list(col.type)
+                or pyarrow.types.is_large_list(col.type)
+                or pyarrow.types.is_struct(col.type)
+                or pyarrow.types.is_map(col.type)
+            ):
+                converted_col = convert_complex_column_to_string(col)
+            converted_columns.append(converted_col)
+
+        return pyarrow.Table.from_arrays(converted_columns, names=rows.column_names)
+
     def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """
         Fetch the next set of rows as an Arrow table.
@@ -571,6 +609,9 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         results = self.results.next_n_rows(size)
         self._next_row_index += results.num_rows
 
+        if not self.backend._use_arrow_native_complex_types:
+            results = self._convert_complex_types_to_string(results)
+
         return results
 
     def fetchall_arrow(self) -> "pyarrow.Table":
@@ -580,6 +621,9 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
 
+        if not self.backend._use_arrow_native_complex_types:
+            results = self._convert_complex_types_to_string(results)
+
         return results
 
     def fetchone(self) -> Optional[Row]: