Exploring HepPlanner for Apache Calcite

Reading Time: 3 minutes

In this blog, we will see different ways to manipulate a rel node tree using a Hep planner. A basic understanding of Apache Calcite is necessary for this. Check out the homepage here https://calcite.apache.org/

What is a HepPlanner?

It is a rule-based planner to transform a relational expression represented as a tree-like structure. It allows us to specify a condition to identify particular nodes of the tree and an operation to perform on them, both of which make a rule.
From a given set of rules, it does an exhaustive iteration on a given tree till no more rules are applicable. There is also a VolcanoPlanner for cost-based optimisation but that’s a topic for another day. This blog will illustrate HepPlanner through an example.

When to use it?

Its ideal for using when the rules need to run in a deterministic order. For example, if rule B is only applicable when rule A is applicable, then rule A is defined before rule B.
For mutually exclusive rules , the ordering doesn’t matter.

Example use case

Suppose we have a simple tree like

ColoredRel(Color=[RED])
  ColoredRel(Color=[GREEN])
     LeafRel(Color=[BLUE])

We’ll showcase 3 different ways to manipulate this tree:-
1) We want all GREEN nodes to be preceded by some node
2) We want all GREEN nodes to always be followed by some node
3) We want to replace all GREEN nodes.

Let us start by defining those nodes and their color.

The Color trait

sealed trait ColorTrait extends RelTrait {
  val color: String
}
object ColorDef extends RelTraitDef[ColorTrait] {
  ...implementation
}
abstract class Color(val color: String) extends ColorTrait {
  override def getTraitDef: RelTraitDef[_ <: RelTrait] = ColorDef
  ...implementation
}
case object Red extends Color("RED")
case object Green extends Color("GREEN")
case object Blue extends Color("BLUE")
case object None extends Color("NONE")

Rel nodes

sealed trait CustomRel extends AbstractRelNode {
  override def explainTerms(pw: RelWriter): RelWriter = {
    val writer = super.explainTerms(pw)
    writer.item("Color", getTraitSet.getTrait(ColorDef))
  }
  ...implementation
}

class ColoredRel(val cluster: RelOptCluster, val traits: RelTraitSet, input: RelNode) extends SingleRel(cluster, traits, input) with CustomRel {
  override def copy(traitSet: RelTraitSet, inputs: util.List[RelNode]): RelNode = new ColoredRel(getCluster, traitSet, sole(inputs))
}

class LeafRel(val cluster: RelOptCluster, val traits: RelTraitSet) extends AbstractRelNode(cluster, traits) with CustomRel {
  override def copy(traitSet: RelTraitSet, inputs: util.List[RelNode]): RelNode =
    if (inputs.isEmpty) new LeafRel(getCluster, traitSet)
    else new ColoredRel(getCluster, traitSet, sole(inputs))
}

We made traits representing Red, Green and Blue and None colour, a tree node as ColoredRel and a leaf node as LeafRel. For simplicity, we’ll restrict to a single child by utilizing a SingleRel.

Next,we introduce a rule for all 3 cases which fires itself on every green coloured node. Note that we replaced original traits otherwise, the planner would keep firing the rule infinitely. If we don’t want this, we can also write an at-most-once semantic rule.

class RuleForGreen(config: RuleForGreen.Config) extends RelRule[RuleForGreen.Config](config) {
  override def onMatch(call: RelOptRuleCall): Unit = config.matchHandler().accept(this, call)
}

object RuleForGreen {

  val PRECEDED_BY_RULE: RelOptRule = withOperandFor(addParent()).toRule
  val SUCCEEDED_BY_RULE: RelOptRule = withOperandFor(addChild()).toRule
  val REPLACED_BY_RULE: RelOptRule = withOperandFor(replace()).toRule

  def withOperandFor(handler: MatchHandler[RuleForGreen]): RelRule.Config =
    new RuleForGreen.Config()
      .withMatchHandler(handler)
      .withRelBuilderFactory(RelFactories.LOGICAL_BUILDER)
      .withOperandSupplier(b0 =>
        b0.operand(classOf[CustomRel]).`trait`(ColorTrait.Green).anyInputs())


  def addParent(): RelRule.MatchHandler[RuleForGreen] =
    (_, call) => {
      val rel: RelNode = call.rel(0)
      val originalTraits = rel.getTraitSet
      val enforcedTraits = originalTraits.replace(ColorTrait.None)
      val enforce = new RuledRel(rel.getCluster, originalTraits, rel.copy(enforcedTraits, rel.getInputs))
      call.transformTo(enforce)
    }

  def addChild(): RelRule.MatchHandler[RuleForGreen] = (_, call) => {
    val rel: RelNode = call.rel(0)
    val originalTraits = rel.getTraitSet
    val enforcedTraits = originalTraits.replace(ColorTrait.None)
    val enforce = new RuledRel(rel.getCluster, originalTraits, rel.getInput(0))
    call.transformTo(rel.copy(enforcedTraits, ImmutableList.of(enforce)))
  }

  def replace(): RelRule.MatchHandler[RuleForGreen] = (_, call) => {
    val rel: RelNode = call.rel(0)
    val desiredTraits = rel.getTraitSet.replace(ColorTrait.None)
    val enforce = new RuledRel(rel.getCluster, desiredTraits, rel.getInput(0))
    call.transformTo(enforce)
  }

  class Config extends RelRule.Config {
    ...implementation
    override def toRule: RelOptRule = new RuleForGreen(this)
  }
}

Running it all:-

def hepPlanner(rules: Set[RelOptRule]): RelOptPlanner = {
    val builder = new HepProgramBuilder()
    rules.foreach(builder.addRuleInstance)
    new HepPlanner(builder.build)
  }

Case 1

val planner = hepPlanner(Set(RuleForGreen.PRECEDED_BY_RULE))
planner.setRoot(SampleCase)
val result = planner.findBestExp()

//OUTPUT
ColoredRel(Color=[RED])
  RuledRel(Color=[GREEN])
    ColoredRel(Color=[NONE])
      LeafRel(Color=[BLUE])

Case 2

val planner = hepPlanner(Set(RuleForGreen.SUCCEEDED_BY_RULE))
planner.setRoot(SampleCase)
val result = planner.findBestExp()

//OUTPUT
ColoredRel(Color=[RED])
  ColoredRel(Color=[NONE])
    RuledRel(Color=[GREEN])
      LeafRel(Color=[BLUE])

Case 3

val planner = hepPlanner(Set(RuleForGreen.REPLACED_BY_RULE))
planner.setRoot(SampleCase)
val result = planner.findBestExp()

//OUTPUT
ColoredRel(Color=[RED])
  RuledRel(Color=[NONE])
    LeafRel(Color=[BLUE])

Benefits of using HepPlanner

The HepPlanner can do much more than optimization. We could transform a rel node tree in unprecedented ways. For example, we could want our own interpretation/execution tree for a relational expression. It can allow us to replace every logical node with an equivalent physical node or enforce some semantics like having some particular traits on a node or re-writing an expression completely.

All we need is to come up with some rules and let the planner do the rest.

Here is the full code example: – https://github.com/knoldus/calcite-rules-example
Next time, we’ll explore VolcanoPlanner using the same example. So stay tuned for more!


Knoldus-blog-footer-image

Leave a Reply