Skip to content

Commit 326554d

Browse files
falloficarussntkathole
falloficaruss
authored and
ntkathole
committed
fix: Handle parameterized and complex Trino types in type map
Signed-off-by: Abhishek Shinde <norizzabhii@gmail.com>
1 parent a42dc85 commit 326554d

2 files changed

Lines changed: 253 additions & 20 deletions

File tree

‎sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_type_map.py‎

Lines changed: 49 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@ def trino_to_feast_value_type(trino_type_as_str: str) -> ValueType:
2222
"boolean": ValueType.BOOL,
2323
"real": ValueType.FLOAT,
2424
"date": ValueType.STRING,
25+
"binary": ValueType.STRING,
26+
"varbinary": ValueType.STRING,
27+
"json": ValueType.STRING,
2528
}
2629
_trino_type_as_str: str = trino_type_as_str
2730
trino_type_as_str = trino_type_as_str.lower()
@@ -36,13 +39,18 @@ def trino_to_feast_value_type(trino_type_as_str: str) -> ValueType:
3639
trino_type_as_str = "decimal64"
3740
else:
3841
trino_type_as_str = "decimal32"
42+
else:
43+
trino_type_as_str = "decimal64"
3944

4045
elif trino_type_as_str.startswith("timestamp"):
4146
trino_type_as_str = "timestamp"
4247

4348
elif trino_type_as_str.startswith("varchar"):
4449
trino_type_as_str = "varchar"
4550

51+
elif trino_type_as_str.startswith("char"):
52+
trino_type_as_str = "char"
53+
4654
if trino_type_as_str not in type_map:
4755
raise ValueError(f"Trino type not supported by feast {_trino_type_as_str}")
4856
return type_map[trino_type_as_str]
@@ -55,7 +63,11 @@ def pa_to_trino_value_type(pa_type_as_str: str) -> str:
5563
trino_type = "{}"
5664
if pa_type_as_str.startswith("list"):
5765
trino_type = "array<{}>"
58-
pa_type_as_str = re.search(r"^list<item:\s(.+)>$", pa_type_as_str).group(1)
66+
match = re.search(r"^list<item:\s(.+)>$", pa_type_as_str)
67+
if match:
68+
pa_type_as_str = match.group(1)
69+
else:
70+
return trino_type.format("varchar")
5971

6072
if pa_type_as_str.startswith("date"):
6173
return trino_type.format("date")
@@ -67,7 +79,10 @@ def pa_to_trino_value_type(pa_type_as_str: str) -> str:
6779
return trino_type.format("timestamp")
6880

6981
if pa_type_as_str.startswith("decimal"):
70-
return trino_type.format(pa_type_as_str)
82+
# PyArrow renders decimal types as decimal128(10, 2) or decimal256(10, 2),
83+
# but Trino expects just decimal(10, 2)
84+
normalized = re.sub(r"^decimal\d+", "decimal", pa_type_as_str)
85+
return trino_type.format(normalized)
7186

7287
if pa_type_as_str.startswith("map<"):
7388
return trino_type.format("varchar")
@@ -92,33 +107,43 @@ def pa_to_trino_value_type(pa_type_as_str: str) -> str:
92107
"float": "double",
93108
"double": "double",
94109
"binary": "binary",
110+
"varbinary": "binary",
95111
"string": "varchar",
112+
"char": "varchar",
96113
}
97114
return trino_type.format(type_map[pa_type_as_str])
98115

99116

100-
_TRINO_TO_PA_TYPE_MAP = {
117+
_TRINO_TO_PA_TYPE_MAP: Dict[str, pa.DataType] = {
101118
"null": pa.null(),
102119
"boolean": pa.bool_(),
103120
"date": pa.date32(),
104121
"tinyint": pa.int8(),
105122
"smallint": pa.int16(),
106123
"integer": pa.int32(),
124+
"int": pa.int32(),
107125
"bigint": pa.int64(),
108126
"double": pa.float64(),
109127
"binary": pa.binary(),
128+
"varbinary": pa.binary(),
110129
"char": pa.string(),
130+
"json": pa.string(),
111131
"real": pa.float32(),
112132
}
113133

114134

135+
def _trino_array_item_type(trino_type_as_str: str) -> str | None:
136+
if trino_type_as_str.startswith("array(") and trino_type_as_str.endswith(")"):
137+
return trino_type_as_str[6:-1].strip()
138+
return None
139+
140+
115141
def trino_to_pa_value_type(trino_type_as_str: str) -> pa.DataType:
116-
trino_type_as_str = trino_type_as_str.lower()
142+
trino_type_as_str = trino_type_as_str.lower().strip()
117143

118-
_is_list: bool = False
119-
if trino_type_as_str.startswith("array"):
120-
_is_list = True
121-
trino_type_as_str = re.search(r"^array\((\w+)\)$", trino_type_as_str).group(1)
144+
array_item_type = _trino_array_item_type(trino_type_as_str)
145+
if array_item_type is not None:
146+
return pa.list_(trino_to_pa_value_type(array_item_type))
122147

123148
if trino_type_as_str.startswith("decimal"):
124149
search_precision = re.search(
@@ -127,20 +152,24 @@ def trino_to_pa_value_type(trino_type_as_str: str) -> pa.DataType:
127152
if search_precision:
128153
precision = int(search_precision.group(1))
129154
if precision > 32:
130-
pa_type = pa.float64()
155+
return pa.float64()
131156
else:
132-
pa_type = pa.float32()
157+
return pa.float32()
158+
return pa.float64()
133159

134-
elif trino_type_as_str.startswith("timestamp"):
135-
pa_type = pa.timestamp("us")
160+
if trino_type_as_str.startswith("timestamp"):
161+
return pa.timestamp("us")
136162

137-
elif trino_type_as_str.startswith("varchar"):
138-
pa_type = pa.string()
163+
if trino_type_as_str.startswith("varchar"):
164+
return pa.string()
165+
166+
if trino_type_as_str.startswith("char"):
167+
return pa.string()
168+
169+
if trino_type_as_str.startswith("row("):
170+
return pa.string()
139171

140-
else:
141-
pa_type = _TRINO_TO_PA_TYPE_MAP[trino_type_as_str]
172+
if trino_type_as_str.startswith("map("):
173+
return pa.string()
142174

143-
if _is_list:
144-
return pa.list_(pa_type)
145-
else:
146-
return pa_type
175+
return _TRINO_TO_PA_TYPE_MAP[trino_type_as_str]

‎sdk/python/tests/unit/infra/offline_stores/contrib/trino_offline_store/test_trino_type_map.py‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
import pyarrow as pa
2+
import pytest
3+
4+
from feast import ValueType
5+
from feast.infra.offline_stores.contrib.trino_offline_store.trino_type_map import (
6+
_trino_array_item_type,
7+
pa_to_trino_value_type,
8+
trino_to_feast_value_type,
9+
trino_to_pa_value_type,
10+
)
11+
12+
13+
class TestTrinoArrayItemType:
14+
def test_simple_type(self) -> None:
15+
assert _trino_array_item_type("array(bigint)") == "bigint"
16+
17+
def test_parameterized_type(self) -> None:
18+
assert _trino_array_item_type("array(varchar(10))") == "varchar(10)"
19+
20+
def test_parameterized_with_comma(self) -> None:
21+
assert _trino_array_item_type("array(decimal(10, 2))") == "decimal(10, 2)"
22+
23+
def test_nested_array(self) -> None:
24+
assert _trino_array_item_type("array(array(varchar))") == "array(varchar)"
25+
26+
def test_complex_row(self) -> None:
27+
assert (
28+
_trino_array_item_type("array(row(x bigint, y varchar(10)))")
29+
== "row(x bigint, y varchar(10))"
30+
)
31+
32+
def test_not_an_array(self) -> None:
33+
assert _trino_array_item_type("varchar") is None
34+
35+
def test_partial_prefix(self) -> None:
36+
assert _trino_array_item_type("array") is None
37+
assert _trino_array_item_type("array(") is None
38+
39+
40+
class TestTrinoToFeastValueType:
41+
def test_simple_types(self) -> None:
42+
assert trino_to_feast_value_type("boolean") == ValueType.BOOL
43+
assert trino_to_feast_value_type("bigint") == ValueType.INT64
44+
assert trino_to_feast_value_type("integer") == ValueType.INT32
45+
assert trino_to_feast_value_type("int") == ValueType.INT32
46+
assert trino_to_feast_value_type("double") == ValueType.DOUBLE
47+
assert trino_to_feast_value_type("real") == ValueType.FLOAT
48+
assert trino_to_feast_value_type("date") == ValueType.STRING
49+
assert trino_to_feast_value_type("tinyint") == ValueType.INT32
50+
assert trino_to_feast_value_type("smallint") == ValueType.INT32
51+
52+
def test_parameterized_varchar(self) -> None:
53+
assert trino_to_feast_value_type("varchar(10)") == ValueType.STRING
54+
55+
def test_parameterized_char(self) -> None:
56+
assert trino_to_feast_value_type("char(10)") == ValueType.STRING
57+
assert trino_to_feast_value_type("char") == ValueType.STRING
58+
59+
def test_timestamp_with_precision(self) -> None:
60+
assert trino_to_feast_value_type("timestamp(3)") == ValueType.UNIX_TIMESTAMP
61+
assert trino_to_feast_value_type("timestamp") == ValueType.UNIX_TIMESTAMP
62+
63+
def test_decimal_with_precision(self) -> None:
64+
assert trino_to_feast_value_type("decimal(10, 2)") == ValueType.FLOAT
65+
assert trino_to_feast_value_type("decimal(38, 2)") == ValueType.DOUBLE
66+
assert trino_to_feast_value_type("decimal(32)") == ValueType.FLOAT
67+
assert trino_to_feast_value_type("decimal(33)") == ValueType.DOUBLE
68+
69+
def test_bare_decimal(self) -> None:
70+
assert trino_to_feast_value_type("decimal") == ValueType.DOUBLE
71+
72+
def test_binary_types(self) -> None:
73+
assert trino_to_feast_value_type("binary") == ValueType.STRING
74+
assert trino_to_feast_value_type("varbinary") == ValueType.STRING
75+
76+
def test_json(self) -> None:
77+
assert trino_to_feast_value_type("json") == ValueType.STRING
78+
79+
def test_unsupported_type(self) -> None:
80+
with pytest.raises(ValueError, match="Trino type not supported"):
81+
trino_to_feast_value_type("unknown_type")
82+
83+
84+
class TestTrinoToPaValueType:
85+
def test_simple_types(self) -> None:
86+
assert trino_to_pa_value_type("boolean") == pa.bool_()
87+
assert trino_to_pa_value_type("bigint") == pa.int64()
88+
assert trino_to_pa_value_type("integer") == pa.int32()
89+
assert trino_to_pa_value_type("int") == pa.int32()
90+
assert trino_to_pa_value_type("double") == pa.float64()
91+
assert trino_to_pa_value_type("real") == pa.float32()
92+
assert trino_to_pa_value_type("date") == pa.date32()
93+
assert trino_to_pa_value_type("tinyint") == pa.int8()
94+
assert trino_to_pa_value_type("smallint") == pa.int16()
95+
96+
def test_parameterized_varchar(self) -> None:
97+
assert trino_to_pa_value_type("varchar(10)") == pa.string()
98+
99+
def test_parameterized_char(self) -> None:
100+
assert trino_to_pa_value_type("char(10)") == pa.string()
101+
assert trino_to_pa_value_type("char") == pa.string()
102+
103+
def test_binary_types(self) -> None:
104+
assert trino_to_pa_value_type("binary") == pa.binary()
105+
assert trino_to_pa_value_type("varbinary") == pa.binary()
106+
107+
def test_json(self) -> None:
108+
assert trino_to_pa_value_type("json") == pa.string()
109+
110+
def test_timestamp(self) -> None:
111+
assert trino_to_pa_value_type("timestamp") == pa.timestamp("us")
112+
assert trino_to_pa_value_type("timestamp(3)") == pa.timestamp("us")
113+
114+
def test_decimal_bare(self) -> None:
115+
assert trino_to_pa_value_type("decimal") == pa.float64()
116+
117+
def test_decimal_with_precision(self) -> None:
118+
assert trino_to_pa_value_type("decimal(10, 2)") == pa.float32()
119+
assert trino_to_pa_value_type("decimal(38, 2)") == pa.float64()
120+
assert trino_to_pa_value_type("decimal(32)") == pa.float32()
121+
assert trino_to_pa_value_type("decimal(33)") == pa.float64()
122+
123+
def test_array_simple(self) -> None:
124+
assert trino_to_pa_value_type("array(bigint)") == pa.list_(pa.int64())
125+
126+
def test_array_parameterized_varchar(self) -> None:
127+
assert trino_to_pa_value_type("array(varchar(10))") == pa.list_(pa.string())
128+
129+
def test_array_parameterized_decimal(self) -> None:
130+
assert trino_to_pa_value_type("array(decimal(10, 2))") == pa.list_(pa.float32())
131+
132+
def test_array_nested(self) -> None:
133+
assert trino_to_pa_value_type("array(array(bigint))") == pa.list_(
134+
pa.list_(pa.int64())
135+
)
136+
137+
def test_row_type(self) -> None:
138+
assert trino_to_pa_value_type("row(x bigint)") == pa.string()
139+
assert trino_to_pa_value_type("row(x bigint, y varchar)") == pa.string()
140+
141+
def test_map_type(self) -> None:
142+
assert trino_to_pa_value_type("map(varchar, bigint)") == pa.string()
143+
144+
def test_array_of_row(self) -> None:
145+
assert trino_to_pa_value_type(
146+
"array(row(x bigint, y varchar(10)))"
147+
) == pa.list_(pa.string())
148+
149+
def test_unsupported_type(self) -> None:
150+
with pytest.raises(KeyError):
151+
trino_to_pa_value_type("unknown_type")
152+
153+
154+
class TestPaToTrinoValueType:
155+
def test_simple_types(self) -> None:
156+
assert pa_to_trino_value_type(str(pa.bool_())) == "boolean"
157+
assert pa_to_trino_value_type(str(pa.int8())) == "tinyint"
158+
assert pa_to_trino_value_type(str(pa.int16())) == "smallint"
159+
assert pa_to_trino_value_type(str(pa.int32())) == "int"
160+
assert pa_to_trino_value_type(str(pa.int64())) == "bigint"
161+
assert pa_to_trino_value_type(str(pa.float32())) == "double"
162+
assert pa_to_trino_value_type(str(pa.float64())) == "double"
163+
assert pa_to_trino_value_type(str(pa.binary())) == "binary"
164+
165+
def test_string(self) -> None:
166+
assert pa_to_trino_value_type(str(pa.string())) == "varchar"
167+
assert pa_to_trino_value_type("large_string") == "varchar"
168+
assert pa_to_trino_value_type("char") == "varchar"
169+
170+
def test_varbinary(self) -> None:
171+
assert pa_to_trino_value_type("varbinary") == "binary"
172+
173+
def test_date(self) -> None:
174+
assert pa_to_trino_value_type(str(pa.date32())) == "date"
175+
176+
def test_timestamp(self) -> None:
177+
assert pa_to_trino_value_type(str(pa.timestamp("us"))) == "timestamp"
178+
assert (
179+
pa_to_trino_value_type(str(pa.timestamp("us", tz="UTC")))
180+
== "timestamp with time zone"
181+
)
182+
183+
def test_decimal128(self) -> None:
184+
assert pa_to_trino_value_type(str(pa.decimal128(10, 2))) == "decimal(10, 2)"
185+
186+
def test_decimal256(self) -> None:
187+
assert pa_to_trino_value_type(str(pa.decimal256(10, 2))) == "decimal(10, 2)"
188+
189+
def test_list(self) -> None:
190+
assert pa_to_trino_value_type(str(pa.list_(pa.int64()))) == "array<bigint>"
191+
192+
def test_list_of_string(self) -> None:
193+
assert pa_to_trino_value_type(str(pa.list_(pa.string()))) == "array<varchar>"
194+
195+
def test_map_degrades_to_varchar(self) -> None:
196+
type_str = str(pa.map_(pa.string(), pa.int64()))
197+
assert pa_to_trino_value_type(type_str) == "varchar"
198+
199+
def test_struct_degrades_to_varchar(self) -> None:
200+
type_str = str(pa.struct([("x", pa.int64())]))
201+
assert pa_to_trino_value_type(type_str) == "varchar"
202+
203+
def test_null(self) -> None:
204+
assert pa_to_trino_value_type(str(pa.null())) == "null"

0 commit comments

Comments
 (0)