Skip to content

Commit a9aaefc

Browse files
Marcus-Rostintkathole
authored andcommitted
fix: More exhaustive athena types
Signed-off-by: Marcus Rosti <marcus.rosti@baton.io> fix: Double quotes and type Signed-off-by: Marcus Rosti <marcus.rosti@baton.io>
1 parent 1e7134f commit a9aaefc

1 file changed

Lines changed: 40 additions & 3 deletions

File tree

sdk/python/feast/type_map.py

Lines changed: 40 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2198,25 +2198,62 @@ def pg_type_code_to_arrow(code: int) -> str:
21982198

21992199
def 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

22222259
def pa_to_athena_value_type(pa_type: "pyarrow.DataType") -> str:

0 commit comments

Comments
 (0)