Skip to content

Commit 3425783

Browse files
nikolauspschuetzntkathole
nikolauspschuetz
authored and
ntkathole
committed
fix: Widen Athena integer type mapping for unsigned ints
pa_to_athena_value_type mapped every unsigned Arrow int (uint8/16/32/64) to Athena tinyint, a signed 8-bit type (-128..127). Any unsigned value above 127 (e.g. a uint32 column) overflows the type in the generated CREATE TABLE DDL built by aws_utils.py. Widen each unsigned type to the next-larger signed Athena type, matching the widening already used for Postgres in arrow_to_pg_type in this module (uint8->smallint, uint16->int, uint32/uint64->bigint). Signed inputs are unchanged. Adds a regression test for both the widened and signed cases. Signed-off-by: Nikolaus Schuetz <nikolauspschuetz@gmail.com>
1 parent 934d341 commit 3425783

2 files changed

Lines changed: 28 additions & 4 deletions

File tree

‎sdk/python/feast/type_map.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2291,10 +2291,14 @@ def pa_to_athena_value_type(pa_type: "pyarrow.DataType") -> str:
22912291
"int16": "smallint",
22922292
"int32": "int",
22932293
"int64": "bigint",
2294-
"uint8": "tinyint",
2295-
"uint16": "tinyint",
2296-
"uint32": "tinyint",
2297-
"uint64": "tinyint",
2294+
# Athena integer types are signed, so unsigned Arrow types must be
2295+
# widened to the next-larger signed type to avoid overflow in the
2296+
# generated DDL (e.g. uint32 exceeds signed int's max). This mirrors
2297+
# the widening already done in arrow_to_pg_type for Postgres.
2298+
"uint8": "smallint",
2299+
"uint16": "int",
2300+
"uint32": "bigint",
2301+
"uint64": "bigint",
22982302
"float": "float",
22992303
"double": "double",
23002304
"binary": "binary",

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
arrow_to_pg_type,
1515
feast_value_type_to_pa,
1616
feast_value_type_to_python_type,
17+
pa_to_athena_value_type,
1718
pa_to_feast_value_type,
1819
pa_to_redshift_value_type,
1920
pg_type_to_feast_value_type,
@@ -531,6 +532,25 @@ def test_arrow_to_pg_type_map(self):
531532
assert arrow_to_pg_type("map<string, string>") == "jsonb"
532533
assert arrow_to_pg_type("map<string, int64>") == "jsonb"
533534

535+
def test_pa_to_athena_value_type_unsigned_ints_widen(self):
536+
"""Unsigned Arrow ints must widen to a signed Athena type that can
537+
hold their full range. Athena has no unsigned integer types, so
538+
mapping every uintN to tinyint (signed -128..127) overflows for any
539+
value above 127 in the generated CREATE TABLE DDL. Each uintN must map
540+
to the next-larger signed type, matching arrow_to_pg_type's widening.
541+
"""
542+
assert pa_to_athena_value_type(pyarrow.uint8()) == "smallint"
543+
assert pa_to_athena_value_type(pyarrow.uint16()) == "int"
544+
assert pa_to_athena_value_type(pyarrow.uint32()) == "bigint"
545+
assert pa_to_athena_value_type(pyarrow.uint64()) == "bigint"
546+
547+
def test_pa_to_athena_value_type_signed_ints_unchanged(self):
548+
"""Signed Arrow ints keep their same-width Athena type (no regression)."""
549+
assert pa_to_athena_value_type(pyarrow.int8()) == "tinyint"
550+
assert pa_to_athena_value_type(pyarrow.int16()) == "smallint"
551+
assert pa_to_athena_value_type(pyarrow.int32()) == "int"
552+
assert pa_to_athena_value_type(pyarrow.int64()) == "bigint"
553+
534554
def test_pg_type_to_feast_value_type_json(self):
535555
"""Test that Postgres json/jsonb types convert to ValueType.MAP."""
536556
assert pg_type_to_feast_value_type("json") == ValueType.MAP

0 commit comments

Comments
 (0)