variant type detection

Author: shivam2680

src/databricks/sql/backend/thrift_backend.py

@@ -735,7 +735,7 @@ def convert_col(t_column_desc):
         return pyarrow.schema([convert_col(col) for col in t_table_schema.columns])
 
     @staticmethod
-    def _col_to_description(col, field, session_id_hex=None):
+    def _col_to_description(col, field=None, session_id_hex=None):
         type_entry = col.typeDesc.types[0]
 
         if type_entry.primitiveEntry:
@@ -778,7 +778,9 @@ def _col_to_description(col, field, session_id_hex=None):
         return col.columnName, cleaned_type, None, None, precision, scale, None
 
     @staticmethod
-    def _hive_schema_to_description(t_table_schema, schema_bytes=None, session_id_hex=None):
+    def _hive_schema_to_description(
+        t_table_schema, schema_bytes=None, session_id_hex=None
+    ):
         field_dict = {}
         if pyarrow and schema_bytes:
             try:
@@ -788,8 +790,13 @@ def _hive_schema_to_description(t_table_schema, schema_bytes=None, session_id_he
                     field_dict[field.name] = field
             except Exception as e:
                 logger.debug(f"Could not parse arrow schema: {e}")
+
         return [
-            ThriftDatabricksClient._col_to_description(col, field_dict.get(col.columnName), session_id_hex)
+            ThriftDatabricksClient._col_to_description(
+                col,
+                field_dict.get(col.columnName) if field_dict else None,
+                session_id_hex,
+            )
             for col in t_table_schema.columns
         ]
 
@@ -822,11 +829,6 @@ def _results_message_to_execute_response(self, resp, operation_state):
             or direct_results.resultSet.hasMoreRows
         )
 
-        description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema,
-            self._session_id_hex,
-        )
-
         if pyarrow:
             schema_bytes = (
                 t_result_set_metadata_resp.arrowSchema
@@ -840,7 +842,9 @@ def _results_message_to_execute_response(self, resp, operation_state):
             schema_bytes = None
 
         description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema, schema_bytes
+            t_result_set_metadata_resp.schema,
+            schema_bytes,
+            self._session_id_hex,
         )
 
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
@@ -887,11 +891,6 @@ def get_execution_result(
 
         t_result_set_metadata_resp = resp.resultSetMetadata
 
-        description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema,
-            self._session_id_hex,
-        )
-
         if pyarrow:
             schema_bytes = (
                 t_result_set_metadata_resp.arrowSchema
@@ -904,6 +903,12 @@ def get_execution_result(
         else:
             schema_bytes = None
 
+        description = self._hive_schema_to_description(
+            t_result_set_metadata_resp.schema,
+            schema_bytes,
+            self._session_id_hex,
+        )
+
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
         is_staging_operation = t_result_set_metadata_resp.isStagingOperation
         has_more_rows = resp.hasMoreRows

tests/unit/test_thrift_backend.py

@@ -2370,7 +2370,7 @@ def test_col_to_description_with_variant_type(self):
             metadata={b'Spark:DataType:SqlName': b'VARIANT'}
         )
 
-        result = ThriftBackend._col_to_description(col, field)
+        result = ThriftDatabricksClient._col_to_description(col, field)
 
         # Verify the result has variant as the type
         self.assertEqual(result[0], "variant_col")  # Column name
@@ -2395,7 +2395,7 @@ def test_col_to_description_without_variant_type(self):
             metadata={}
         )
 
-        result = ThriftBackend._col_to_description(col, field)
+        result = ThriftDatabricksClient._col_to_description(col, field)
 
         # Verify the result has string as the type (unchanged)
         self.assertEqual(result[0], "normal_col")  # Column name
@@ -2414,7 +2414,7 @@ def test_col_to_description_with_null_field(self):
         )
 
         # Pass None as the field
-        result = ThriftBackend._col_to_description(col, None)
+        result = ThriftDatabricksClient._col_to_description(col, None)
 
         # Verify the result has string as the type (unchanged)
         self.assertEqual(result[0], "missing_field")  # Column name
@@ -2452,7 +2452,7 @@ def test_hive_schema_to_description_with_arrow_schema(self):
         schema_bytes = arrow_schema.serialize().to_pybytes()
 
         # Get the description
-        description = ThriftBackend._hive_schema_to_description(t_table_schema, schema_bytes)
+        description = ThriftDatabricksClient._hive_schema_to_description(t_table_schema, schema_bytes)
 
         # Verify regular column type
         self.assertEqual(description[0][0], "regular_col")
@@ -2473,7 +2473,7 @@ def test_hive_schema_to_description_with_null_schema_bytes(self):
         t_table_schema = ttypes.TTableSchema(columns=columns)
 
         # Get the description with null schema_bytes
-        description = ThriftBackend._hive_schema_to_description(t_table_schema, None)
+        description = ThriftDatabricksClient._hive_schema_to_description(t_table_schema, None)
 
         # Verify column type remains unchanged
         self.assertEqual(description[0][0], "regular_col")
@@ -2493,7 +2493,7 @@ def test_col_to_description_with_malformed_metadata(self):
             metadata={b'Spark:DataType:SqlName': b'Some unexpected value'}
         )
 
-        result = ThriftBackend._col_to_description(col, field)
+        result = ThriftDatabricksClient._col_to_description(col, field)
 
         # Verify the type remains unchanged
         self.assertEqual(result[0], "weird_field")  # Column name