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: ```text 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. ```hint Direction of information Decide whether the set of needed columns is known at the top of the tree or at the leaves, and let the traversal carry it in that direction. ``` #### 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. ```hint Split before you push A filter with several conditions does not have to move as one unit. ``` #### 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 ### What a Strong Answer Covers - Recognizing standard logical rewrites without being told which ones to apply - Clean recursive tree rewrites that return a new plan and preserve the output schema - A correctness argument for each rewrite, including join semantics - Executable code with a tiny evaluator or test that compares results before and after - The complexity of each pass relative to the number of nodes and columns ### Follow-up Questions - How would you push a filter through a projection that computes expressions, such as `total = price * qty`? - Which conjuncts could become join conditions, and why does that matter for performance? - How would you decide the order in which rewrite rules run, and when to stop?