Compiling Plan Nodes

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

Introduction

The PartiQL Library allows for the injection of custom compilation strategies for plan nodes.

Who is this for?

This usage guide is aimed at developers who want to take the PartiQL compiler out-of-the-box, yet, provide custom implementations of either scalar or relational expressions.

Examples of custom scalar or relational expressions include:

  1. An overridden implementation of a standard plan node (i.e. SCAN, FILTER, PROJECT, etc.)

  2. A custom implementation of a custom plan node (e.g. SCAN_FILTER, etc.). See Custom Plan Nodes for more information.

Implementation

Overriding a Node’s Compilation

For the following example, we will be writing a custom compilation strategy for a standard plan node (Scan).

The code snippet below shows how you can create a strategy for a Scan, retrieve its children, compile its children, and create an executable node.

ScanStrategy.java
import org.jetbrains.annotations.NotNull;
import org.partiql.eval.Expr;
import org.partiql.eval.ExprValue;
import org.partiql.eval.Mode;
import org.partiql.eval.compiler.Match;
import org.partiql.eval.compiler.Pattern;
import org.partiql.eval.compiler.Strategy;
import org.partiql.plan.Operand;
import org.partiql.plan.rel.RelScan;

/**
 * Compiles Scan
 */
public class ScanStrategy extends Strategy {

    private static final Pattern pattern = new Pattern(RelScan.class);

    public ScanStrategy() {
        super(pattern);
    }

    @NotNull
    @Override
    public Expr apply(@NotNull Match match, @NotNull Mode mode, @NotNull Callback callback) {
        // Retrieve the Scan
        Operand.Single op = (Operand.Single) match.getOperand(); // For the simplest cases, you can cast to an Operand.Single
        RelScan scan = (RelScan) op.operator; // For the simplest cases, you can cast the operator to the node specified by your pattern.

        // Compile the Children
        ExprValue expr = (ExprValue) callback.apply(scan.getRex());

        // Create an Execution Node. If the mode implicates differing behavior, please take that into consideration.
        return new ScanImpl(expr);
    }
}

Below, you may find the example implementation of ScanImpl, an executable node. Note that the main APIs to implement are open(), close(), hasNext(), and next(). This specific example evaluates the scan’s underlying expression, checks that the data is iterable, and passes the iterator to a class (DatumToRelationIterator) that converts the collection’s elements to a relational Row (containing the single element).

ScanImpl.java
import org.jetbrains.annotations.NotNull;
import org.partiql.eval.Environment;
import org.partiql.eval.ExprRelation;
import org.partiql.eval.ExprValue;
import org.partiql.eval.Mode;
import org.partiql.eval.Row;
import org.partiql.spi.errors.PError;
import org.partiql.spi.errors.PErrorKind;
import org.partiql.spi.errors.PRuntimeException;
import org.partiql.spi.errors.Severity;
import org.partiql.spi.types.PType;
import org.partiql.spi.value.Datum;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;

public class ScanImpl implements ExprRelation {

    private final ExprValue rex;
    private final Mode mode;

    private Iterator<Row> inputIter;

    public ScanImpl(@NotNull ExprValue rex, @NotNull Mode mode) {
        this.rex = rex;
        this.mode = mode;
    }

    @Override
    public void open(@NotNull Environment environment) {
        // Evaluate the data, which may be lazily materialized.
        Datum data = rex.eval(environment);

        // Retrieve the type
        PType dataType = data.getType();
        int typeCode = dataType.code();

        // Check that the type can be iterated upon (array/bag)
        switch (typeCode) {
            case PType.ARRAY:
            case PType.BAG:
                this.inputIter = new DatumToRelationIterator(data.iterator());
                break;
            default:
                // Since it can't be iterated upon, check the compilation mode.
                switch (mode.code()) {
                    // If permissive mode, wrap the value in a bag, and use the bag as the data to iterate
                    case Mode.PERMISSIVE:
                        this.inputIter = new DatumToRelationIterator(Datum.bag(data).iterator());
                        break;
                    // If strict mode, throw a PRuntimeException with the TYPE_UNEXPECTED code.
                    case Mode.STRICT:
                        throw unexpectedTypeForScan(dataType);
                    default:
                        throw new RuntimeException("Mode " + mode + " not supported.");
                }
        }
    }

    @NotNull
    @Override
    public Row next() {
        return inputIter.next();
    }

    @Override
    public boolean hasNext() {
        return inputIter.hasNext();
    }

    @Override
    public void close() {
        this.inputIter = null;
    }

    private PRuntimeException unexpectedTypeForScan(@NotNull PType actualType) {
        return new PRuntimeException(
                new PError(
                        PError.TYPE_UNEXPECTED,
                        Severity.ERROR(),
                        PErrorKind.EXECUTION(),
                        null,
                        new HashMap<>() {{
                            put("EXPECTED_TYPES", new ArrayList<>() {{ add(PType.bag()); add(PType.array()); }});
                            put("ACTUAL_TYPE", actualType);
                        }}
                )
        );
    }

    private static class DatumToRelationIterator implements Iterator<Row> {

        private final Iterator<Datum> elements;

        public DatumToRelationIterator(@NotNull Iterator<Datum> elements) {
            this.elements = elements;
        }

        @Override
        public boolean hasNext() {
            return this.elements.hasNext();
        }

        @Override
        public Row next() {
            Datum[] fields = new Datum[1];
            fields[0] = elements.next();
            return new Row(fields);
        }
    }
}

Compiling of a Custom Plan Node

For the following example, we will be writing a custom compilation strategy for a custom plan node (ScanFilter) discussed in the Custom Plan Nodes usage guide.

The code snippet below shows how you can create a strategy for a ScanFilter, retrieve its children, compile its children, and create an executable node.

ScanFilterStrategy.java
import org.jetbrains.annotations.NotNull;
import org.partiql.eval.Expr;
import org.partiql.eval.ExprRelation;
import org.partiql.eval.ExprValue;
import org.partiql.eval.Mode;
import org.partiql.eval.compiler.Match;
import org.partiql.eval.compiler.Pattern;
import org.partiql.eval.compiler.Strategy;
import org.partiql.plan.Operand;

/**
 * Compiles ScanFilter
 */
public class ScanFilterStrategy extends Strategy {

    private static final Pattern pattern = new Pattern(ScanFilter.class);

    public ScanFilterStrategy() {
        super(pattern);
    }

    @NotNull
    @Override
    public Expr apply(@NotNull Match match, @NotNull Mode mode, @NotNull Callback callback) {
        // Retrieve the ScanFilter
        Operand.Single op = (Operand.Single) match.getOperand(); // For the simplest cases, you can cast to an Operand.Single
        ScanFilter scanFilter = (ScanFilter) op.operator; // For the simplest cases, you can cast the operator to the node specified by your pattern.

        // Compile the Children
        ExprRelation inputRel = (ExprRelation) callback.apply(scanFilter.getInput());
        ExprValue predicate = (ExprValue) callback.apply(scanFilter.getPredicate());

        // Create an Execution Node. If the mode implicates differing behavior, please take that into consideration.
        return new ScanFilterImpl(inputRel, predicate);
    }
}

Now that we have created a strategy, we need to implement the executable node (ScanFilterImpl). The main APIs to implement are open(), close(), hasNext(), and next(). The implementation of the executable ScanFilterImpl may look like this:

ScanFilterImpl.java
import org.jetbrains.annotations.NotNull;
import org.partiql.eval.Environment;
import org.partiql.eval.ExprRelation;
import org.partiql.eval.ExprValue;
import org.partiql.eval.Row;
import org.partiql.spi.types.PType;
import org.partiql.spi.value.Datum;

import java.util.Iterator;

public class ScanFilterImpl implements ExprRelation {

    private final ExprRelation input;
    private final ExprValue predicate;

    private Environment env;
    private Iterator<Row> inputIter;
    private Row _next;

    public ScanFilterImpl(@NotNull ExprRelation input, @NotNull ExprValue predicate) {
        this.input = input;
        this.predicate = predicate;
    }

    @Override
    public void open(@NotNull Environment environment) {
        this.input.open(environment);
    }

    private void peek() {
        while (inputIter.hasNext()) {
            Row inputRow = inputIter.next();
            Environment newEnv = env.push(inputRow);
            Datum result = predicate.eval(newEnv);
            if (!result.isNull() && !result.isMissing() && result.getType().code() == PType.BOOL && result.getBoolean()) {
                this._next = inputRow;
                return;
            }
        }
    }

    @NotNull
    @Override
    public Row next() {
        if (_next == null) {
            peek();
        }
        if (_next == null) {
            throw new RuntimeException("Attempting to call next() when no elements remaining. Please use hasNext().");
        }
        Row _temp = _next;
        _next = null;
        return _temp;
    }

    @Override
    public boolean hasNext() {
        if (_next != null) {
            return true;
        }
        peek();
        return _next != null;
    }

    @Override
    public void close() {
        this.input.close();
        this.env = null;
        this.inputIter = null;
        this._next = null;
    }
}

Integration

Please see the Using the Compiler usage guide.