Add Athena filter() function support to SQLAlchemy dialect

Implements support for Amazon Athena's filter() function with lambda
expressions in PyAthena's SQLAlchemy dialect, addressing issue #480.

Features:
- Basic filter() function compilation with lambda expressions
- Support for complex lambda conditions and nested field access
- Comprehensive error handling for invalid argument counts
- Type-safe implementation using isinstance checks
- Full test coverage with 7 test cases

Examples:
- filter(array_col, 'x -> x > 0')
- filter(data_col, 'x -> x["field"] > value')
- count(filter(action_col, 'x -> x["timestamp"] BETWEEN dates'))

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

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

Author: laughingman7743

pyathena/sqlalchemy/compiler.py

@@ -10,6 +10,7 @@
     IdentifierPreparer,
     SQLCompiler,
 )
+from sqlalchemy.sql.elements import BindParameter
 from sqlalchemy.sql.schema import Column
 
 from pyathena.model import (
@@ -178,6 +179,33 @@ class AthenaStatementCompiler(SQLCompiler):
     def visit_char_length_func(self, fn: "FunctionElement[Any]", **kw):
         return f"length{self.function_argspec(fn, **kw)}"
 
+    def visit_filter_func(self, fn: "FunctionElement[Any]", **kw) -> str:
+        """Compile Athena filter() function with lambda expressions.
+
+        Supports syntax: filter(array_expr, lambda_expr)
+        Example: filter(ARRAY[1, 2, 3], x -> x > 1)
+        """
+        if len(fn.clauses.clauses) != 2:
+            raise exc.CompileError(
+                f"filter() function expects exactly 2 arguments, got {len(fn.clauses.clauses)}"
+            )
+
+        array_expr = fn.clauses.clauses[0]
+        lambda_expr = fn.clauses.clauses[1]
+
+        # Process the array expression normally
+        array_sql = self.process(array_expr, **kw)
+
+        # Process lambda expression - handle string literals as lambda expressions
+        if isinstance(lambda_expr, BindParameter) and isinstance(lambda_expr.value, str):
+            # Handle string literal lambda expressions like 'x -> x > 0'
+            lambda_sql = lambda_expr.value
+        else:
+            # Process as regular SQL expression
+            lambda_sql = self.process(lambda_expr, **kw)
+
+        return f"filter({array_sql}, {lambda_sql})"
+
     def visit_cast(self, cast: "Cast[Any]", **kwargs):
         if (isinstance(cast.type, types.VARCHAR) and cast.type.length is None) or isinstance(
             cast.type, types.String

tests/pyathena/sqlalchemy/test_compiler.py

@@ -2,10 +2,14 @@
 
 from unittest.mock import Mock
 
-from sqlalchemy import Integer, String
+import pytest
+from sqlalchemy import Column, Integer, MetaData, String, Table, exc, func, select
+from sqlalchemy.sql import literal
 
+from pyathena.sqlalchemy.base import AthenaDialect
 from pyathena.sqlalchemy.compiler import AthenaTypeCompiler
 from pyathena.sqlalchemy.types import ARRAY, MAP, STRUCT, AthenaArray, AthenaMap, AthenaStruct
+from tests import ENV
 
 
 class TestAthenaTypeCompiler:
@@ -109,3 +113,107 @@ def test_visit_array_no_attributes(self):
         array_type = type("MockArray", (), {})()
         result = compiler.visit_array(array_type)
         assert result == "ARRAY<STRING>"
+
+
+class TestAthenaStatementCompiler:
+    """Test cases for Athena statement compiler functionality."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.dialect = AthenaDialect()
+        self.metadata = MetaData(schema=ENV.schema)
+        self.test_table = Table(
+            "test_athena_statement_compiler",
+            self.metadata,
+            Column("id", Integer),
+            Column("data", ARRAY(String)),
+            Column("numbers", ARRAY(Integer)),
+        )
+
+    def test_visit_filter_func_basic(self):
+        """Test basic filter() function compilation."""
+        # Test basic filter with string lambda expression
+        stmt = select(func.filter(self.test_table.c.numbers, literal("x -> x > 0")))
+        compiled = stmt.compile(dialect=self.dialect)
+
+        sql_str = str(compiled)
+        assert "filter(" in sql_str
+        assert "x -> x > 0" in sql_str
+
+    def test_visit_filter_func_array_literal(self):
+        """Test filter() function with array literal."""
+        # Test filter with array literal - using ARRAY constructor
+        stmt = select(
+            func.filter(
+                func.array(literal(1), literal(2), literal(3), literal(-1)), literal("x -> x > 0")
+            )
+        )
+        compiled = stmt.compile(dialect=self.dialect)
+
+        sql_str = str(compiled)
+        assert "filter(" in sql_str
+        assert "x -> x > 0" in sql_str
+
+    def test_visit_filter_func_complex_lambda(self):
+        """Test filter() function with complex lambda expression."""
+        # Test complex lambda expression
+        complex_lambda = literal("x -> x IS NOT NULL AND x > 5")
+        stmt = select(func.filter(self.test_table.c.numbers, complex_lambda))
+        compiled = stmt.compile(dialect=self.dialect)
+
+        sql_str = str(compiled)
+        assert "filter(" in sql_str
+        assert "x -> x IS NOT NULL AND x > 5" in sql_str
+
+    def test_visit_filter_func_nested_access(self):
+        """Test filter() function with nested field access."""
+        # Test lambda with nested field access (for complex types)
+        nested_lambda = literal("x -> x['timestamp'] > '2023-01-01'")
+        stmt = select(func.filter(self.test_table.c.data, nested_lambda))
+        compiled = stmt.compile(dialect=self.dialect)
+
+        sql_str = str(compiled)
+        assert "filter(" in sql_str
+        assert "x -> x['timestamp'] > '2023-01-01'" in sql_str
+
+    def test_visit_filter_func_wrong_argument_count(self):
+        """Test filter() function with wrong number of arguments."""
+        # Test error when wrong number of arguments provided
+        with pytest.raises(
+            exc.CompileError, match="filter\\(\\) function expects exactly 2 arguments"
+        ):
+            stmt = select(func.filter(self.test_table.c.numbers))
+            stmt.compile(dialect=self.dialect)
+
+        with pytest.raises(
+            exc.CompileError, match="filter\\(\\) function expects exactly 2 arguments"
+        ):
+            stmt = select(
+                func.filter(self.test_table.c.numbers, literal("x -> x > 0"), literal("extra_arg"))
+            )
+            stmt.compile(dialect=self.dialect)
+
+    def test_visit_filter_func_integration_example(self):
+        """Test filter() function with the original issue example."""
+        # Test the example from the GitHub issue
+        lambda_expr = literal(
+            "x -> x['timestamp'] <= '2023-10-10' AND x['timestamp'] >= '2023-10-01' "
+            "AND x['action_count'] >= 2"
+        )
+        stmt = select(func.count(func.filter(self.test_table.c.data, lambda_expr)))
+        compiled = stmt.compile(dialect=self.dialect)
+
+        sql_str = str(compiled)
+        assert "count(" in sql_str
+        assert "filter(" in sql_str
+        assert "x -> x['timestamp'] <= '2023-10-10'" in sql_str
+        assert "x['action_count'] >= 2" in sql_str
+
+    def test_visit_char_length_func_existing(self):
+        """Test existing char_length function still works."""
+        # Ensure existing functionality isn't broken
+        stmt = select(func.char_length(self.test_table.c.data))
+        compiled = stmt.compile(dialect=self.dialect)
+
+        sql_str = str(compiled)
+        assert "length(" in sql_str