IGNITE-27822 Support recursive query for Calcite engine

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/LogicalRelImplementor.java

@@ -22,7 +22,9 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.UUID;
 import java.util.function.BiPredicate;
@@ -74,6 +76,9 @@
 import org.apache.ignite.internal.processors.query.calcite.exec.rel.Node;
 import org.apache.ignite.internal.processors.query.calcite.exec.rel.Outbox;
 import org.apache.ignite.internal.processors.query.calcite.exec.rel.ProjectNode;
+import org.apache.ignite.internal.processors.query.calcite.exec.rel.RecursiveCteState;
+import org.apache.ignite.internal.processors.query.calcite.exec.rel.RecursiveTableSpoolNode;
+import org.apache.ignite.internal.processors.query.calcite.exec.rel.RepeatUnionNode;
 import org.apache.ignite.internal.processors.query.calcite.exec.rel.ScanNode;
 import org.apache.ignite.internal.processors.query.calcite.exec.rel.ScanStorageNode;
 import org.apache.ignite.internal.processors.query.calcite.exec.rel.ScanTableRowNode;
@@ -100,8 +105,11 @@
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteNestedLoopJoin;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteProject;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteReceiver;
+import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRecursiveTableScan;
+import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRecursiveTableSpool;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRel;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRelVisitor;
+import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRepeatUnion;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteSender;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteSort;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteSortedIndexSpool;
@@ -162,6 +170,9 @@ public class LogicalRelImplementor<Row> implements IgniteRelVisitor<Node<Row>> {
     /** */
     private final ExpressionFactory<Row> expressionFactory;
 
+    /** Query-local recursive CTE states, keyed by transient table identifier. */
+    private final Map<String, RecursiveCteState<Row>> recursiveStates = new HashMap<>();
+
     /**
      * @param ctx Root context.
      * @param affSrvc Affinity service.
@@ -614,6 +625,11 @@ private boolean hasExchange(RelNode rel) {
         }
     }
 
+    /** {@inheritDoc} */
+    @Override public Node<Row> visit(IgniteRecursiveTableScan rel) {
+        return new ScanNode<>(ctx, rel.getRowType(), recursiveState(rel.stateId()).current());
+    }
+
     /** {@inheritDoc} */
     @Override public Node<Row> visit(IgniteValues rel) {
         List<RexLiteral> vals = Commons.flat(Commons.cast(rel.getTuples()));
@@ -634,6 +650,17 @@ private boolean hasExchange(RelNode rel) {
         return node;
     }
 
+    /** {@inheritDoc} */
+    @Override public Node<Row> visit(IgniteRepeatUnion rel) {
+        RecursiveCteState<Row> state = recursiveState(rel.stateId());
+        RepeatUnionNode<Row> node = new RepeatUnionNode<>(ctx, rel.getRowType(), state, rel.iterationLimit());
+
+        state.clear();
+        node.register(F.asList(visit(rel.getLeft()), visit(rel.getRight())));
+
+        return node;
+    }
+
     /** {@inheritDoc} */
     @Override public Node<Row> visit(IgniteLimit rel) {
         long offset = validateAndGetOffset(rel.offset(), LimitNode.OFFSET_DEFAULT);
@@ -686,6 +713,16 @@ private long validateAndGetFetch(RexNode node, long defaultVal) {
         return node;
     }
 
+    /** {@inheritDoc} */
+    @Override public Node<Row> visit(IgniteRecursiveTableSpool rel) {
+        RecursiveTableSpoolNode<Row> node =
+            new RecursiveTableSpoolNode<>(ctx, rel.getRowType(), recursiveState(rel.stateId()));
+
+        node.register(visit(rel.getInput()));
+
+        return node;
+    }
+
     /** {@inheritDoc} */
     @Override public Node<Row> visit(IgniteSortedIndexSpool rel) {
         RelCollation collation = rel.collation();
@@ -1025,6 +1062,11 @@ public <T extends Node<Row>> T go(IgniteRel rel) {
         return (T)visit(rel);
     }
 
+    /** Returns the state shared by the repeat union, spool and scan of one recursive CTE. */
+    private RecursiveCteState<Row> recursiveState(String stateId) {
+        return recursiveStates.computeIfAbsent(stateId, key -> new RecursiveCteState<>(ctx));
+    }
+
     /** */
     private ScanStorageNode<Row> createStorageScan(
         String storageName,

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/rel/RecursiveCteState.java

@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ignite.internal.processors.query.calcite.exec.rel;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.ignite.internal.processors.query.calcite.exec.ExecutionContext;
+import org.apache.ignite.internal.processors.query.calcite.exec.tracker.RowTracker;
+import org.apache.ignite.internal.util.GridUnsafe;
+
+/** Query-local current and next deltas of a recursive CTE. */
+public class RecursiveCteState<Row> {
+    /** Rows visible to the recursive table scan. */
+    private List<Row> current = Collections.emptyList();
+
+    /** Rows produced by the active seed or recursive term. */
+    private List<Row> next;
+
+    /** Memory tracker for rows in the current delta. */
+    private RowTracker<Row> currentMemoryTracker;
+
+    /** Memory tracker for rows in the next delta. */
+    private RowTracker<Row> nextMemoryTracker;
+
+    /** */
+    public RecursiveCteState(ExecutionContext<Row> ctx) {
+        currentMemoryTracker = ctx.createNodeMemoryTracker(GridUnsafe.OBJ_REF_SIZE);
+        nextMemoryTracker = ctx.createNodeMemoryTracker(GridUnsafe.OBJ_REF_SIZE);
+    }
+
+    /** Starts collecting the next delta. */
+    public void beginWrite() {
+        assert next == null;
+
+        next = new ArrayList<>();
+    }
+
+    /** Adds a row to the next delta. */
+    public void add(Row row) {
+        assert next != null;
+
+        next.add(row);
+        nextMemoryTracker.onRowAdded(row);
+    }
+
+    /** Makes the collected delta visible to recursive scans. */
+    public void commit() {
+        assert next != null;
+
+        currentMemoryTracker.reset();
+        current = next;
+        next = null;
+
+        RowTracker<Row> tracker = currentMemoryTracker;
+
+        currentMemoryTracker = nextMemoryTracker;
+        nextMemoryTracker = tracker;
+    }
+
+    /** Current delta. */
+    public Iterable<Row> current() {
+        return () -> current.iterator();
+    }
+
+    /** Clears all query-local rows. */
+    public void clear() {
+        current = Collections.emptyList();
+        next = null;
+
+        currentMemoryTracker.reset();
+        nextMemoryTracker.reset();
+    }
+}

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/rel/RecursiveTableSpoolNode.java

@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ignite.internal.processors.query.calcite.exec.rel;
+
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.ignite.internal.processors.query.calcite.exec.ExecutionContext;
+import org.apache.ignite.internal.util.typedef.F;
+
+/** Pass-through spool that atomically replaces the recursive CTE delta on end-of-input. */
+public class RecursiveTableSpoolNode<Row> extends AbstractNode<Row> implements SingleNode<Row>, Downstream<Row> {
+    /** Query-local recursive state. */
+    private final RecursiveCteState<Row> state;
+
+    /** Whether the current input cycle has started writing. */
+    private boolean writing;
+
+    /** */
+    public RecursiveTableSpoolNode(
+        ExecutionContext<Row> ctx,
+        RelDataType rowType,
+        RecursiveCteState<Row> state
+    ) {
+        super(ctx, rowType);
+
+        this.state = state;
+    }
+
+    /** {@inheritDoc} */
+    @Override public void request(int rowsCnt) throws Exception {
+        assert !F.isEmpty(sources()) && sources().size() == 1;
+        assert rowsCnt > 0;
+
+        checkState();
+
+        if (!writing) {
+            state.beginWrite();
+            writing = true;
+        }
+
+        source().request(rowsCnt);
+    }
+
+    /** {@inheritDoc} */
+    @Override public void push(Row row) throws Exception {
+        assert downstream() != null;
+        assert writing;
+
+        checkState();
+
+        state.add(row);
+        downstream().push(row);
+    }
+
+    /** {@inheritDoc} */
+    @Override public void end() throws Exception {
+        assert downstream() != null;
+        assert writing;
+
+        checkState();
+
+        state.commit();
+        writing = false;
+
+        downstream().end();
+    }
+
+    /** {@inheritDoc} */
+    @Override protected void rewindInternal() {
+        writing = false;
+    }
+
+    /** {@inheritDoc} */
+    @Override protected Downstream<Row> requestDownstream(int idx) {
+        if (idx != 0)
+            throw new IndexOutOfBoundsException();
+
+        return this;
+    }
+}