Spaces:

anu151105
/

myenv

Sleeping

prem1505

Initialize space with OpenEnv SQL Optimizer

1e97968 8 months ago

10.2 kB

"""
tasks.py — The three SQL optimization tasks (easy → medium → hard).
Each task mirrors a real-world scenario a DBA or backend engineer would face:
  1. Easy   — SELECT * + missing WHERE index on an orders table
  2. Medium — Correlated subquery anti-pattern on a user analytics table
  3. Hard   — Inefficient aggregation with missing CTE / window function optimization
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any


@dataclass
class Task:
    task_id: str
    difficulty: str          # easy | medium | hard
    description: str
    schema_ddl: str
    slow_query: str
    reference_query: str     # canonical optimized solution (for grading)
    execution_stats: dict[str, Any]
    # Structural signals the grader checks for in the optimized query
    expected_improvements: list[str] = field(default_factory=list)
    forbidden_patterns: list[str] = field(default_factory=list)
    # Maximum estimated_cost for full performance score
    target_cost: float = 100.0


# ---------------------------------------------------------------------------
# Task 1 — Easy: SELECT * + no index on filtered column
# ---------------------------------------------------------------------------

TASK_EASY = Task(
    task_id="simple_select_optimization",
    difficulty="easy",
    description=(
        "Optimize a query that uses SELECT * and filters on an un-indexed column. "
        "The agent should replace SELECT * with specific columns and suggest an index hint."
    ),
    schema_ddl="""
-- E-commerce orders table (50M rows in production)
CREATE TABLE orders (
    order_id     BIGINT PRIMARY KEY,
    customer_id  BIGINT NOT NULL,
    status       VARCHAR(20) NOT NULL,   -- 'pending','shipped','delivered','cancelled'
    total_amount DECIMAL(12,2) NOT NULL,
    created_at   TIMESTAMP NOT NULL,
    updated_at   TIMESTAMP NOT NULL,
    shipping_address TEXT,
    notes        TEXT
);
-- Existing indexes
CREATE INDEX idx_orders_customer ON orders(customer_id);
-- NOTE: No index on (status, created_at)
""",
    slow_query="""
SELECT *
FROM orders
WHERE status = 'pending'
  AND created_at >= '2024-01-01'
ORDER BY created_at DESC;
""",
    reference_query="""
SELECT order_id, customer_id, status, total_amount, created_at
FROM orders
WHERE status = 'pending'
  AND created_at >= '2024-01-01'
ORDER BY created_at DESC;
-- Recommended: CREATE INDEX idx_orders_status_created ON orders(status, created_at DESC);
""",
    execution_stats={
        "estimated_cost": 9800.0,
        "scan_type": "seq_scan",
        "rows_examined": 50_000_000,
        "has_anti_patterns": ["SELECT *", "seq_scan on large table"],
    },
    expected_improvements=["specific columns instead of *", "index hint or note"],
    forbidden_patterns=["SELECT \\*"],
    target_cost=200.0,
)


# ---------------------------------------------------------------------------
# Task 2 — Medium: Correlated subquery anti-pattern
# ---------------------------------------------------------------------------

TASK_MEDIUM = Task(
    task_id="join_optimization",
    difficulty="medium",
    description=(
        "Rewrite a correlated subquery that runs once per row into an efficient "
        "JOIN or CTE. The query finds top customers who spent more than the average "
        "order value in their region."
    ),
    schema_ddl="""
CREATE TABLE customers (
    customer_id  BIGINT PRIMARY KEY,
    name         VARCHAR(255) NOT NULL,
    email        VARCHAR(255) UNIQUE NOT NULL,
    region       VARCHAR(50) NOT NULL,
    created_at   TIMESTAMP NOT NULL
);
CREATE TABLE orders (
    order_id     BIGINT PRIMARY KEY,
    customer_id  BIGINT NOT NULL REFERENCES customers(customer_id),
    region       VARCHAR(50) NOT NULL,
    total_amount DECIMAL(12,2) NOT NULL,
    status       VARCHAR(20) NOT NULL,
    created_at   TIMESTAMP NOT NULL
);
CREATE INDEX idx_orders_customer   ON orders(customer_id);
CREATE INDEX idx_orders_region     ON orders(region);
CREATE INDEX idx_customers_region  ON customers(region);
""",
    slow_query="""
SELECT
    c.customer_id,
    c.name,
    c.region,
    SUM(o.total_amount) AS total_spent
FROM customers c
JOIN orders o ON o.customer_id = c.customer_id
WHERE o.status = 'delivered'
  AND o.total_amount > (
      SELECT AVG(o2.total_amount)
      FROM orders o2
      WHERE o2.region = o.region
        AND o2.status = 'delivered'
  )
GROUP BY c.customer_id, c.name, c.region
ORDER BY total_spent DESC
LIMIT 100;
""",
    reference_query="""
WITH regional_avg AS (
    SELECT region, AVG(total_amount) AS avg_amount
    FROM orders
    WHERE status = 'delivered'
    GROUP BY region
)
SELECT
    c.customer_id,
    c.name,
    c.region,
    SUM(o.total_amount) AS total_spent
FROM customers c
JOIN orders o ON o.customer_id = c.customer_id
JOIN regional_avg ra ON ra.region = o.region
WHERE o.status = 'delivered'
  AND o.total_amount > ra.avg_amount
GROUP BY c.customer_id, c.name, c.region
ORDER BY total_spent DESC
LIMIT 100;
""",
    execution_stats={
        "estimated_cost": 48_750.0,
        "scan_type": "seq_scan",
        "rows_examined": 12_000_000,
        "has_anti_patterns": [
            "correlated subquery (runs once per row)",
            "missing CTE for repeated aggregation",
        ],
    },
    expected_improvements=["CTE or derived table for regional_avg", "JOIN instead of correlated subquery"],
    forbidden_patterns=["SELECT.*AVG.*FROM.*WHERE.*=.*\\.region"],  # correlated subquery pattern
    target_cost=1_500.0,
)


# ---------------------------------------------------------------------------
# Task 3 — Hard: Complex aggregation with window functions needed
# ---------------------------------------------------------------------------

TASK_HARD = Task(
    task_id="aggregation_optimization",
    difficulty="hard",
    description=(
        "Optimize a complex monthly cohort revenue report that uses repeated "
        "self-joins and subqueries where window functions and a single-pass CTE "
        "pipeline would be dramatically faster."
    ),
    schema_ddl="""
CREATE TABLE events (
    event_id    BIGINT PRIMARY KEY,
    user_id     BIGINT NOT NULL,
    event_type  VARCHAR(50) NOT NULL,   -- 'purchase','refund','view','signup'
    amount      DECIMAL(12,2),
    occurred_at TIMESTAMP NOT NULL,
    session_id  UUID,
    metadata    JSONB
);
CREATE TABLE users (
    user_id      BIGINT PRIMARY KEY,
    signup_date  DATE NOT NULL,
    plan         VARCHAR(20) NOT NULL,  -- 'free','pro','enterprise'
    country      VARCHAR(3) NOT NULL
);
CREATE INDEX idx_events_user       ON events(user_id);
CREATE INDEX idx_events_type_date  ON events(event_type, occurred_at);
CREATE INDEX idx_users_signup      ON users(signup_date);
CREATE INDEX idx_users_plan        ON users(plan);
""",
    slow_query="""
SELECT
    DATE_TRUNC('month', u.signup_date) AS cohort_month,
    u.plan,
    COUNT(DISTINCT u.user_id) AS cohort_size,
    (
        SELECT COALESCE(SUM(e.amount), 0)
        FROM events e
        WHERE e.user_id = u.user_id
          AND e.event_type = 'purchase'
          AND e.occurred_at >= DATE_TRUNC('month', u.signup_date)
          AND e.occurred_at <  DATE_TRUNC('month', u.signup_date) + INTERVAL '30 days'
    ) AS m0_revenue,
    (
        SELECT COALESCE(SUM(e.amount), 0)
        FROM events e
        WHERE e.user_id = u.user_id
          AND e.event_type = 'purchase'
          AND e.occurred_at >= DATE_TRUNC('month', u.signup_date) + INTERVAL '30 days'
          AND e.occurred_at <  DATE_TRUNC('month', u.signup_date) + INTERVAL '60 days'
    ) AS m1_revenue,
    (
        SELECT COALESCE(SUM(e.amount), 0)
        FROM events e
        WHERE e.user_id = u.user_id
          AND e.event_type = 'purchase'
          AND e.occurred_at >= DATE_TRUNC('month', u.signup_date) + INTERVAL '60 days'
          AND e.occurred_at <  DATE_TRUNC('month', u.signup_date) + INTERVAL '90 days'
    ) AS m2_revenue
FROM users u
WHERE u.signup_date >= '2023-01-01'
GROUP BY cohort_month, u.plan, u.user_id, u.signup_date
ORDER BY cohort_month, u.plan;
""",
    reference_query="""
WITH user_purchases AS (
    SELECT
        e.user_id,
        e.amount,
        e.occurred_at,
        u.signup_date,
        u.plan,
        DATE_TRUNC('month', u.signup_date) AS cohort_month,
        (e.occurred_at - u.signup_date::TIMESTAMP) AS days_since_signup
    FROM users u
    JOIN events e ON e.user_id = u.user_id
    WHERE u.signup_date >= '2023-01-01'
      AND e.event_type = 'purchase'
),
cohort_revenue AS (
    SELECT
        cohort_month,
        plan,
        user_id,
        SUM(amount) FILTER (WHERE days_since_signup < INTERVAL '30 days')  AS m0_revenue,
        SUM(amount) FILTER (WHERE days_since_signup >= INTERVAL '30 days'
                              AND days_since_signup <  INTERVAL '60 days') AS m1_revenue,
        SUM(amount) FILTER (WHERE days_since_signup >= INTERVAL '60 days'
                              AND days_since_signup <  INTERVAL '90 days') AS m2_revenue
    FROM user_purchases
    GROUP BY cohort_month, plan, user_id
)
SELECT
    cohort_month,
    plan,
    COUNT(DISTINCT user_id)  AS cohort_size,
    COALESCE(SUM(m0_revenue), 0) AS m0_revenue,
    COALESCE(SUM(m1_revenue), 0) AS m1_revenue,
    COALESCE(SUM(m2_revenue), 0) AS m2_revenue
FROM cohort_revenue
GROUP BY cohort_month, plan
ORDER BY cohort_month, plan;
""",
    execution_stats={
        "estimated_cost": 892_000.0,
        "scan_type": "seq_scan",
        "rows_examined": 450_000_000,
        "has_anti_patterns": [
            "3x correlated scalar subqueries",
            "GROUP BY includes user_id causing over-aggregation",
            "no FILTER aggregate — using WHERE in subqueries",
        ],
    },
    expected_improvements=[
        "CTE to pre-join events once",
        "FILTER aggregate or CASE WHEN instead of correlated subqueries",
        "single GROUP BY at cohort level",
    ],
    forbidden_patterns=[],
    target_cost=15_000.0,
)


TASKS: dict[str, Task] = {
    TASK_EASY.task_id:   TASK_EASY,
    TASK_MEDIUM.task_id: TASK_MEDIUM,
    TASK_HARD.task_id:   TASK_HARD,
}

	"""
	tasks.py — The three SQL optimization tasks (easy → medium → hard).

	Each task mirrors a real-world scenario a DBA or backend engineer would face:
	1. Easy — SELECT * + missing WHERE index on an orders table
	2. Medium — Correlated subquery anti-pattern on a user analytics table
	3. Hard — Inefficient aggregation with missing CTE / window function optimization
	"""

	from __future__ import annotations
	from dataclasses import dataclass, field
	from typing import Any


	@dataclass
	class Task:
	task_id: str
	difficulty: str # easy \| medium \| hard
	description: str
	schema_ddl: str
	slow_query: str
	reference_query: str # canonical optimized solution (for grading)
	execution_stats: dict[str, Any]
	# Structural signals the grader checks for in the optimized query
	expected_improvements: list[str] = field(default_factory=list)
	forbidden_patterns: list[str] = field(default_factory=list)
	# Maximum estimated_cost for full performance score
	target_cost: float = 100.0


	# ---------------------------------------------------------------------------
	# Task 1 — Easy: SELECT * + no index on filtered column
	# ---------------------------------------------------------------------------

	TASK_EASY = Task(
	task_id="simple_select_optimization",
	difficulty="easy",
	description=(
	"Optimize a query that uses SELECT * and filters on an un-indexed column. "
	"The agent should replace SELECT * with specific columns and suggest an index hint."
	),
	schema_ddl="""
	-- E-commerce orders table (50M rows in production)
	CREATE TABLE orders (
	order_id BIGINT PRIMARY KEY,
	customer_id BIGINT NOT NULL,
	status VARCHAR(20) NOT NULL, -- 'pending','shipped','delivered','cancelled'
	total_amount DECIMAL(12,2) NOT NULL,
	created_at TIMESTAMP NOT NULL,
	updated_at TIMESTAMP NOT NULL,
	shipping_address TEXT,
	notes TEXT
	);

	-- Existing indexes
	CREATE INDEX idx_orders_customer ON orders(customer_id);
	-- NOTE: No index on (status, created_at)
	""",
	slow_query="""
	SELECT *
	FROM orders
	WHERE status = 'pending'
	AND created_at >= '2024-01-01'
	ORDER BY created_at DESC;
	""",
	reference_query="""
	SELECT order_id, customer_id, status, total_amount, created_at
	FROM orders
	WHERE status = 'pending'
	AND created_at >= '2024-01-01'
	ORDER BY created_at DESC;
	-- Recommended: CREATE INDEX idx_orders_status_created ON orders(status, created_at DESC);
	""",
	execution_stats={
	"estimated_cost": 9800.0,
	"scan_type": "seq_scan",
	"rows_examined": 50_000_000,
	"has_anti_patterns": ["SELECT *", "seq_scan on large table"],
	},
	expected_improvements=["specific columns instead of *", "index hint or note"],
	forbidden_patterns=["SELECT \\*"],
	target_cost=200.0,
	)


	# ---------------------------------------------------------------------------
	# Task 2 — Medium: Correlated subquery anti-pattern
	# ---------------------------------------------------------------------------

	TASK_MEDIUM = Task(
	task_id="join_optimization",
	difficulty="medium",
	description=(
	"Rewrite a correlated subquery that runs once per row into an efficient "
	"JOIN or CTE. The query finds top customers who spent more than the average "
	"order value in their region."
	),
	schema_ddl="""
	CREATE TABLE customers (
	customer_id BIGINT PRIMARY KEY,
	name VARCHAR(255) NOT NULL,
	email VARCHAR(255) UNIQUE NOT NULL,
	region VARCHAR(50) NOT NULL,
	created_at TIMESTAMP NOT NULL
	);

	CREATE TABLE orders (
	order_id BIGINT PRIMARY KEY,
	customer_id BIGINT NOT NULL REFERENCES customers(customer_id),
	region VARCHAR(50) NOT NULL,
	total_amount DECIMAL(12,2) NOT NULL,
	status VARCHAR(20) NOT NULL,
	created_at TIMESTAMP NOT NULL
	);

	CREATE INDEX idx_orders_customer ON orders(customer_id);
	CREATE INDEX idx_orders_region ON orders(region);
	CREATE INDEX idx_customers_region ON customers(region);
	""",
	slow_query="""
	SELECT
	c.customer_id,
	c.name,
	c.region,
	SUM(o.total_amount) AS total_spent
	FROM customers c
	JOIN orders o ON o.customer_id = c.customer_id
	WHERE o.status = 'delivered'
	AND o.total_amount > (
	SELECT AVG(o2.total_amount)
	FROM orders o2
	WHERE o2.region = o.region
	AND o2.status = 'delivered'
	)
	GROUP BY c.customer_id, c.name, c.region
	ORDER BY total_spent DESC
	LIMIT 100;
	""",
	reference_query="""
	WITH regional_avg AS (
	SELECT region, AVG(total_amount) AS avg_amount
	FROM orders
	WHERE status = 'delivered'
	GROUP BY region
	)
	SELECT
	c.customer_id,
	c.name,
	c.region,
	SUM(o.total_amount) AS total_spent
	FROM customers c
	JOIN orders o ON o.customer_id = c.customer_id
	JOIN regional_avg ra ON ra.region = o.region
	WHERE o.status = 'delivered'
	AND o.total_amount > ra.avg_amount
	GROUP BY c.customer_id, c.name, c.region
	ORDER BY total_spent DESC
	LIMIT 100;
	""",
	execution_stats={
	"estimated_cost": 48_750.0,
	"scan_type": "seq_scan",
	"rows_examined": 12_000_000,
	"has_anti_patterns": [
	"correlated subquery (runs once per row)",
	"missing CTE for repeated aggregation",
	],
	},
	expected_improvements=["CTE or derived table for regional_avg", "JOIN instead of correlated subquery"],
	forbidden_patterns=["SELECT.AVG.FROM.WHERE.=.*\\.region"], # correlated subquery pattern
	target_cost=1_500.0,
	)


	# ---------------------------------------------------------------------------
	# Task 3 — Hard: Complex aggregation with window functions needed
	# ---------------------------------------------------------------------------

	TASK_HARD = Task(
	task_id="aggregation_optimization",
	difficulty="hard",
	description=(
	"Optimize a complex monthly cohort revenue report that uses repeated "
	"self-joins and subqueries where window functions and a single-pass CTE "
	"pipeline would be dramatically faster."
	),
	schema_ddl="""
	CREATE TABLE events (
	event_id BIGINT PRIMARY KEY,
	user_id BIGINT NOT NULL,
	event_type VARCHAR(50) NOT NULL, -- 'purchase','refund','view','signup'
	amount DECIMAL(12,2),
	occurred_at TIMESTAMP NOT NULL,
	session_id UUID,
	metadata JSONB
	);

	CREATE TABLE users (
	user_id BIGINT PRIMARY KEY,
	signup_date DATE NOT NULL,
	plan VARCHAR(20) NOT NULL, -- 'free','pro','enterprise'
	country VARCHAR(3) NOT NULL
	);

	CREATE INDEX idx_events_user ON events(user_id);
	CREATE INDEX idx_events_type_date ON events(event_type, occurred_at);
	CREATE INDEX idx_users_signup ON users(signup_date);
	CREATE INDEX idx_users_plan ON users(plan);
	""",
	slow_query="""
	SELECT
	DATE_TRUNC('month', u.signup_date) AS cohort_month,
	u.plan,
	COUNT(DISTINCT u.user_id) AS cohort_size,
	(
	SELECT COALESCE(SUM(e.amount), 0)
	FROM events e
	WHERE e.user_id = u.user_id
	AND e.event_type = 'purchase'
	AND e.occurred_at >= DATE_TRUNC('month', u.signup_date)
	AND e.occurred_at < DATE_TRUNC('month', u.signup_date) + INTERVAL '30 days'
	) AS m0_revenue,
	(
	SELECT COALESCE(SUM(e.amount), 0)
	FROM events e
	WHERE e.user_id = u.user_id
	AND e.event_type = 'purchase'
	AND e.occurred_at >= DATE_TRUNC('month', u.signup_date) + INTERVAL '30 days'
	AND e.occurred_at < DATE_TRUNC('month', u.signup_date) + INTERVAL '60 days'
	) AS m1_revenue,
	(
	SELECT COALESCE(SUM(e.amount), 0)
	FROM events e
	WHERE e.user_id = u.user_id
	AND e.event_type = 'purchase'
	AND e.occurred_at >= DATE_TRUNC('month', u.signup_date) + INTERVAL '60 days'
	AND e.occurred_at < DATE_TRUNC('month', u.signup_date) + INTERVAL '90 days'
	) AS m2_revenue
	FROM users u
	WHERE u.signup_date >= '2023-01-01'
	GROUP BY cohort_month, u.plan, u.user_id, u.signup_date
	ORDER BY cohort_month, u.plan;
	""",
	reference_query="""
	WITH user_purchases AS (
	SELECT
	e.user_id,
	e.amount,
	e.occurred_at,
	u.signup_date,
	u.plan,
	DATE_TRUNC('month', u.signup_date) AS cohort_month,
	(e.occurred_at - u.signup_date::TIMESTAMP) AS days_since_signup
	FROM users u
	JOIN events e ON e.user_id = u.user_id
	WHERE u.signup_date >= '2023-01-01'
	AND e.event_type = 'purchase'
	),
	cohort_revenue AS (
	SELECT
	cohort_month,
	plan,
	user_id,
	SUM(amount) FILTER (WHERE days_since_signup < INTERVAL '30 days') AS m0_revenue,
	SUM(amount) FILTER (WHERE days_since_signup >= INTERVAL '30 days'
	AND days_since_signup < INTERVAL '60 days') AS m1_revenue,
	SUM(amount) FILTER (WHERE days_since_signup >= INTERVAL '60 days'
	AND days_since_signup < INTERVAL '90 days') AS m2_revenue
	FROM user_purchases
	GROUP BY cohort_month, plan, user_id
	)
	SELECT
	cohort_month,
	plan,
	COUNT(DISTINCT user_id) AS cohort_size,
	COALESCE(SUM(m0_revenue), 0) AS m0_revenue,
	COALESCE(SUM(m1_revenue), 0) AS m1_revenue,
	COALESCE(SUM(m2_revenue), 0) AS m2_revenue
	FROM cohort_revenue
	GROUP BY cohort_month, plan
	ORDER BY cohort_month, plan;
	""",
	execution_stats={
	"estimated_cost": 892_000.0,
	"scan_type": "seq_scan",
	"rows_examined": 450_000_000,
	"has_anti_patterns": [
	"3x correlated scalar subqueries",
	"GROUP BY includes user_id causing over-aggregation",
	"no FILTER aggregate — using WHERE in subqueries",
	],
	},
	expected_improvements=[
	"CTE to pre-join events once",
	"FILTER aggregate or CASE WHEN instead of correlated subqueries",
	"single GROUP BY at cohort level",
	],
	forbidden_patterns=[],
	target_cost=15_000.0,
	)


	TASKS: dict[str, Task] = {
	TASK_EASY.task_id: TASK_EASY,
	TASK_MEDIUM.task_id: TASK_MEDIUM,
	TASK_HARD.task_id: TASK_HARD,
	}