nl2sql-bench / server /tasks /medium.py
ritvik360's picture
Upload folder using huggingface_hub
a39d8ef verified
"""
nl2sql-bench/server/tasks/medium.py
=====================================
Task 2 — Join & Aggregation (difficulty: medium)
Questions require at least one JOIN and GROUP BY / HAVING.
Expect most frontier models to succeed in 2–3 steps.
"""
from __future__ import annotations
from .base import BaseTask, TaskExample, register
@register
class JoinAggregationTask(BaseTask):
name = "join-aggregation"
difficulty = "medium"
examples = [
TaskExample(
question=(
"How many orders has each customer placed? "
"Return columns: customer_name, order_count. "
"Include customers with zero orders. "
"Sort by order_count descending, then customer_name ascending."
),
sql=(
"SELECT c.name AS customer_name, COUNT(o.id) AS order_count "
"FROM customers c "
"LEFT JOIN orders o ON c.id = o.customer_id "
"GROUP BY c.id, c.name "
"ORDER BY order_count DESC, customer_name ASC"
),
notes="LEFT JOIN to include zero-order customers, COUNT aggregate.",
),
TaskExample(
question=(
"What is the average product rating per category? "
"Only include categories that have at least one review. "
"Return columns: category_name, avg_rating. "
"Round avg_rating to 2 decimal places. "
"Sort by avg_rating descending."
),
sql=(
"SELECT c.name AS category_name, "
" ROUND(AVG(r.rating), 2) AS avg_rating "
"FROM categories c "
"JOIN products p ON p.category_id = c.id "
"JOIN reviews r ON r.product_id = p.id "
"GROUP BY c.id, c.name "
"ORDER BY avg_rating DESC"
),
notes="Two JOINs, AVG aggregate, ROUND function.",
),
TaskExample(
question=(
"Which categories have more than 5 products in stock "
"(i.e., stock_quantity > 0)? "
"Return columns: category_name, in_stock_count. "
"Sort by in_stock_count descending."
),
sql=(
"SELECT c.name AS category_name, "
" COUNT(p.id) AS in_stock_count "
"FROM categories c "
"JOIN products p ON p.category_id = c.id "
"WHERE p.stock_quantity > 0 "
"GROUP BY c.id, c.name "
"HAVING COUNT(p.id) > 5 "
"ORDER BY in_stock_count DESC"
),
notes="WHERE before GROUP BY, HAVING filter on aggregate.",
),
TaskExample(
question=(
"Which customers have spent more than $500 total on delivered orders? "
"Return columns: customer_name, total_spent. "
"Round total_spent to 2 decimal places. "
"Sort by total_spent descending."
),
sql=(
"SELECT c.name AS customer_name, "
" ROUND(SUM(o.total_amount), 2) AS total_spent "
"FROM customers c "
"JOIN orders o ON o.customer_id = c.id "
"WHERE o.status = 'delivered' "
"GROUP BY c.id, c.name "
"HAVING SUM(o.total_amount) > 500 "
"ORDER BY total_spent DESC"
),
notes="SUM aggregate, HAVING on SUM, status filter.",
),
TaskExample(
question=(
"Show the total quantity sold for each product. "
"Only include products that appear in at least one order item. "
"Return columns: product_name, total_quantity_sold. "
"Sort by total_quantity_sold descending."
),
sql=(
"SELECT p.name AS product_name, "
" SUM(oi.quantity) AS total_quantity_sold "
"FROM products p "
"JOIN order_items oi ON oi.product_id = p.id "
"GROUP BY p.id, p.name "
"ORDER BY total_quantity_sold DESC"
),
notes="JOIN on order_items, SUM aggregate.",
),
]
def description(self) -> str:
return (
"Multi-table JOIN queries with GROUP BY, HAVING, and aggregation "
"functions (COUNT, SUM, AVG, ROUND). Tests relational reasoning."
)