[multistage] UNION/INTERSECT/EXCEPT implementation

pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotQueryRuleSets.java

@@ -122,6 +122,7 @@ private PinotQueryRuleSets() {
 
       PinotJoinExchangeNodeInsertRule.INSTANCE,
       PinotAggregateExchangeNodeInsertRule.INSTANCE,
-      PinotWindowExchangeNodeInsertRule.INSTANCE
+      PinotWindowExchangeNodeInsertRule.INSTANCE,
+      PinotSetOpExchangeNodeInsertRule.INSTANCE
   );
 }

pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSetOpExchangeNodeInsertRule.java

@@ -0,0 +1,87 @@
+/**
+ * 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.calcite.rel.rules;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.SetOp;
+import org.apache.calcite.rel.logical.LogicalExchange;
+import org.apache.calcite.rel.logical.LogicalIntersect;
+import org.apache.calcite.rel.logical.LogicalMinus;
+import org.apache.calcite.rel.logical.LogicalUnion;
+import org.apache.calcite.tools.RelBuilderFactory;
+
+
+/**
+ * Special rule for Pinot, this rule is fixed to always insert exchange after SetOp node.
+ */
+public class PinotSetOpExchangeNodeInsertRule extends RelOptRule {
+  public static final PinotSetOpExchangeNodeInsertRule INSTANCE =
+      new PinotSetOpExchangeNodeInsertRule(PinotRuleUtils.PINOT_REL_FACTORY);
+
+  public PinotSetOpExchangeNodeInsertRule(RelBuilderFactory factory) {
+    super(operand(SetOp.class, any()), factory, null);
+  }
+
+  @Override
+  public boolean matches(RelOptRuleCall call) {
+    if (call.rels.length < 1) {
+      return false;
+    }
+    if (call.rel(0) instanceof SetOp) {
+      SetOp setOp = call.rel(0);
+      for (RelNode input : setOp.getInputs()) {
+        if (PinotRuleUtils.isExchange(input)) {
+          return false;
+        }
+      }
+      return true;
+    }
+    return false;
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    SetOp setOp = call.rel(0);
+    List<RelNode> newInputs = new ArrayList<>();
+    List<Integer> hashFields =
+        IntStream.range(0, setOp.getRowType().getFieldCount()).boxed().collect(Collectors.toCollection(ArrayList::new));
+    for (RelNode input : setOp.getInputs()) {
+      RelNode exchange = LogicalExchange.create(input, RelDistributions.hash(hashFields));
+      newInputs.add(exchange);
+    }
+    SetOp newSetOpNode;
+    if (setOp instanceof LogicalUnion) {
+      newSetOpNode = new LogicalUnion(setOp.getCluster(), setOp.getTraitSet(), newInputs, setOp.all);
+    } else if (setOp instanceof LogicalIntersect) {
+      newSetOpNode = new LogicalIntersect(setOp.getCluster(), setOp.getTraitSet(), newInputs, setOp.all);
+    } else if (setOp instanceof LogicalMinus) {
+      newSetOpNode = new LogicalMinus(setOp.getCluster(), setOp.getTraitSet(), newInputs, setOp.all);
+    } else {
+      throw new UnsupportedOperationException("Unsupported set op node: " + setOp);
+    }
+    call.transformTo(newSetOpNode);
+  }
+}

pinot-query-planner/src/test/resources/queries/SetOpPlans.json

@@ -7,10 +7,12 @@
         "output": [
           "Execution Plan",
           "\nLogicalUnion(all=[true])",
-          "\n  LogicalProject(col1=[$0], col2=[$1])",
-          "\n    LogicalTableScan(table=[[a]])",
-          "\n  LogicalProject(col1=[$0], col2=[$1])",
-          "\n    LogicalTableScan(table=[[b]])",
+          "\n  LogicalExchange(distribution=[hash[0, 1]])",
+          "\n    LogicalProject(col1=[$0], col2=[$1])",
+          "\n      LogicalTableScan(table=[[a]])",
+          "\n  LogicalExchange(distribution=[hash[0, 1]])",
+          "\n    LogicalProject(col1=[$0], col2=[$1])",
+          "\n      LogicalTableScan(table=[[b]])",
           "\n"
         ]
       },
@@ -20,13 +22,17 @@
         "output": [
           "Execution Plan",
           "\nLogicalUnion(all=[true])",
-          "\n  LogicalUnion(all=[true])",
-          "\n    LogicalProject(col1=[$0], col2=[$1])",
-          "\n      LogicalTableScan(table=[[a]])",
+          "\n  LogicalExchange(distribution=[hash[0, 1]])",
+          "\n    LogicalUnion(all=[true])",
+          "\n      LogicalExchange(distribution=[hash[0, 1]])",
+          "\n        LogicalProject(col1=[$0], col2=[$1])",
+          "\n          LogicalTableScan(table=[[a]])",
+          "\n      LogicalExchange(distribution=[hash[0, 1]])",
+          "\n        LogicalProject(col1=[$0], col2=[$1])",
+          "\n          LogicalTableScan(table=[[b]])",
+          "\n  LogicalExchange(distribution=[hash[0, 1]])",
           "\n    LogicalProject(col1=[$0], col2=[$1])",
-          "\n      LogicalTableScan(table=[[b]])",
-          "\n  LogicalProject(col1=[$0], col2=[$1])",
-          "\n    LogicalTableScan(table=[[c]])",
+          "\n      LogicalTableScan(table=[[c]])",
           "\n"
         ]
       },
@@ -39,13 +45,17 @@
           "\n  LogicalExchange(distribution=[hash[0, 1]])",
           "\n    LogicalAggregate(group=[{0, 1}])",
           "\n      LogicalUnion(all=[true])",
-          "\n        LogicalUnion(all=[true])",
-          "\n          LogicalProject(col1=[$0], col2=[$1])",
-          "\n            LogicalTableScan(table=[[a]])",
+          "\n        LogicalExchange(distribution=[hash[0, 1]])",
+          "\n          LogicalUnion(all=[true])",
+          "\n            LogicalExchange(distribution=[hash[0, 1]])",
+          "\n              LogicalProject(col1=[$0], col2=[$1])",
+          "\n                LogicalTableScan(table=[[a]])",
+          "\n            LogicalExchange(distribution=[hash[0, 1]])",
+          "\n              LogicalProject(col1=[$0], col2=[$1])",
+          "\n                LogicalTableScan(table=[[b]])",
+          "\n        LogicalExchange(distribution=[hash[0, 1]])",
           "\n          LogicalProject(col1=[$0], col2=[$1])",
-          "\n            LogicalTableScan(table=[[b]])",
-          "\n        LogicalProject(col1=[$0], col2=[$1])",
-          "\n          LogicalTableScan(table=[[c]])",
+          "\n            LogicalTableScan(table=[[c]])",
           "\n"
         ]
       },
@@ -55,13 +65,17 @@
         "output": [
           "Execution Plan",
           "\nLogicalIntersect(all=[false])",
-          "\n  LogicalIntersect(all=[false])",
-          "\n    LogicalProject(col1=[$0], col2=[$1])",
-          "\n      LogicalTableScan(table=[[a]])",
+          "\n  LogicalExchange(distribution=[hash[0, 1]])",
+          "\n    LogicalIntersect(all=[false])",
+          "\n      LogicalExchange(distribution=[hash[0, 1]])",
+          "\n        LogicalProject(col1=[$0], col2=[$1])",
+          "\n          LogicalTableScan(table=[[a]])",
+          "\n      LogicalExchange(distribution=[hash[0, 1]])",
+          "\n        LogicalProject(col1=[$0], col2=[$1])",
+          "\n          LogicalTableScan(table=[[b]])",
+          "\n  LogicalExchange(distribution=[hash[0, 1]])",
           "\n    LogicalProject(col1=[$0], col2=[$1])",
-          "\n      LogicalTableScan(table=[[b]])",
-          "\n  LogicalProject(col1=[$0], col2=[$1])",
-          "\n    LogicalTableScan(table=[[c]])",
+          "\n      LogicalTableScan(table=[[c]])",
           "\n"
         ]
       },
@@ -71,13 +85,17 @@
         "output": [
           "Execution Plan",
           "\nLogicalMinus(all=[false])",
-          "\n  LogicalMinus(all=[false])",
-          "\n    LogicalProject(col1=[$0], col2=[$1])",
-          "\n      LogicalTableScan(table=[[a]])",
+          "\n  LogicalExchange(distribution=[hash[0, 1]])",
+          "\n    LogicalMinus(all=[false])",
+          "\n      LogicalExchange(distribution=[hash[0, 1]])",
+          "\n        LogicalProject(col1=[$0], col2=[$1])",
+          "\n          LogicalTableScan(table=[[a]])",
+          "\n      LogicalExchange(distribution=[hash[0, 1]])",
+          "\n        LogicalProject(col1=[$0], col2=[$1])",
+          "\n          LogicalTableScan(table=[[b]])",
+          "\n  LogicalExchange(distribution=[hash[0, 1]])",
           "\n    LogicalProject(col1=[$0], col2=[$1])",
-          "\n      LogicalTableScan(table=[[b]])",
-          "\n  LogicalProject(col1=[$0], col2=[$1])",
-          "\n    LogicalTableScan(table=[[c]])",
+          "\n      LogicalTableScan(table=[[c]])",
           "\n"
         ]
       }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/IntersectOperator.java

@@ -0,0 +1,49 @@
+/**
+ * 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.pinot.query.runtime.operator;
+
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.core.data.table.Record;
+import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+
+
+/**
+ * Intersect operator.
+ */
+public class IntersectOperator extends SetOperator {
+  private static final String EXPLAIN_NAME = "INTERSECT";
+
+  public IntersectOperator(OpChainExecutionContext opChainExecutionContext, List<MultiStageOperator> upstreamOperators,
+      DataSchema dataSchema) {
+    super(opChainExecutionContext, upstreamOperators, dataSchema);
+  }
+
+  @Nullable
+  @Override
+  public String toExplainString() {
+    return EXPLAIN_NAME;
+  }
+
+  @Override
+  protected boolean handleRowMatched(Object[] row) {
+    return _rightRowSet.remove(new Record(row));
+  }
+}

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MinusOperator.java

@@ -0,0 +1,49 @@
+/**
+ * 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.pinot.query.runtime.operator;
+
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.core.data.table.Record;
+import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+
+
+/**
+ * Minus/Except operator.
+ */
+public class MinusOperator extends SetOperator {
+  private static final String EXPLAIN_NAME = "MINUS";
+
+  public MinusOperator(OpChainExecutionContext opChainExecutionContext, List<MultiStageOperator> upstreamOperators,
+      DataSchema dataSchema) {
+    super(opChainExecutionContext, upstreamOperators, dataSchema);
+  }
+
+  @Nullable
+  @Override
+  public String toExplainString() {
+    return EXPLAIN_NAME;
+  }
+
+  @Override
+  protected boolean handleRowMatched(Object[] row) {
+    return _rightRowSet.add(new Record(row));
+  }
+}