Skip to content

Commit ae74360

Browse files
raashish1601raashish1601
authored
fix: Map narrow and unsigned Arrow integer types to Feast value types (#6960)
* fix: Map narrow and unsigned Arrow integer types to Feast value types Signed-off-by: Raashish Aggarwal <[email protected]> * fix: Map uint32 to INT64 in pandas dtype inference too Signed-off-by: Raashish Aggarwal <[email protected]> --------- Signed-off-by: Raashish Aggarwal <[email protected]>
1 parent 7b9fad2 commit ae74360

2 files changed

Lines changed: 39 additions & 1 deletion

File tree

‎sdk/python/feast/type_map.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -394,7 +394,7 @@ def python_type_to_feast_value_type(
394394
"int64": ValueType.INT64,
395395
"uint64": ValueType.INT64,
396396
"int32": ValueType.INT32,
397-
"uint32": ValueType.INT32,
397+
"uint32": ValueType.INT64,
398398
"int16": ValueType.INT32,
399399
"uint16": ValueType.INT32,
400400
"uint8": ValueType.INT32,
@@ -1636,8 +1636,14 @@ def pa_to_feast_value_type(pa_type_as_str: str) -> ValueType:
16361636
value_type = ValueType.STRUCT
16371637
else:
16381638
type_map = {
1639+
"int8": ValueType.INT32,
1640+
"int16": ValueType.INT32,
16391641
"int32": ValueType.INT32,
16401642
"int64": ValueType.INT64,
1643+
"uint8": ValueType.INT32,
1644+
"uint16": ValueType.INT32,
1645+
"uint32": ValueType.INT64,
1646+
"uint64": ValueType.INT64,
16411647
"double": ValueType.DOUBLE,
16421648
"float": ValueType.FLOAT,
16431649
"string": ValueType.STRING,

‎sdk/python/tests/unit/test_type_map.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2536,3 +2536,35 @@ def test_spark_to_feast_value_type_small_ints(spark_type, expected):
25362536
"""SparkSource passes ``DataType.simpleString()``, which spells ShortType and
25372537
ByteType as ``smallint`` and ``tinyint``."""
25382538
assert spark_to_feast_value_type(spark_type) == expected
2539+
2540+
2541+
@pytest.mark.parametrize(
2542+
"pa_type, expected",
2543+
[
2544+
(pyarrow.int8(), ValueType.INT32),
2545+
(pyarrow.int16(), ValueType.INT32),
2546+
(pyarrow.uint8(), ValueType.INT32),
2547+
(pyarrow.uint16(), ValueType.INT32),
2548+
(pyarrow.uint32(), ValueType.INT64),
2549+
(pyarrow.uint64(), ValueType.INT64),
2550+
(pyarrow.list_(pyarrow.int16()), ValueType.INT32_LIST),
2551+
(pyarrow.list_(pyarrow.uint32()), ValueType.INT64_LIST),
2552+
],
2553+
)
2554+
def test_pa_to_feast_value_type_small_and_unsigned_ints(pa_type, expected):
2555+
"""Narrow and unsigned Arrow integers, e.g. Parquet int16 columns, widen to
2556+
the smallest Feast integer type that holds every value."""
2557+
assert pa_to_feast_value_type(str(pa_type)) == expected
2558+
2559+
2560+
@pytest.mark.parametrize(
2561+
"dtype", ["int8", "int16", "uint8", "uint16", "uint32", "uint64"]
2562+
)
2563+
def test_pandas_and_arrow_int_inference_agree(dtype):
2564+
"""A column infers the same Feast type from its pandas dtype and its Arrow type."""
2565+
values = pd.Series([1, 2], dtype=dtype)
2566+
from_pandas = python_type_to_feast_value_type(
2567+
"f", value=values.iloc[0], type_name=str(values.dtype)
2568+
)
2569+
from_arrow = pa_to_feast_value_type(str(pyarrow.array(values).type))
2570+
assert from_pandas == from_arrow

0 commit comments

Comments
 (0)