@@ -50,60 +50,65 @@ def _convert_decimal(
5050
5151class SqlType :
5252 """
53- SQL type constants
53+ SQL type constants based on Thrift TTypeId values.
5454
55- The list of types can be found in the SEA REST API Reference:
56- https://docs.databricks.com/api/workspace/statementexecution/executestatement
55+ These correspond to the normalized type names that come from the SEA backend
56+ after normalize_sea_type_to_thrift processing (lowercase, without _TYPE suffix).
5757 """
5858
5959 # Numeric types
60- BYTE = "byte"
61- SHORT = "short"
62- INT = "int"
63- LONG = "long"
64- FLOAT = "float"
65- DOUBLE = "double"
66- DECIMAL = "decimal"
60+ TINYINT = "tinyint" # Maps to TTypeId.TINYINT_TYPE
61+ SMALLINT = "smallint" # Maps to TTypeId.SMALLINT_TYPE
62+ INT = "int" # Maps to TTypeId.INT_TYPE
63+ BIGINT = "bigint" # Maps to TTypeId.BIGINT_TYPE
64+ FLOAT = "float" # Maps to TTypeId.FLOAT_TYPE
65+ DOUBLE = "double" # Maps to TTypeId.DOUBLE_TYPE
66+ DECIMAL = "decimal" # Maps to TTypeId.DECIMAL_TYPE
6767
6868 # Boolean type
69- BOOLEAN = "boolean"
69+ BOOLEAN = "boolean" # Maps to TTypeId.BOOLEAN_TYPE
7070
7171 # Date/Time types
72- DATE = "date"
73- TIMESTAMP = "timestamp"
74- INTERVAL = "interval"
72+ DATE = "date" # Maps to TTypeId.DATE_TYPE
73+ TIMESTAMP = "timestamp" # Maps to TTypeId.TIMESTAMP_TYPE
74+ INTERVAL_YEAR_MONTH = (
75+ "interval_year_month" # Maps to TTypeId.INTERVAL_YEAR_MONTH_TYPE
76+ )
77+ INTERVAL_DAY_TIME = "interval_day_time" # Maps to TTypeId.INTERVAL_DAY_TIME_TYPE
7578
7679 # String types
77- CHAR = "char"
78- STRING = "string"
80+ CHAR = "char" # Maps to TTypeId.CHAR_TYPE
81+ VARCHAR = "varchar" # Maps to TTypeId.VARCHAR_TYPE
82+ STRING = "string" # Maps to TTypeId.STRING_TYPE
7983
8084 # Binary type
81- BINARY = "binary"
85+ BINARY = "binary" # Maps to TTypeId.BINARY_TYPE
8286
8387 # Complex types
84- ARRAY = "array"
85- MAP = "map"
86- STRUCT = "struct"
88+ ARRAY = "array" # Maps to TTypeId.ARRAY_TYPE
89+ MAP = "map" # Maps to TTypeId.MAP_TYPE
90+ STRUCT = "struct" # Maps to TTypeId.STRUCT_TYPE
8791
8892 # Other types
89- NULL = "null"
90- USER_DEFINED_TYPE = "user_defined_type"
93+ NULL = "null" # Maps to TTypeId.NULL_TYPE
94+ UNION = "union" # Maps to TTypeId.UNION_TYPE
95+ USER_DEFINED = "user_defined" # Maps to TTypeId.USER_DEFINED_TYPE
9196
9297
9398class SqlTypeConverter :
9499 """
95100 Utility class for converting SQL types to Python types.
96- Based on the types supported by the Databricks SDK .
101+ Based on the Thrift TTypeId types after normalization .
97102 """
98103
99104 # SQL type to conversion function mapping
100105 # TODO: complex types
101106 TYPE_MAPPING : Dict [str , Callable ] = {
102107 # Numeric types
103- SqlType .BYTE : lambda v : int (v ),
104- SqlType .SHORT : lambda v : int (v ),
108+ SqlType .TINYINT : lambda v : int (v ),
109+ SqlType .SMALLINT : lambda v : int (v ),
105110 SqlType .INT : lambda v : int (v ),
106- SqlType .LONG : lambda v : int (v ),
111+ SqlType .BIGINT : lambda v : int (v ),
107112 SqlType .FLOAT : lambda v : float (v ),
108113 SqlType .DOUBLE : lambda v : float (v ),
109114 SqlType .DECIMAL : _convert_decimal ,
@@ -112,30 +117,34 @@ class SqlTypeConverter:
112117 # Date/Time types
113118 SqlType .DATE : lambda v : datetime .date .fromisoformat (v ),
114119 SqlType .TIMESTAMP : lambda v : parser .parse (v ),
115- SqlType .INTERVAL : lambda v : v , # Keep as string for now
120+ SqlType .INTERVAL_YEAR_MONTH : lambda v : v , # Keep as string for now
121+ SqlType .INTERVAL_DAY_TIME : lambda v : v , # Keep as string for now
116122 # String types - no conversion needed
117123 SqlType .CHAR : lambda v : v ,
124+ SqlType .VARCHAR : lambda v : v ,
118125 SqlType .STRING : lambda v : v ,
119126 # Binary type
120127 SqlType .BINARY : lambda v : bytes .fromhex (v ),
121128 # Other types
122129 SqlType .NULL : lambda v : None ,
123130 # Complex types and user-defined types return as-is
124- SqlType .USER_DEFINED_TYPE : lambda v : v ,
131+ SqlType .USER_DEFINED : lambda v : v ,
125132 }
126133
127134 @staticmethod
128135 def convert_value (
129136 value : str ,
130137 sql_type : str ,
138+ column_name : Optional [str ],
131139 ** kwargs ,
132140 ) -> object :
133141 """
134142 Convert a string value to the appropriate Python type based on SQL type.
135143
136144 Args:
137145 value: The string value to convert
138- sql_type: The SQL type (e.g., 'int', 'decimal')
146+ sql_type: The SQL type (e.g., 'tinyint', 'decimal')
147+ column_name: The name of the column being converted
139148 **kwargs: Additional keyword arguments for the conversion function
140149
141150 Returns:
@@ -155,6 +164,10 @@ def convert_value(
155164 return converter_func (value , precision , scale )
156165 else :
157166 return converter_func (value )
158- except (ValueError , TypeError , decimal .InvalidOperation ) as e :
159- logger .warning (f"Error converting value '{ value } { sql_type } { e } )
167+ except Exception as e :
168+ warning_message = f"Error converting value '{ value } { sql_type }
169+ if column_name :
170+ warning_message += f" in column { column_name }
171+ warning_message += f": { e }
172+ logger .warning (warning_message )
160173 return value
0 commit comments