Merge branch 'main' into fix/reduce-lambda-layer-size-icu

Author: kukushking

awswrangler/_data_types.py

@@ -59,6 +59,8 @@ def pyarrow2athena(  # noqa: PLR0911,PLR0912
         )
     if pa.types.is_map(dtype):
         return f"map<{pyarrow2athena(dtype=dtype.key_type, ignore_null=ignore_null)},{pyarrow2athena(dtype=dtype.item_type, ignore_null=ignore_null)}>"
+    if isinstance(dtype, getattr(pa, "BaseExtensionType", pa.ExtensionType)):
+        return pyarrow2athena(dtype=dtype.storage_type, ignore_null=ignore_null)
     if dtype == pa.null():
         if ignore_null:
             return ""

pyproject.toml

@@ -33,7 +33,7 @@ redshift = ["redshift-connector>=2.0.0,<3"]
 mysql = ["pymysql>=1.0.0,<2"]
 postgres = ["pg8000>=1.29.0,<2"]
 sqlserver = ["pyodbc>=4,<6"]
-oracle = ["oracledb>=1,<4"]
+oracle = ["oracledb>=1,<5"]
 gremlin = [
     "gremlinpython>=3.7.1,<4",
     "requests>=2.33.0,<3",
@@ -52,7 +52,7 @@ opensearch = [
 ]
 openpyxl = ["openpyxl>=3.0.0,<4"]
 progressbar = ["progressbar2>=4.0.0,<5"]
-deltalake = ["deltalake>=0.18.0,<1.6.0"]
+deltalake = ["deltalake>=0.18.0,<1.7.0"]
 pyiceberg = ["pyiceberg[pyarrow]>=0.7.0,<1"]
 geopandas = ["geopandas>=1.0.0,<2", "pyproj>=3.6,<3.7.2"]
 modin = ["modin>=0.31,<0.38"]

tests/unit/test_s3_parquet.py

@@ -84,6 +84,42 @@ def test_read_parquet_metadata_large_dtype(path):
     assert columns_types.get("c1") == "string"
 
 
+def test_pyarrow2athena_uuid_extension_type():
+    from awswrangler._data_types import pyarrow2athena
+
+    assert pyarrow2athena(pa.uuid()) == "binary"
+
+
+def test_pyarrow2athena_custom_extension_type():
+    from awswrangler._data_types import pyarrow2athena
+
+    class _TestExtType(pa.ExtensionType):
+        def __init__(self):
+            super().__init__(pa.int64(), "test.custom_ext")
+
+        def __arrow_ext_serialize__(self):
+            return b""
+
+        @classmethod
+        def __arrow_ext_deserialize__(cls, storage_type, serialized):
+            return cls()
+
+    assert pyarrow2athena(_TestExtType()) == "bigint"
+
+
+def test_athena_types_from_pyarrow_schema_with_extension():
+    from awswrangler._data_types import athena_types_from_pyarrow_schema
+
+    schema = pa.schema(
+        [
+            pa.field("id", pa.uuid()),
+            pa.field("value", pa.int64()),
+        ]
+    )
+    result = athena_types_from_pyarrow_schema(schema)
+    assert result == {"id": "binary", "value": "bigint"}
+
+
 @pytest.mark.parametrize(
     "partition_cols",
     [