Add ARRAY type support with compiler and converter functionality

- Add _to_array() converter function supporting JSON and native formats
- Create AthenaArray and ARRAY type classes in SQLAlchemy types
- Implement visit_array() and visit_ARRAY() methods in compiler
- Add comprehensive tests for ARRAY compiler functionality
- Follow same architectural patterns as existing STRUCT and MAP types

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: laughingman7743

pyathena/converter.py

@@ -8,7 +8,7 @@
 from copy import deepcopy
 from datetime import date, datetime, time
 from decimal import Decimal
-from typing import Any, Callable, Dict, Optional, Type
+from typing import Any, Callable, Dict, List, Optional, Type
 
 from dateutil.tz import gettz
 
@@ -78,6 +78,51 @@ def _to_json(varchar_value: Optional[str]) -> Optional[Any]:
     return json.loads(varchar_value)
 
 
+def _to_array(varchar_value: Optional[str]) -> Optional[List[Any]]:
+    """Convert array data to Python list.
+
+    Supports two formats:
+    1. JSON format: '[1, 2, 3]' or '["a", "b", "c"]' (recommended)
+    2. Athena native format: '[1, 2, 3]' (basic cases only)
+
+    For complex arrays, use CAST(array_column AS JSON) in your SQL query.
+
+    Args:
+        varchar_value: String representation of array data
+
+    Returns:
+        List representation of array, or None if parsing fails
+    """
+    if varchar_value is None:
+        return None
+
+    # Quick check: if it doesn't look like an array, return None
+    if not (varchar_value.startswith("[") and varchar_value.endswith("]")):
+        return None
+
+    # Optimize: Try JSON parsing first (most reliable)
+    try:
+        result = json.loads(varchar_value)
+        return result if isinstance(result, list) else None
+    except json.JSONDecodeError:
+        # If JSON parsing fails, fall back to basic parsing for simple cases
+        pass
+
+    inner = varchar_value[1:-1].strip()
+    if not inner:
+        return []
+
+    try:
+        # Simple array format: [1, 2, 3] or [a, b, c]
+        # For complex structures, return None to keep as string
+        if any(char in inner for char in "{}()[]"):
+            # Contains complex structures, skip parsing
+            return None
+        return _parse_array_native(inner)
+    except Exception:
+        return None
+
+
 def _to_map(varchar_value: Optional[str]) -> Optional[Dict[str, Any]]:
     """Convert map data to Python dictionary.
 
@@ -179,6 +224,35 @@ def _to_struct(varchar_value: Optional[str]) -> Optional[Dict[str, Any]]:
         return None
 
 
+def _parse_array_native(inner: str) -> Optional[List[Any]]:
+    """Parse array native format: 1, 2, 3 or a, b, c.
+
+    Args:
+        inner: Interior content of array without brackets.
+
+    Returns:
+        List with parsed values, or None if no valid values found.
+    """
+    result = []
+
+    # Simple split by comma for basic cases
+    items = [item.strip() for item in inner.split(",")]
+
+    for item in items:
+        if not item:
+            continue
+
+        # Skip items with special characters (safety check)
+        if any(char in item for char in '{}[]()="'):
+            continue
+
+        # Convert item to appropriate type
+        converted_item = _convert_value(item)
+        result.append(converted_item)
+
+    return result if result else None
+
+
 def _parse_map_native(inner: str) -> Optional[Dict[str, Any]]:
     """Parse map native format: key1=value1, key2=value2.
 
@@ -302,7 +376,7 @@ def _to_default(varchar_value: Optional[str]) -> Optional[str]:
     "date": _to_date,
     "time": _to_time,
     "varbinary": _to_binary,
-    "array": _to_default,
+    "array": _to_array,
     "map": _to_map,
     "row": _to_struct,
     "decimal": _to_decimal,

pyathena/sqlalchemy/compiler.py

@@ -17,7 +17,7 @@
     AthenaPartitionTransform,
     AthenaRowFormatSerde,
 )
-from pyathena.sqlalchemy.types import AthenaMap, AthenaStruct
+from pyathena.sqlalchemy.types import AthenaArray, AthenaMap, AthenaStruct
 
 if TYPE_CHECKING:
     from sqlalchemy import (
@@ -164,6 +164,15 @@ def visit_map(self, type_, **kw):  # noqa: N802
     def visit_MAP(self, type_, **kw):  # noqa: N802
         return self.visit_map(type_, **kw)
 
+    def visit_array(self, type_, **kw):  # noqa: N802
+        if isinstance(type_, AthenaArray):
+            item_type_str = self.process(type_.item_type, **kw)
+            return f"ARRAY<{item_type_str}>"
+        return "ARRAY<STRING>"
+
+    def visit_ARRAY(self, type_, **kw):  # noqa: N802
+        return self.visit_array(type_, **kw)
+
 
 class AthenaStatementCompiler(SQLCompiler):
     def visit_char_length_func(self, fn: "FunctionElement[Any]", **kw):

pyathena/sqlalchemy/types.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from datetime import date, datetime
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 from sqlalchemy.sql import sqltypes
 from sqlalchemy.sql.type_api import TypeEngine
@@ -106,3 +106,24 @@ def python_type(self) -> type:
 
 class MAP(AthenaMap):
     __visit_name__ = "MAP"
+
+
+class AthenaArray(TypeEngine[List[Any]]):
+    __visit_name__ = "array"
+
+    def __init__(self, item_type: Any = None) -> None:
+        if item_type is None:
+            self.item_type: TypeEngine[Any] = sqltypes.String()
+        elif isinstance(item_type, TypeEngine):
+            self.item_type = item_type
+        else:
+            # Assume it's a SQLAlchemy type class and instantiate it
+            self.item_type = item_type()
+
+    @property
+    def python_type(self) -> type:
+        return list
+
+
+class ARRAY(AthenaArray):
+    __visit_name__ = "ARRAY"

tests/pyathena/sqlalchemy/test_compiler.py

@@ -5,7 +5,7 @@
 from sqlalchemy import Integer, String
 
 from pyathena.sqlalchemy.compiler import AthenaTypeCompiler
-from pyathena.sqlalchemy.types import MAP, STRUCT, AthenaMap, AthenaStruct
+from pyathena.sqlalchemy.types import ARRAY, MAP, STRUCT, AthenaArray, AthenaMap, AthenaStruct
 
 
 class TestAthenaTypeCompiler:
@@ -80,3 +80,32 @@ def test_visit_map_no_attributes(self):
         map_type = type("MockMap", (), {})()
         result = compiler.visit_map(map_type)
         assert result == "MAP<STRING, STRING>"
+
+    def test_visit_array_default(self):
+        dialect = Mock()
+        compiler = AthenaTypeCompiler(dialect)
+        array_type = AthenaArray()
+        result = compiler.visit_array(array_type)
+        assert result == "ARRAY<STRING>"
+
+    def test_visit_array_with_type(self):
+        dialect = Mock()
+        compiler = AthenaTypeCompiler(dialect)
+        array_type = AthenaArray(Integer)
+        result = compiler.visit_array(array_type)
+        assert result == "ARRAY<INTEGER>"
+
+    def test_visit_array_uppercase(self):
+        dialect = Mock()
+        compiler = AthenaTypeCompiler(dialect)
+        array_type = ARRAY(String)
+        result = compiler.visit_ARRAY(array_type)
+        assert result == "ARRAY<STRING>" or result == "ARRAY<VARCHAR>"
+
+    def test_visit_array_no_attributes(self):
+        # Test array type without item_type attribute
+        dialect = Mock()
+        compiler = AthenaTypeCompiler(dialect)
+        array_type = type("MockArray", (), {})()
+        result = compiler.visit_array(array_type)
+        assert result == "ARRAY<STRING>"