Custom Plan Nodes

If there are inaccuracies discovered with this documentation, please submit a GitHub issue.

Introduction

The PartiQL Library allows for the integration of custom plan nodes.

Who is this for?

This usage guide is aimed at developers who want to add their own plan nodes to a PartiQL Plan.

Prerequisites

To get the APIs discussed in this usage guide, please take a dependency on the SPI package.

build.gradle.kts

dependencies {
    implementation("org.partiql:partiql-plan:1.+")
}

Custom Plan Node Basics

Implementing a custom plan node is as simple as extending the plan’s base interfaces, specifically Rel (or RelBase) and Rex. These represent relational algebra expressions and scalar expressions, respectively.

That being said, it’s not entirely as simple as creating the implementation. Depending on what your application does to the node, you may need to:

Transform an existing plan to use your node.
Create a custom visitor (specific to your application) that can visit your node.
Create a custom plan rewriter (specific to your application) that can transform plans may contain your node.
Integrate custom compilation strategies to handle your custom node.

In the following sections, we will create a node, named ScanFilter, whose functionality is more complex than a normal RelScan. This custom node — instead of scanning arbitrary PartiQL values — can only scan columnar data. On top of this, it has pushed down a predicate to apply a filter before passing on the results of the scan to another relation.

Example Relation Node

Node Implementation

In the below code snippet, you’ll see the basics of ScanFilter. It is comprised of a reference to a RexTable and a predicate that will be applied at runtime. During the construction of a ScanFilter, it will also ensure that the table’s underlying data is schema-full and columnar.

ScanFilter

import org.jetbrains.annotations.NotNull;
import org.partiql.plan.Operand;
import org.partiql.plan.OperatorVisitor;
import org.partiql.plan.rel.RelBase;
import org.partiql.plan.rel.RelType;
import org.partiql.plan.rex.Rex;
import org.partiql.plan.rex.RexTable;
import org.partiql.spi.types.PType;
import org.partiql.spi.types.PTypeField;

import java.util.List;

public class ScanFilter extends RelBase {

    private final RexTable input;
    private final Rex predicate;
    private final List<PTypeField> fields;

    public ScanFilter(RexTable input, Rex predicate) {
        this.input = input;
        this.predicate = predicate;
        PType type = this.input.getTable().getSchema();
        assert type.code() == PType.BAG;
        PType elementType = type.getTypeParameter();
        assert elementType.code() == PType.ROW;
        this.fields = elementType.getFields().stream().toList();
    }

    @NotNull
    public Rex getPredicate() {
        return predicate;
    }

    @NotNull
    public RexTable getInput() {
        return input;
    }

    @Override
    public <R, C> R accept(OperatorVisitor<R, C> operatorVisitor, C c) {
        if (operatorVisitor instanceof MyApplicationVisitor<R,C>) {
            return ((MyApplicationVisitor<R, C>) operatorVisitor).visitScanFilter(this, c);
        }
        throw new IllegalStateException("ScanFilter can't be processed by: " + operatorVisitor.getClass().getName());
    }

    @NotNull
    @Override
    protected RelType type() {
        return RelType.of(fields.toArray(new PTypeField[0]));
    }

    @NotNull
    @Override
    protected List<Operand> operands() {
        return List.of(Operand.single(input), Operand.single(predicate));
    }
}

Visitor Implementation

If your application is going to perform analysis (via visiting) on the plan or rewrite the plan (via a PlanRewriter), then you will need to extend the base OperatorVisitor to include your custom nodes.

MyApplicationVisitor.java

import org.jetbrains.annotations.NotNull;
import org.partiql.plan.OperatorVisitor;

public interface MyApplicationVisitor <R,C> extends OperatorVisitor<R, C> {

    default R visitScanFilter(@NotNull ScanFilter rel, C ctx) {
        return this.defaultVisit(rel, ctx);
    }
}

Rewriter Implementation

If your application rewrites plans containing your custom node (via a PlanRewriter), then you’ll need to extend the default PlanRewriter to include your custom visitor and nodes.

MyApplicationRewriter.java

import org.jetbrains.annotations.NotNull;
import org.partiql.plan.Operator;
import org.partiql.plan.OperatorRewriter;
import org.partiql.plan.rex.Rex;
import org.partiql.plan.rex.RexTable;

public abstract class MyApplicationRewriter <C> extends OperatorRewriter<C> implements MyApplicationVisitor<Operator, C> {

    @NotNull
    @Override
    public Operator visitScanFilter(@NotNull ScanFilter rel, C ctx) {
        // Rewrite input
        RexTable input = rel.getInput();
        RexTable input_new = visit(input, ctx, RexTable.class);

        // Rewrite predicate
        Rex predicate = rel.getPredicate();
        Rex predicate_new = visit(predicate, ctx, Rex.class);

        // Rewrite ScanFilter (only if visiting the children produces new classes)
        if (input != input_new || predicate != predicate_new) {
            return new ScanFilter(input_new, predicate_new);
        }
        return rel;
    }
}

Notice how the MyApplicationRewriter extends the OperatorRewriter and implements the MyApplicationVisitor.

Integration

To integrate your custom nodes, all it takes is to create a new plan with your nodes or modify an existing plan to include your nodes.

Your custom visits and rewrites, however, will not be integrated into the PartiQL compilation process. These will exist in your own codebase.

To compile your custom nodes, please see the Compiling Plan Nodes usage guide.