Fix mapping of MAP types with boolean or integer keys

Author: hashhar

Diff for: tests/integration/test_types_integration.py

@@ -799,11 +799,9 @@ def test_map(trino_connection):
         SqlTest(trino_connection)
         .add_field(sql="CAST(null AS MAP(VARCHAR, INTEGER))", python=None)
         .add_field(sql="MAP()", python={})
-        # TODO: boolean map keys get returned as strings
-        # .add_field(sql="MAP(ARRAY[true, false], ARRAY[false, true])", python={True: False, False: True})
-        # .add_field(sql="MAP(ARRAY[true, false], ARRAY[true, null])", python={True: True, False: None})
-        # TODO: integer map keys get returned as strings
-        # .add_field(sql="MAP(ARRAY[1, 2], ARRAY[1, null])", python={1: 1, 2: None})
+        .add_field(sql="MAP(ARRAY[true, false], ARRAY[false, true])", python={True: False, False: True})
+        .add_field(sql="MAP(ARRAY[true, false], ARRAY[true, null])", python={True: True, False: None})
+        .add_field(sql="MAP(ARRAY[1, 2], ARRAY[1, null])", python={1: 1, 2: None})
         .add_field(sql="MAP("
                        "ARRAY[CAST('NaN' AS REAL), CAST('-Infinity' AS REAL), CAST(3.4028235E38 AS REAL), CAST(1.4E-45 AS REAL), CAST('Infinity' AS REAL), CAST(1 AS REAL)], "  # noqa: E501
                        "ARRAY[CAST('NaN' AS REAL), CAST('-Infinity' AS REAL), CAST(3.4028235E38 AS REAL), CAST(1.4E-45 AS REAL), CAST('Infinity' AS REAL), null])",  # noqa: E501

Diff for: trino/mapper.py

@@ -31,6 +31,29 @@ def map(self, value: Any) -> Optional[T]:
         pass
 
 
+class BooleanValueMapper(ValueMapper[bool]):
+    def map(self, value: Any) -> Optional[bool]:
+        if value is None:
+            return None
+        if isinstance(value, bool):
+            return value
+        if str(value).lower() == 'true':
+            return True
+        if str(value).lower() == 'false':
+            return False
+        raise ValueError(f"Server sent unexpected value {value} of type {type(value)} for boolean")
+
+
+class IntegerValueMapper(ValueMapper[int]):
+    def map(self, value: Any) -> Optional[int]:
+        if value is None:
+            return None
+        if isinstance(value, int):
+            return value
+        # int(3.1) == 3 but server won't send such values for integer types
+        return int(value)
+
+
 class DoubleValueMapper(ValueMapper[float]):
     def map(self, value) -> Optional[float]:
         if value is None:
@@ -221,6 +244,10 @@ def _create_value_mapper(self, column) -> ValueMapper:
         col_type = column['rawType']
 
         # primitive types
+        if col_type == 'boolean':
+            return BooleanValueMapper()
+        if col_type in {'tinyint', 'smallint', 'integer', 'bigint'}:
+            return IntegerValueMapper()
         if col_type in {'double', 'real'}:
             return DoubleValueMapper()
         if col_type == 'decimal':