Implement Column Pruning and Filter Pushdown on a Query Plan Tree

A query engine represents a SQL query as a tree of plan nodes. You are handed such a tree and asked to "optimize" it. The interviewer does not say which optimizations are expected, and the code must actually run on a sample plan, not just be described. The rewrites were not named up front in the interview; the two that were eventually required are the Parts below.

The exact node classes were not reported. Assume this minimal logical plan model (Python dataclasses are fine):

• Scan(table, columns): reads the listed columns of a base table.
• Filter(conjuncts, child): keeps the child's rows for which every conjunct is true. A conjunct is a comparison column op literal.
• Project(columns, child): outputs only the listed columns of the child (plain column references, no expressions).
• Join(left, right, left_key, right_key): an inner equi-join on left.left_key = right.right_key that outputs all columns of both sides.

Column names are globally unique across tables, so a column name identifies its table. A typical input is:

Project([order_id, name],
 Filter([amount > 100, country = 'US'],
 Join(Scan(orders, [order_id, customer_id, amount, status, created_at]),
 Scan(customers, [cust_id, name, country, signup_date]),
 left_key=customer_id, right_key=cust_id)))

Clarifying Questions

• Can conjuncts reference more than one column (for example, comparing two columns), and can projections compute expressions?
• Can the plan contain outer joins, aggregations or sorts, or only the four node types above?
• Should the optimizer mutate the tree in place or return a new tree?
• Must the optimized plan preserve the exact column order and row order of the original output?

Part 1 — Column pruning

Implement a pass that rewrites the plan so every Scan reads only the columns that some node above it actually needs, without changing the query's output.

What This Part Should Cover

• Which columns each node type needs from its child, including filter columns and join keys
• Keeping the root's output unchanged while shrinking every scan
• Runnable code, and its cost in terms of plan size

Part 2 — Filter pushdown

Implement a pass that moves filter conditions as close to the scans as possible (predicate pushdown): through projections and into the correct side of joins, without changing the query result.

Clarifying Questions for this Part

• If a conjunct references columns from both sides of a join, where should it stay?

What This Part Should Cover

• Splitting conjunctions, and deciding for each conjunct how far down it may legally move
• Why pushing into either side of an inner join is safe, and when it would not be safe for other join types
• How the two passes compose, plus a runnable check that the optimized plan returns the same rows