Katkat
Back

Optimize a SQL query plan tree

DatabricksSoftware Engineer · Technical Screen
HardSoftware EngineerTechnical ScreenCoding & Algorithms

Last updated

This problem tests your ability to manipulate an Abstract Syntax Tree (AST) representing a relational query plan to reduce data processing costs through predicate pushdown and projection pruning.

Query Plan Optimization

In a database engine, a query plan is represented as a tree of relational operators: Scan, Filter, Project, and Join. To optimize performance, we want to push Filter operations as far down the tree as possible (closer to the Scan) and move Project operations down to prune unnecessary columns early.

Task

Given a tree of nodes, implement a function optimize_plan(root) that performs two passes on the tree:

  1. Predicate Pushdown: Move Filter nodes below Join nodes if the filter predicate only references columns from one side of the join.
  2. Projection Pruning: Move Project nodes down the tree, discarding columns that are not used by the parent nodes or the current filter/join condition.

Node Structure

  • Scan(table_name, columns): Leaves of the tree.
  • Filter(predicate, child): Filters rows based on a condition.
  • Project(columns, child): Selects a subset of columns.
  • Join(condition, left_child, right_child): Joins two relations.

Example

Consider the plan: Project([a, b], Filter(a > 10, Join(c == d, Scan(T1, [a, c]), Scan(T2, [b, d]))))

After optimization, the Filter(a > 10) should be pushed below the Join onto the Scan(T1, [a, c]) side because a is only present in T1.

Constraints

  • The tree is a standard binary tree where Join is the only node with two children.
  • Assume predicate is a simple string identifying the column involved (e.g., "a > 10").
  • You must return the root of the modified tree.
class Node:
    def __init__(self, type, children=None, **kwargs):
        self.type = type
        self.children = children or []
        self.kwargs = kwargs # Stores columns, predicate, condition, etc.

Hints

Hint 1

Use a recursive post-order traversal to process children before modifying the current node.

Hint 2

For predicate pushdown, check if the predicate string contains only column names that exist in the schema of the left or right child respectively.

Hint 3

For projection pruning, maintain a 'required columns' set that bubbles up from the root to the leaves, allowing you to drop unused columns at every `Scan` or `Project` node.

Follow-up questions

  • How would you handle a cross-join (Cartesian product) if the predicate involves columns from both tables?
  • What changes are required if the plan includes common sub-expressions, effectively making the plan a Directed Acyclic Graph (DAG) instead of a tree?
  • How do cost-based optimization (CBO) approaches differ from the rule-based optimizations implemented here?