Custom traits in Apache Calcite
Tags: calcite, sql, database
Published: 2021-04-16 15:10:15

# Abstract
Physical properties are an essential part of the optimization process that allows you to explore more alternative plans.

Apache Calcite comes with convention and collation (sort order) properties. Many query engines require custom properties. For example, distributed and heterogeneous engines that we often see in our daily practice need to carefully plan the movement of data between machines and devices, which requires a custom property to describe data location.

In this blog post, we will explore how to define, register and enforce a custom property, also known as a *trait*, with Apache Calcite cost-based optimizer.

# Physical Properties
We start our journey by looking at the example of common physical property - sort order.

Query optimizers work with relational operators, such as `Scan`, `Project`, `Filter`, and `Join`. During the optimization, an operator may require its input to satisfy a specific condition. To check whether the condition is satisfied, operators may expose **physical properties** - plain values associated with an operator. Operators may compare the desired and actual properties of their inputs and enforce the desired property by injecting a special **enforcer** operator on top of the input.

Consider the join operator `t1 JOIN t2 ON t1.a = t2.b`. We could use a merge join if both inputs are sorted on their join attributes, `t1.a` and `t2.b`, respectively. We may define the collation property for every operator, describing the sort order of produced rows:
```
Join[t1.a=t2.b]
  Input[t1]      [SORTED by a]
  Input[t2]      [NOT SORTED]
```

The merge join operator may enforce the sorting on `t1.a` and `t2.b` on its inputs. Since the first input is already sorted on `t1.a`, it remains unchanged. The second input is not sorted, so the enforcer `Sort` operator is injected, making a merge join possible:
```
MergeJoin[t1.a=t2.b]  
  Input[t1]           [SORTED by t1.a]
  Sort[t2.a]          [SORTED by t2.b]
    Input[t2]         [NOT SORTED]
```

# Apache Calcite API
In Apache Calcite, properties are defined by the `RelTrait` and `RelTraitDef` classes. `RelTrait` is a concrete value of the property. RelTraitDef is a property definition, which describes the property name, expected Java class of the property, the default value of the property, and how to enforce the property. Property definitions are registered in the planner via the `RelOptPlanner.addRelTraitDef` method. The planner will ensure that every operator has a specific value for every registered property definition, whether the default or not.
All properties of a node are organized in an immutable data structure `RelTraitSet`. This class has convenient methods to add and update properties with copying semantics. You may access the properties of a concrete operator using the `RelOptNode.getTraitSet` method.
To enforce a specific property on the operator during planning, you should do the following from within the rule:
1. Get the current properties of a node using `RelOptNode.getTraitSet` method.
1. Create a new instance of `RelTraitSet` with updated properties.
1. Enforce the properties by calling the `RelOptRule.convert` method.

Finally, before invoking the planner program, you may define the desired properties of the root operator of the optimized relational tree. After the optimization, the planner will either return the operator that satisfies these properties or throw an exception.

Internally, the Apache Calcite enforces properties by adding a special `AbstractConverter` operator with the desired traits on top of the target operator.

```
AbstractConverter [SORTED by a]
  Input[t2]       [NOT SORTED]
```

To transform the `AbstractConverter` into a real enforcer node, such as `Sort`, you should add the built-in `ExpandConversionRule` rule to your optimization program. This rule will attempt to expand the `AbstractConverter` into a sequence of enforcers to satisfy the desired traits. We have only one unsatisfied property in our example, so the converter expands into a single `Sort` operator.

```
Sort[t2.a]        [SORTED by a]
  Input[t2]       [NOT SORTED]
```

You may use your custom expansion rule if needed. See Apache Flink [custom rule](https://github.com/apache/flink/blob/release-1.12/flink-table/flink-table-planner-blink/src/main/scala/org/apache/flink/table/planner/plan/rules/physical/FlinkExpandConversionRule.scala) as an example.

# Custom Property
As we understand the purpose of properties and which Apache Calcite API to use, we will define, register, and enforce our custom property.

Consider that we have a distributed database, where every relational operator might be distributed between nodes in one of two ways:

1. `PARTITIONED` - relation is partitioned between nodes. Every tuple (row) resides on one of the nodes. An example is a typical distributed data structure.
1. `SINGLETON` - relation is located on a single node. An example is a cursor that delivers the final result to the user application.

In our example, we would like to ensure that the top operator always has a `SINGLETON` distribution, simulating the results' delivery to a single node.

## Enforcer
First, we define the enforcer operator. To ensure the `SINGLETON` distribution, we need to move from all nodes to a single node. In distributed databases, data movement operators are often called `Exchange`. The minimal requirement for a custom operator in Apache Calcite is to define the constructor and the `copy` method.
```java
public class ExchangeRel extends SingleRel {
    public RedistributeRel(
        RelOptCluster cluster,
        RelTraitSet traits,
        RelNode input
    ) {
        super(cluster, traits, input);
    }

    @Override
    public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
        return new ExchangeRel(getCluster(), traitSet, inputs.get(0));
    }
}
```

## Trait
Next, we define our custom trait and trait definition. Our implementation must adhere to the following rules:

1. The trait must refer to a common trait definition instance in the method `getTraitDef`.
1. The trait must override the `satisfies` method to define whether the current trait satisfies the target trait. If not, the enforcer will be used.
1. The trait definition must declare the expected Java class of the trait in the `getTraitClass` method.
1. The trait definition must declare the default value of the trait in the `getDefault` method.
1. The trait definition must implement the method `convert`, which Apache Calcite will invoke to create the enforcer if the current trait doesn't satisfy the desired trait. If there is no valid conversion between traits, `null` should be returned.

Below is the source code of our trait. We define two concrete values, `PARTITIONED` and `SINGLETON`. We also define the special value `ANY`, which we use as the default. We say that both `PARTITIONED` and `SINGLETON` satisfy `ANY` but `PARTITIONED` and `SINGLETON` do not satisfy each other.

```java
public class Distribution implements RelTrait {

    public static final Distribution ANY = new Distribution(Type.ANY);
    public static final Distribution PARTITIONED = new Distribution(Type.PARTITIONED);
    public static final Distribution SINGLETON = new Distribution(Type.SINGLETON);

    private final Type type;

    private Distribution(Type type) {
        this.type = type;
    }

    @Override
    public RelTraitDef getTraitDef() {
        return DistributionTraitDef.INSTANCE;
    }

    @Override
    public boolean satisfies(RelTrait toTrait) {
        Distribution toTrait0 = (Distribution) toTrait;

        if (toTrait0.type == Type.ANY) {
            return true;
        }

        return this.type.equals(toTrait0.type);
    }

    enum Type {
        ANY,
        PARTITIONED,
        SINGLETON
    }
}
```

Our trait definition defines the `convert` function, which injects the `ExchangeRel` enforcer if the current property doesn't satisfy the target one.

```java
public class DistributionTraitDef extends RelTraitDef<Distribution> {

    public static DistributionTraitDef INSTANCE = new DistributionTraitDef();

    private DistributionTraitDef() {
        // No-op.
    }

    @Override
    public Class<Distribution> getTraitClass() {
        return Distribution.class;
    }

    @Override
    public String getSimpleName() {
        return "DISTRIBUTION";
    }

    @Override
    public RelNode convert(
        RelOptPlanner planner,
        RelNode rel,
        Distribution toTrait,
        boolean allowInfiniteCostConverters
    ) {
        Distribution fromTrait = rel.getTraitSet().getTrait(DistributionTraitDef.INSTANCE);

        if (fromTrait.satisfies(toTrait)) {
            return rel;
        }

        return new ExchangeRel(
            rel.getCluster(),
            rel.getTraitSet().plus(toTrait),
            rel
        );
    }

    @Override
    public boolean canConvert(
        RelOptPlanner planner,
        Distribution fromTrait,
        Distribution toTrait
    ) {
        return true;
    }

    @Override
    public Distribution getDefault() {
        return Distribution.ANY;
    }
}
```

You would likely have more distribution types, dedicated distribution columns, and different exchange types in production implementations. You may refer to Apache Flink as an example of a real [distribution trait](https://github.com/apache/flink/blob/release-1.12/flink-table/flink-table-planner-blink/src/main/scala/org/apache/flink/table/planner/plan/trait/FlinkRelDistribution.scala#L63).

## Putting It All Together
Let's see the new trait in action. The complete source code is available [here](https://github.com/querifylabs/querifylabs-blog/tree/main/02-custom-calcite-trait).

First, we create a schema with a couple of tables - one with `PARTITIONED` distribution and another with `SINGLETON` distribution. We use custom table and schema implementation, similar to the ones we used in the previous [blog post](https://www.querifylabs.com/blog/assembling-a-query-optimizer-with-apache-calcite).

```java
// Table with PARTITIONED distribution.
Table table1 = Table.newBuilder("table1", Distribution.PARTITIONED)
  .addField("field", SqlTypeName.DECIMAL).build();

// Table with SINGLETON distribution.
Table table2 = Table.newBuilder("table2", Distribution.SINGLETON)
  .addField("field", SqlTypeName.DECIMAL).build();

Schema schema = Schema.newBuilder("schema").addTable(table1).addTable(table2).build();
```

Then we create a planner instance and register our custom trait definition in it.

```java
VolcanoPlanner planner = new VolcanoPlanner();

planner.addRelTraitDef(ConventionTraitDef.INSTANCE);
planner.addRelTraitDef(DistributionTraitDef.INSTANCE);
```

Finally, we create a table scan operator for each of our tables and enforce the `SINGLETON` distribution. Notice that we use the aforementioned `ExpandConversionRule` in our optimization program. Otherwise, the enforcement will not work.

```java
// Use the built-in rule that will expand abstract converters.
RuleSet rules = RuleSets.ofList(AbstractConverter.ExpandConversionRule.INSTANCE);

// Prepare the desired traits with the SINGLETON distribution.
RelTraitSet desiredTraits = node.getTraitSet().plus(Distribution.SINGLETON);
        
// Use the planner to enforce the desired traits
RelNode optimizedNode = Programs.of(rules).run(
    planner,
    node,
    desiredTraits,
    Collections.emptyList(),
    Collections.emptyList()
);
```

Now we run the [TraitTest](https://github.com/querifylabs/querifylabs-blog/commit/1bb7d1c8942e7df00916d20de0e51f2a1965ac23#diff-5c30c410d3e7db69b7b8375de2802b3c04f94799dc4b30239fba532001d8eea4) from the sample project to see this in action. For the `PARTITIONED` table, the planner has added the `ExchangeRel` to enforce the `SINGLETON` distribution.

```
BEFORE:
2:LogicalTableScan(table=[[schema, partitioned]])

AFTER:
7:ExchangeRel
  2:LogicalTableScan(table=[[schema, partitioned]])
```

But the table with the `SINGLETON` distribution remains unchanged because it already has the desired distribution.

```
BEFORE:
0:LogicalTableScan(table=[[schema, singleton]])

AFTER:
0:LogicalTableScan(table=[[schema, singleton]])
```

Congratulations! Our custom property is ready.

# Summary
Physical properties are an important concept in query optimization that allows you to explore more alternative plans.

In this blog post, we demonstrated how to define the custom physical property in Apache Calcite. We created custom `RelTraitDef` and `RelTrait` classes, registered them in the planner, and used the custom operator to enforce the desired value of the property. 

However, we omitted one crucial question - how to **propagate** properties between operators? It turns out, Apache Calcite cannot do this well, and you will have to make a tough decision choosing between several non-ideal solutions. We will discuss property propagation in detail in future posts. Stay tuned!

We are always ready to help you with your SQL query optimizer design. Just [let us know](https://www.querifylabs.com/#contact-form).