ritvik360/nl2sql-bench

1"""2nl2sql-bench/server/tasks/medium.py3=====================================4Task 2 — Join & Aggregation  (difficulty: medium)5 6Questions require at least one JOIN and GROUP BY / HAVING.7Expect most frontier models to succeed in 2–3 steps.8"""9 10from __future__ import annotations11 12from .base import BaseTask, TaskExample, register13 14 15@register16class JoinAggregationTask(BaseTask):17    name = "join-aggregation"18    difficulty = "medium"19 20    examples = [21        TaskExample(22            question=(23                "How many orders has each customer placed? "24                "Return columns: customer_name, order_count. "25                "Include customers with zero orders. "26                "Sort by order_count descending, then customer_name ascending."27            ),28            sql=(29                "SELECT c.name AS customer_name, COUNT(o.id) AS order_count "30                "FROM customers c "31                "LEFT JOIN orders o ON c.id = o.customer_id "32                "GROUP BY c.id, c.name "33                "ORDER BY order_count DESC, customer_name ASC"34            ),35            notes="LEFT JOIN to include zero-order customers, COUNT aggregate.",36        ),37        TaskExample(38            question=(39                "What is the average product rating per category? "40                "Only include categories that have at least one review. "41                "Return columns: category_name, avg_rating. "42                "Round avg_rating to 2 decimal places. "43                "Sort by avg_rating descending."44            ),45            sql=(46                "SELECT c.name AS category_name, "47                "       ROUND(AVG(r.rating), 2) AS avg_rating "48                "FROM categories c "49                "JOIN products p ON p.category_id = c.id "50                "JOIN reviews r ON r.product_id = p.id "51                "GROUP BY c.id, c.name "52                "ORDER BY avg_rating DESC"53            ),54            notes="Two JOINs, AVG aggregate, ROUND function.",55        ),56        TaskExample(57            question=(58                "Which categories have more than 5 products in stock "59                "(i.e., stock_quantity > 0)? "60                "Return columns: category_name, in_stock_count. "61                "Sort by in_stock_count descending."62            ),63            sql=(64                "SELECT c.name AS category_name, "65                "       COUNT(p.id) AS in_stock_count "66                "FROM categories c "67                "JOIN products p ON p.category_id = c.id "68                "WHERE p.stock_quantity > 0 "69                "GROUP BY c.id, c.name "70                "HAVING COUNT(p.id) > 5 "71                "ORDER BY in_stock_count DESC"72            ),73            notes="WHERE before GROUP BY, HAVING filter on aggregate.",74        ),75        TaskExample(76            question=(77                "Which customers have spent more than $500 total on delivered orders? "78                "Return columns: customer_name, total_spent. "79                "Round total_spent to 2 decimal places. "80                "Sort by total_spent descending."81            ),82            sql=(83                "SELECT c.name AS customer_name, "84                "       ROUND(SUM(o.total_amount), 2) AS total_spent "85                "FROM customers c "86                "JOIN orders o ON o.customer_id = c.id "87                "WHERE o.status = 'delivered' "88                "GROUP BY c.id, c.name "89                "HAVING SUM(o.total_amount) > 500 "90                "ORDER BY total_spent DESC"91            ),92            notes="SUM aggregate, HAVING on SUM, status filter.",93        ),94        TaskExample(95            question=(96                "Show the total quantity sold for each product. "97                "Only include products that appear in at least one order item. "98                "Return columns: product_name, total_quantity_sold. "99                "Sort by total_quantity_sold descending."100            ),101            sql=(102                "SELECT p.name AS product_name, "103                "       SUM(oi.quantity) AS total_quantity_sold "104                "FROM products p "105                "JOIN order_items oi ON oi.product_id = p.id "106                "GROUP BY p.id, p.name "107                "ORDER BY total_quantity_sold DESC"108            ),109            notes="JOIN on order_items, SUM aggregate.",110        ),111    ]112 113    def description(self) -> str:114        return (115            "Multi-table JOIN queries with GROUP BY, HAVING, and aggregation "116            "functions (COUNT, SUM, AVG, ROUND). Tests relational reasoning."117        )118