Skip to content

Commit 155bc24

Browse files
dbbvitorclaude
authored and
shuchu
committed
fix: Map Trino DECIMAL columns to pyarrow decimal128
trino_to_pa_value_type() mapped every DECIMAL column to a lossy float32/float64, silently discarding precision. Return decimal128 instead, matching what Table.from_pandas() already infers from the cursor's decimal.Decimal values with no schema forced on it. Co-Authored-By: Claude Sonnet 5 <[email protected]> Signed-off-by: dbbvitor <[email protected]>
1 parent b60f9c7 commit 155bc24

2 files changed

Lines changed: 35 additions & 15 deletions

File tree

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

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -146,16 +146,15 @@ def trino_to_pa_value_type(trino_type_as_str: str) -> pa.DataType:
146146
return pa.list_(trino_to_pa_value_type(array_item_type))
147147

148148
if trino_type_as_str.startswith("decimal"):
149-
search_precision = re.search(
150-
r"^decimal\((\d+)(?>,\s?\d+)?\)$", trino_type_as_str
149+
search_precision_scale = re.search(
150+
r"^decimal\((\d+)(?:,\s?(\d+))?\)$", trino_type_as_str
151151
)
152-
if search_precision:
153-
precision = int(search_precision.group(1))
154-
if precision > 32:
155-
return pa.float64()
156-
else:
157-
return pa.float32()
158-
return pa.float64()
152+
if search_precision_scale:
153+
precision = int(search_precision_scale.group(1))
154+
scale = int(search_precision_scale.group(2) or 0)
155+
# Trino's DECIMAL precision tops out at 38, which is decimal128's max too.
156+
return pa.decimal128(min(precision, 38), scale)
157+
return pa.decimal128(38, 0)
159158

160159
if trino_type_as_str.startswith("timestamp"):
161160
return pa.timestamp("us")

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

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -112,13 +112,32 @@ def test_timestamp(self) -> None:
112112
assert trino_to_pa_value_type("timestamp(3)") == pa.timestamp("us")
113113

114114
def test_decimal_bare(self) -> None:
115-
assert trino_to_pa_value_type("decimal") == pa.float64()
115+
assert trino_to_pa_value_type("decimal") == pa.decimal128(38, 0)
116116

117117
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()
118+
assert trino_to_pa_value_type("decimal(10, 2)") == pa.decimal128(10, 2)
119+
assert trino_to_pa_value_type("decimal(10,2)") == pa.decimal128(10, 2)
120+
assert trino_to_pa_value_type("decimal(38, 10)") == pa.decimal128(38, 10)
121+
assert trino_to_pa_value_type("decimal(32)") == pa.decimal128(32, 0)
122+
assert trino_to_pa_value_type("decimal(9)") == pa.decimal128(9, 0)
123+
124+
def test_decimal_schema_accepts_cursor_values(self) -> None:
125+
from decimal import Decimal
126+
127+
import pandas as pd
128+
129+
schema = pa.schema(
130+
[pa.field("amount", trino_to_pa_value_type("decimal(10, 2)"))]
131+
)
132+
df = pd.DataFrame({"amount": [Decimal("12.34"), None, Decimal("-0.01")]})
133+
134+
table = pa.Table.from_pandas(df, schema=schema, preserve_index=False)
135+
136+
assert table.column("amount").to_pylist() == [
137+
Decimal("12.34"),
138+
None,
139+
Decimal("-0.01"),
140+
]
122141

123142
def test_array_simple(self) -> None:
124143
assert trino_to_pa_value_type("array(bigint)") == pa.list_(pa.int64())
@@ -127,7 +146,9 @@ def test_array_parameterized_varchar(self) -> None:
127146
assert trino_to_pa_value_type("array(varchar(10))") == pa.list_(pa.string())
128147

129148
def test_array_parameterized_decimal(self) -> None:
130-
assert trino_to_pa_value_type("array(decimal(10, 2))") == pa.list_(pa.float32())
149+
assert trino_to_pa_value_type("array(decimal(10, 2))") == pa.list_(
150+
pa.decimal128(10, 2)
151+
)
131152

132153
def test_array_nested(self) -> None:
133154
assert trino_to_pa_value_type("array(array(bigint))") == pa.list_(

0 commit comments

Comments
 (0)