Spark Catalyst

catalyst/
├── analysis/           # Query analysis and resolution
├── catalog/           # Catalog management and metadata
├── expressions/       # Expression definitions and evaluation
├── optimizer/         # Query optimisation rules
├── parser/            # SQL parsing
├── planning/          # Query planning strategies
├── plans/             # Query plan representations
│   ├── logical/       # Logical plan operators
│   └── physical/      # Physical plan operators
├── rules/             # Rule execution framework
├── trees/             # Tree node infrastructure
├── types/             # Data type utilities
└── util/              # Utility classes

abstract class TreeNode[BaseType <: TreeNode[BaseType]] extends Product {
  // Children of this node.
  def children: Seq[BaseType]

  // Transform this tree by applying a function to all nodes.
  def transform(rule: PartialFunction[BaseType, BaseType]): BaseType

  // Transform all nodes bottom-up.
  def transformDown(rule: PartialFunction[BaseType, BaseType]): BaseType

  // Transform all nodes top-down.
  def transformUp(rule: PartialFunction[BaseType, BaseType]): BaseType

  // Fast equality check.
  def fastEquals(other: TreeNode[_]): Boolean

  // Tree pattern matching for efficient traversal.
  def treePatternBits: BitSet
}

plan.transformUp {
  case Filter(condition, child) if isAlwaysTrue(condition) =>
    child
}

plan.transformDown {
  case Project(projectList, child) =>
    // Transform project first, then children
    optimiseProject(projectList, child)
}

val filters = plan.collect {
  case f @ Filter(_, _) => f
}

abstract class Expression extends TreeNode[Expression] {
  // Can this expression be evaluated at query planning time?
  def foldable: Boolean

  // Does this expression always return the same result for fixed inputs?
  def deterministic: Boolean

  // Can this expression evaluate to null?
  def nullable: Boolean

  // Data type of the expression result.
  def dataType: DataType

  // Evaluate this expression given an input row.
  def eval(input: InternalRow): Any

  // Generate code for this expression.
  def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode

  // Attributes referenced by this expression.
  def references: AttributeSet
}

Literal(42)                    // Integer literal
Literal("hello")               // String literal
Literal(null, StringType)      // Null literal

AttributeReference("name", StringType, nullable = true)()
AttributeReference("age", IntegerType, nullable = false)()

EqualTo(left, right)           // left = right
GreaterThan(left, right)       // left > right
LessThanOrEqual(left, right)   // left <= right
And(left, right)               // left AND right
Or(left, right)                // left OR right
Not(child)                     // NOT child

Add(left, right)               // left + right
Subtract(left, right)          // left - right
Multiply(left, right)          // left * right
Divide(left, right)            // left / right

Substring(str, pos, len)       // substring(str, pos, len)
Upper(child)                   // upper(child)
Lower(child)                   // lower(child)
Concat(children)               // concat(child1, child2, ...)

Cast(child, targetType)        // cast(child as targetType)

abstract class QueryPlan[PlanType <: QueryPlan[PlanType]]
  extends TreeNode[PlanType] {

  // Output schema of this plan node.
  def output: Seq[Attribute]

  // Set of output attributes.
  def outputSet: AttributeSet

  // Set of attributes from all children.
  def inputSet: AttributeSet

  // Attributes produced by this node.
  def producedAttributes: AttributeSet

  // Attributes referenced by expressions.
  def references: AttributeSet

  // Attributes referenced but not provided by children.
  def missingInput: AttributeSet

  // All expressions in this plan node.
  def expressions: Seq[Expression]

  // Transform expressions in this plan.
  def transformExpressions(rule: PartialFunction[Expression, Expression]): PlanType
}

// Read from a relation.
LogicalRelation(relation, output, catalogTable)

// Local in-memory data.
LocalRelation(output, data)

// Empty relation.
EmptyRelation(output)

// Project specific columns.
Project(projectList: Seq[NamedExpression], child: LogicalPlan)

// Filter rows.
Filter(condition: Expression, child: LogicalPlan)

// Select distinct rows.
Distinct(child: LogicalPlan)

// Group by and aggregate.
Aggregate(
  groupingExpressions: Seq[Expression],
  aggregateExpressions: Seq[NamedExpression],
  child: LogicalPlan
)

// Join two relations.
Join(
  left: LogicalPlan,
  right: LogicalPlan,
  joinType: JoinType,
  condition: Option[Expression],
  hint: JoinHint
)

// Sort rows.
Sort(
  order: Seq[SortOrder],
  global: Boolean,
  child: LogicalPlan
)

// Limit number of rows.
Limit(limitExpr: Expression, child: LogicalPlan)

// Union of two relations.
Union(children: Seq[LogicalPlan])

// Intersection.
Intersect(left: LogicalPlan, right: LogicalPlan, isAll: Boolean)

// Difference.
Except(left: LogicalPlan, right: LogicalPlan, isAll: Boolean)

abstract class InternalRow extends SpecializedGetters {
  // Number of fields in this row.
  def numFields: Int

  // Check if field at ordinal is null.
  def isNullAt(ordinal: Int): Boolean

  // Get value at ordinal.
  def get(ordinal: Int, dataType: DataType): Any

  // Specialised getters for primitive types.
  def getBoolean(ordinal: Int): Boolean
  def getByte(ordinal: Int): Byte
  def getShort(ordinal: Int): Short
  def getInt(ordinal: Int): Int
  def getLong(ordinal: Int): Long
  def getFloat(ordinal: Int): Float
  def getDouble(ordinal: Int): Double
  def getDecimal(ordinal: Int, precision: Int, scale: Int): Decimal
  def getUTF8String(ordinal: Int): UTF8String

  // Update value at ordinal.
  def update(ordinal: Int, value: Any): Unit

  // Set field to null.
  def setNullAt(ordinal: Int): Unit

  // Create a copy of this row.
  def copy(): InternalRow

  // Convert to Scala sequence.
  def toSeq(schema: StructType): Seq[Any]
}

// Basic rule.
object MyOptimisationRule extends Rule[LogicalPlan] {
  def apply(plan: LogicalPlan): LogicalPlan = plan.transformUp {
    case Filter(condition, child) if isAlwaysTrue(condition) =>
      child
  }
}

// Configurable rule.
case class MyParameterisedRule(conf: SQLConf) extends Rule[LogicalPlan] {
  def apply(plan: LogicalPlan): LogicalPlan = {
    if (conf.myFeatureEnabled) {
      optimisePlan(plan)
    } else {
      plan
    }
  }
}

abstract class RuleExecutor[TreeType <: TreeNode[_]] {
  // Define batches of rules to execute.
  protected def batches: Seq[Batch]

  // Execute all batches on the plan.
  def execute(plan: TreeType): TreeType
}

// Batch execution strategies.
abstract class Strategy
case class Once extends Strategy
case class FixedPoint(maxIterations: Int) extends Strategy

object MyOptimiser extends RuleExecutor[LogicalPlan] {
  val batches = Seq(
    Batch("Normalisation", Once,
      EliminateSubqueryAliases,
      RemoveRedundantAliases
    ),
    Batch("Operator Optimisation", FixedPoint(100),
      PushDownPredicate,
      ConstantFolding,
      ColumnPruning
    ),
    Batch("Join Reordering", Once,
      CostBasedJoinReorder
    )
  )
}

val analyzer = new Analyzer(catalogManager)
val analysedPlan = analyzer.execute(unresolvedPlan)

// Push filters below projections.
PushDownPredicate

// Push filters into join conditions.
PushPredicateThroughJoin

// Push filters to data sources.
PushDownPredicates

// Eliminate unnecessary columns.
ColumnPruning

// Combine adjacent projections.
CollapseProject

// Evaluate constant expressions.
ConstantFolding

// Simplify expressions.
SimplifyConditionals
SimplifyCasts

// Reorder joins for better performance.
CostBasedJoinReorder

// Eliminate redundant joins.
EliminateOuterJoin

// Decorrelate correlated subqueries.
DecorrelateInnerQuery

// Merge scalar subqueries.
MergeScalarSubqueries

case class ExprCode(
  code: Block,              // Generated code block
  isNull: ExprValue,        // Variable for null check
  value: ExprValue          // Variable for result value
)

trait CodegenFallback extends Expression {
  // Fallback to interpreted evaluation.
  protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
    ctx.references += this
    val objectTerm = ctx.addReferenceObj("expression", this)
    ExprCode(
      code = code"""
        boolean ${ev.isNull} = true;
        ${CodeGenerator.javaType(dataType)} ${ev.value} =
          ${CodeGenerator.defaultValue(dataType)};
        Object result = $objectTerm.eval(${ctx.INPUT_ROW});
        if (result != null) {
          ${ev.isNull} = false;
          ${ev.value} = (${CodeGenerator.boxedType(dataType)}) result;
        }
      """,
      isNull = ev.isNull,
      value = ev.value
    )
  }
}

// Generate unsafe projection from expressions.
val projection = GenerateUnsafeProjection.generate(expressions)
val result = projection(inputRow)

// Generate mutable projection.
val mutableProjection = GenerateMutableProjection.generate(expressions)
val outputRow = mutableProjection(inputRow)

// Generate safe projection.
val safeProjection = GenerateSafeProjection.generate(expressions)

// Generate predicate for filter condition.
val predicate = GeneratePredicate.generate(condition)
val passes = predicate.eval(row)

// Generate row comparator.
val ordering = GenerateOrdering.generate(sortOrders)
val compareResult = ordering.compare(row1, row2)

BooleanType
ByteType
ShortType
IntegerType
LongType
FloatType
DoubleType
StringType
BinaryType
DateType
TimestampType
TimestampNTZType

// Array type.
ArrayType(elementType: DataType, containsNull: Boolean)

// Map type.
MapType(keyType: DataType, valueType: DataType, valueContainsNull: Boolean)

// Struct type.
StructType(fields: Seq[StructField])

// Struct field.
StructField(name: String, dataType: DataType, nullable: Boolean, metadata: Metadata)

DecimalType(precision: Int, scale: Int)
DecimalType.SYSTEM_DEFAULT  // Decimal(38, 18)

case class MyCustomFunction(child: Expression) extends UnaryExpression {
  // Define output type.
  override def dataType: DataType = StringType

  // Can the result be null?
  override def nullable: Boolean = child.nullable

  // Evaluate the expression.
  override def eval(input: InternalRow): Any = {
    val value = child.eval(input)
    if (value == null) {
      null
    } else {
      // Custom logic here.
      UTF8String.fromString(value.toString.toUpperCase)
    }
  }

  // Generate code for this expression.
  override protected def doGenCode(
      ctx: CodegenContext,
      ev: ExprCode): ExprCode = {
    val childGen = child.genCode(ctx)
    ev.copy(code = code"""
      ${childGen.code}
      boolean ${ev.isNull} = ${childGen.isNull};
      ${CodeGenerator.javaType(dataType)} ${ev.value} =
        ${CodeGenerator.defaultValue(dataType)};
      if (!${ev.isNull}) {
        ${ev.value} = UTF8String.fromString(
          ${childGen.value}.toString().toUpperCase());
      }
    """)
  }

  // Override for pretty printing.
  override def prettyName: String = "my_custom_function"
}

object EliminateRedundantCasts extends Rule[LogicalPlan] {
  def apply(plan: LogicalPlan): LogicalPlan = plan.transformAllExpressions {
    case Cast(child, dataType, _, _) if child.dataType == dataType =>
      // Remove cast if types match.
      child
  }
}

plan match {
  case Project(projectList, child) =>
    // Handle projection.

  case Filter(condition, Project(projectList, child)) =>
    // Handle filter over projection.

  case Join(left, right, joinType, Some(condition), _) =>
    // Handle join with condition.

  case Aggregate(grouping, aggregates, child) =>
    // Handle aggregation.

  case _ =>
    // Default case.
}

// Find all attribute references.
val attributes = expr.collect {
  case a: AttributeReference => a
}

// Find all subqueries.
val subqueries = expr.collect {
  case s: SubqueryExpression => s
}

// Transform specific expression types.
val transformed = expr.transformUp {
  case Add(Literal(0, _), right) => right
  case Add(left, Literal(0, _)) => left
}

// Resolve attribute by name.
def resolve(attrName: String, input: LogicalPlan): Option[Attribute] = {
  input.output.find(_.name == attrName)
}

// Resolve with qualifier.
def resolveQualified(
    qualifier: Seq[String],
    attrName: String,
    input: LogicalPlan): Option[Attribute] = {
  input.output.find { attr =>
    attr.qualifier.startsWith(qualifier) && attr.name == attrName
  }
}

// Create schema from attributes.
val schema = StructType(attributes.map { attr =>
  StructField(attr.name, attr.dataType, attr.nullable, attr.metadata)
})

// Convert schema to attributes.
val attributes = schema.toAttributes

// Add column to schema.
val newSchema = schema.add("newColumn", StringType, nullable = true)

// Drop column from schema.
val reducedSchema = StructType(schema.filterNot(_.name == "dropColumn"))

test("my custom function") {
  val input = Literal("hello")
  val expr = MyCustomFunction(input)

  // Test evaluation.
  assert(expr.eval(null) == UTF8String.fromString("HELLO"))

  // Test properties.
  assert(expr.dataType == StringType)
  assert(expr.foldable)
  assert(expr.deterministic)
}

test("eliminate redundant casts") {
  val plan = Project(
    Seq(Alias(Cast(AttributeReference("x", IntegerType)(), IntegerType), "y")()),
    testRelation
  )

  val optimised = EliminateRedundantCasts(plan)

  val expected = Project(
    Seq(Alias(AttributeReference("x", IntegerType)(), "y")()),
    testRelation
  )

  comparePlans(optimised, expected)
}