Spaces:
Sleeping
Sleeping
File size: 4,670 Bytes
ede2fa4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """
Task definitions for the SQL Query Environment.
Each task has:
- task_id: unique identifier
- description: natural language question
- difficulty: easy / medium / hard
- expected_sql: a reference SQL that produces the correct answer
- expected_columns: the column names in the expected result
- expected_rows: the exact rows expected (deterministic)
"""
from dataclasses import dataclass, field
@dataclass
class TaskDefinition:
task_id: str
description: str
difficulty: str
expected_sql: str
expected_columns: list[str]
expected_rows: list[tuple]
# ββββββββββββββββββββββββββββββββββββββββββββββ
# TASK 1 β EASY: Single table, simple WHERE + ORDER BY
# ββββββββββββββββββββββββββββββββββββββββββββββ
TASK_1 = TaskDefinition(
task_id="task_1",
description=(
"List the names and salaries of all active employees who earn more "
"than 70000, ordered by salary descending."
),
difficulty="easy",
expected_sql="""
SELECT name, salary
FROM employees
WHERE is_active = 1 AND salary > 70000
ORDER BY salary DESC;
""",
expected_columns=["name", "salary"],
expected_rows=[
("Ravi Krishnan", 95000.0),
("Priya Patel", 92000.0),
("Amit Rao", 88000.0),
("Arjun Sharma", 85000.0),
("Lakshmi Iyer", 82000.0),
("Rahul Verma", 78000.0),
("Anita Desai", 72000.0),
("Nikhil Agarwal", 71000.0),
],
)
# ββββββββββββββββββββββββββββββββββββββββββββββ
# TASK 2 β MEDIUM: JOIN + GROUP BY + HAVING
# ββββββββββββββββββββββββββββββββββββββββββββββ
TASK_2 = TaskDefinition(
task_id="task_2",
description=(
"For each department, show the department name, the number of active "
"employees, and the average salary of active employees. Only include "
"departments that have more than 2 active employees. "
"Order by average salary descending."
),
difficulty="medium",
expected_sql="""
SELECT d.name, COUNT(e.id) AS num_employees, ROUND(AVG(e.salary), 2) AS avg_salary
FROM departments d
JOIN employees e ON d.id = e.department_id
WHERE e.is_active = 1
GROUP BY d.id, d.name
HAVING COUNT(e.id) > 2
ORDER BY avg_salary DESC;
""",
expected_columns=["name", "num_employees", "avg_salary"],
expected_rows=[
("Engineering", 4, 87500.0),
("Sales", 3, 70333.33),
("Marketing", 3, 65666.67),
],
)
# ββββββββββββββββββββββββββββββββββββββββββββββ
# TASK 3 β HARD: Subquery + multiple JOINs + complex logic
# ββββββββββββββββββββββββββββββββββββββββββββββ
TASK_3 = TaskDefinition(
task_id="task_3",
description=(
"Find all active employees who lead at least one active project and "
"whose salary is above the overall average salary of all active employees. "
"For each such employee, show their name, salary, department name, and "
"the total budget of active projects they lead. "
"Order by total project budget descending."
),
difficulty="hard",
expected_sql="""
SELECT
e.name,
e.salary,
d.name AS department_name,
SUM(p.budget) AS total_project_budget
FROM employees e
JOIN departments d ON e.department_id = d.id
JOIN projects p ON p.lead_employee_id = e.id
WHERE e.is_active = 1
AND p.status = 'active'
AND e.salary > (
SELECT AVG(salary) FROM employees WHERE is_active = 1
)
GROUP BY e.id, e.name, e.salary, d.name
ORDER BY total_project_budget DESC;
""",
expected_columns=["name", "salary", "department_name", "total_project_budget"],
expected_rows=[
("Ravi Krishnan", 95000.0, "Engineering", 150000.0),
("Priya Patel", 92000.0, "Engineering", 120000.0),
("Arjun Sharma", 85000.0, "Engineering", 80000.0),
],
)
# All tasks in order
ALL_TASKS = {
"task_1": TASK_1,
"task_2": TASK_2,
"task_3": TASK_3,
}
TASK_LIST = [TASK_1, TASK_2, TASK_3]
|