@@ -2198,25 +2198,62 @@ def pg_type_code_to_arrow(code: int) -> str:
21982198
21992199def athena_to_feast_value_type (athena_type_as_str : str ) -> ValueType :
22002200 # Type names from https://docs.aws.amazon.com/athena/latest/ug/data-types.html
2201+ athena_type = athena_type_as_str .lower ().strip ()
2202+ if athena_type .startswith ("array" ):
2203+ inner_type_match = re .search (r"(?:<|\[)(.+)(?:>|\])" , athena_type )
2204+ if inner_type_match :
2205+ inner_type = inner_type_match .group (1 ).strip ()
2206+ inner_feast_type = athena_to_feast_value_type (inner_type )
2207+
2208+ list_mapping = {
2209+ ValueType .BYTES : ValueType .BYTES_LIST ,
2210+ ValueType .STRING : ValueType .STRING_LIST ,
2211+ ValueType .INT32 : ValueType .INT32_LIST ,
2212+ ValueType .INT64 : ValueType .INT64_LIST ,
2213+ ValueType .DOUBLE : ValueType .DOUBLE_LIST ,
2214+ ValueType .FLOAT : ValueType .FLOAT_LIST ,
2215+ ValueType .BOOL : ValueType .BOOL_LIST ,
2216+ ValueType .UNIX_TIMESTAMP : ValueType .UNIX_TIMESTAMP_LIST ,
2217+ ValueType .MAP : ValueType .MAP_LIST ,
2218+ ValueType .JSON : ValueType .JSON_LIST ,
2219+ ValueType .STRUCT : ValueType .STRUCT_LIST ,
2220+ ValueType .UUID : ValueType .UUID_LIST ,
2221+ ValueType .DECIMAL : ValueType .DECIMAL_LIST ,
2222+ }
2223+ return list_mapping .get (inner_feast_type , ValueType .VALUE_LIST )
2224+ return ValueType .VALUE_LIST
2225+
2226+ base_type = re .split (r"[(<\[]" , athena_type )[0 ].strip ()
2227+
2228+ if "timestamp" in base_type or "time" in base_type or "date" in base_type :
2229+ return ValueType .UNIX_TIMESTAMP
2230+
22012231 type_map = {
2202- "null" : ValueType .UNKNOWN ,
2232+ "null" : ValueType .UNKNOWN , # There is a null type, but this preserves backwards compat
22032233 "boolean" : ValueType .BOOL ,
22042234 "tinyint" : ValueType .INT32 ,
22052235 "smallint" : ValueType .INT32 ,
22062236 "int" : ValueType .INT32 ,
2237+ "integer" : ValueType .INT32 ,
22072238 "bigint" : ValueType .INT64 ,
22082239 "double" : ValueType .DOUBLE ,
22092240 "float" : ValueType .FLOAT ,
2241+ "real" : ValueType .FLOAT ,
2242+ "decimal" : ValueType .DECIMAL ,
22102243 "binary" : ValueType .BYTES ,
2244+ "varbinary" : ValueType .BYTES ,
22112245 "char" : ValueType .STRING ,
22122246 "varchar" : ValueType .STRING ,
22132247 "string" : ValueType .STRING ,
2214- "timestamp" : ValueType .UNIX_TIMESTAMP ,
22152248 "json" : ValueType .JSON ,
22162249 "struct" : ValueType .STRUCT ,
2250+ "row" : ValueType .STRUCT ,
22172251 "map" : ValueType .MAP ,
2252+ "uuid" : ValueType .UUID ,
2253+ "ipaddress" : ValueType .STRING ,
22182254 }
2219- return type_map [athena_type_as_str .lower ()]
2255+
2256+ return type_map .get (base_type , ValueType .UNKNOWN )
22202257
22212258
22222259def pa_to_athena_value_type (pa_type : "pyarrow.DataType" ) -> str :
0 commit comments