Enhance unnest support in decoupled mode

sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java

@@ -35,17 +35,19 @@
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.RelFactories;
 import org.apache.calcite.rel.metadata.DefaultRelMetadataProvider;
+import org.apache.calcite.rel.rules.AggregateProjectMergeRule;
 import org.apache.calcite.rel.rules.CoreRules;
 import org.apache.calcite.rel.rules.DateRangeRules;
+import org.apache.calcite.rel.rules.FilterCorrelateRule;
 import org.apache.calcite.rel.rules.FilterJoinRule.FilterIntoJoinRule.FilterIntoJoinRuleConfig;
+import org.apache.calcite.rel.rules.FilterProjectTransposeRule;
 import org.apache.calcite.rel.rules.JoinExtractFilterRule;
 import org.apache.calcite.rel.rules.JoinPushThroughJoinRule;
 import org.apache.calcite.rel.rules.ProjectMergeRule;
 import org.apache.calcite.rel.rules.PruneEmptyRules;
 import org.apache.calcite.sql.SqlExplainFormat;
 import org.apache.calcite.sql.SqlExplainLevel;
 import org.apache.calcite.sql2rel.RelDecorrelator;
-import org.apache.calcite.sql2rel.RelFieldTrimmer;
 import org.apache.calcite.tools.Program;
 import org.apache.calcite.tools.Programs;
 import org.apache.calcite.tools.RelBuilder;
@@ -280,11 +282,23 @@ private Program buildDecoupledLogicalOptimizationProgram(PlannerContext plannerC
     builder.addMatchLimit(CalciteRulesManager.HEP_DEFAULT_MATCH_LIMIT);
     builder.addRuleCollection(baseRuleSet(plannerContext));
     builder.addRuleInstance(CoreRules.UNION_MERGE);
+    builder.addRuleInstance(FilterCorrelateRule.Config.DEFAULT.toRule());
+    builder.addRuleInstance(FilterProjectTransposeRule.Config.DEFAULT.toRule());
     builder.addRuleInstance(JoinExtractFilterRule.Config.DEFAULT.toRule());
     builder.addRuleInstance(FilterIntoJoinRuleConfig.DEFAULT.withPredicate(DruidJoinRule::isSupportedPredicate).toRule());
+    builder.addRuleInstance(FilterProjectTransposeRule.Config.DEFAULT.toRule());
     builder.addRuleInstance(new LogicalUnnestRule());
     builder.addRuleInstance(new UnnestInputCleanupRule());
-    return Programs.of(builder.build(), true, DefaultRelMetadataProvider.INSTANCE);
+
+    final HepProgramBuilder cleanupRules = HepProgram.builder();
+    cleanupRules.addRuleInstance(FilterProjectTransposeRule.Config.DEFAULT.toRule());
+    cleanupRules.addRuleInstance(CoreRules.PROJECT_MERGE);
+    cleanupRules.addRuleInstance(AggregateProjectMergeRule.Config.DEFAULT.toRule());
+    return Programs.sequence(
+        Programs.of(builder.build(), true, DefaultRelMetadataProvider.INSTANCE),
+        new DruidTrimFieldsProgram(true),
+        Programs.of(cleanupRules.build(), true, DefaultRelMetadataProvider.INSTANCE)
+    );
   }
 
   /**
@@ -539,7 +553,31 @@ public RelNode run(
     {
       final RelBuilder relBuilder = RelFactories.LOGICAL_BUILDER.create(rel.getCluster(), null);
       final RelNode decorrelatedRel = RelDecorrelator.decorrelateQuery(rel, relBuilder);
-      return new RelFieldTrimmer(null, relBuilder).trim(decorrelatedRel);
+      RelNode ret = new DruidRelFieldTrimmer(null, relBuilder, true).trim(decorrelatedRel);
+      return ret;
+    }
+
+  }
+
+  /** Program that trims fields. */
+  private static class DruidTrimFieldsProgram implements Program
+  {
+    private boolean trim;
+
+    public DruidTrimFieldsProgram(boolean trim)
+    {
+      this.trim = trim;
+    }
+
+    @Override
+    public RelNode run(RelOptPlanner planner, RelNode rel,
+        RelTraitSet requiredOutputTraits,
+        List<RelOptMaterialization> materializations,
+        List<RelOptLattice> lattices)
+    {
+      final RelBuilder relBuilder = RelFactories.LOGICAL_BUILDER.create(rel.getCluster(), null);
+      RelNode ret = new DruidRelFieldTrimmer(null, relBuilder, trim).trim(rel);
+      return ret;
     }
   }
 }

sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidRelFieldTrimmer.java

@@ -0,0 +1,313 @@
+/*
+ * 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.druid.sql.calcite.planner;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.rel.RelHomogeneousShuttle;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.CorrelationId;
+import org.apache.calcite.rel.core.TableScan;
+import org.apache.calcite.rel.logical.LogicalCorrelate;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexFieldAccess;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexPermuteInputsShuttle;
+import org.apache.calcite.rex.RexShuttle;
+import org.apache.calcite.rex.RexVisitor;
+import org.apache.calcite.sql.validate.SqlValidator;
+import org.apache.calcite.sql2rel.RelFieldTrimmer;
+import org.apache.calcite.tools.RelBuilder;
+import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.mapping.IntPair;
+import org.apache.calcite.util.mapping.Mapping;
+import org.apache.calcite.util.mapping.MappingType;
+import org.apache.calcite.util.mapping.Mappings;
+import org.apache.druid.sql.calcite.rule.logical.LogicalUnnest;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+
+public class DruidRelFieldTrimmer extends RelFieldTrimmer
+{
+  private final boolean trimTableScan;
+  private final RelBuilder relBuilder;
+
+  public DruidRelFieldTrimmer(@Nullable SqlValidator validator, RelBuilder relBuilder, boolean trimTableScan)
+  {
+    super(validator, relBuilder);
+    this.relBuilder = relBuilder;
+    this.trimTableScan = trimTableScan;
+  }
+
+  @Override
+  protected TrimResult dummyProject(int fieldCount, RelNode input)
+  {
+    return makeIdentityMapping(input);
+  }
+
+  protected TrimResult dummyProject(int fieldCount, RelNode input,
+      @Nullable RelNode originalRelNode)
+  {
+    if (fieldCount != 0 && trimTableScan) {
+      return super.dummyProject(fieldCount, input, originalRelNode);
+    }
+    // workaround to support fieldCount == 0 projections
+    final Mapping mapping = Mappings.create(MappingType.INVERSE_SURJECTION, fieldCount, 0);
+    if (input.getRowType().getFieldCount() == 0) {
+      // there is no need to do anything
+      return result(input, mapping);
+    }
+    relBuilder.push(input);
+    relBuilder.project(Collections.emptyList(), Collections.emptyList());
+    RelNode newProject = relBuilder.build();
+    if (originalRelNode != null) {
+      newProject = RelOptUtil.propagateRelHints(originalRelNode, newProject);
+    }
+    return result(newProject, mapping);
+  }
+
+  private TrimResult makeIdentityMapping(RelNode input)
+  {
+    Mapping mapping = Mappings.createIdentity(input.getRowType().getFieldCount());
+    return result(input, mapping);
+  }
+
+  public TrimResult trimFields(
+      final TableScan tableAccessRel,
+      ImmutableBitSet fieldsUsed,
+      Set<RelDataTypeField> extraFields)
+  {
+    if (trimTableScan) {
+      return super.trimFields(tableAccessRel, fieldsUsed, extraFields);
+    } else {
+      Mapping mapping = Mappings.createIdentity(tableAccessRel.getRowType().getFieldCount());
+      return result(tableAccessRel, mapping, tableAccessRel);
+    }
+  }
+
+  public TrimResult trimFields(LogicalCorrelate correlate,
+      ImmutableBitSet fieldsUsed,
+      Set<RelDataTypeField> extraFields)
+  {
+    if (!extraFields.isEmpty()) {
+      // bail out with generic trim
+      return trimFields((RelNode) correlate, fieldsUsed, extraFields);
+    }
+
+    fieldsUsed = fieldsUsed.union(correlate.getRequiredColumns());
+
+    List<RelNode> newInputs = new ArrayList<>();
+    List<Mapping> inputMappings = new ArrayList<>();
+    int changeCount = 0;
+    int offset = 0;
+    for (RelNode input : correlate.getInputs()) {
+      final RelDataType inputRowType = input.getRowType();
+      final int inputFieldCount = inputRowType.getFieldCount();
+
+      ImmutableBitSet currentInputFieldsUsed = fieldsUsed
+          .intersect(ImmutableBitSet.range(offset, offset + inputFieldCount))
+          .shift(-offset);
+
+      TrimResult trimResult;
+      try {
+        trimResult = dispatchTrimFields(input, currentInputFieldsUsed, extraFields);
+      }
+      catch (RuntimeException e) {
+        throw e;
+      }
+
+      newInputs.add(trimResult.left);
+      if (trimResult.left != input) {
+        changeCount++;
+      }
+
+      final Mapping inputMapping = trimResult.right;
+      inputMappings.add(inputMapping);
+
+      offset += inputFieldCount;
+    }
+
+    if (changeCount == 0) {
+      return result(correlate, Mappings.createIdentity(correlate.getRowType().getFieldCount()));
+    }
+
+    Mapping mapping = makeMapping(inputMappings);
+    RexBuilder rexBuilder = correlate.getCluster().getRexBuilder();
+
+    final LogicalCorrelate newCorrelate = correlate.copy(
+        correlate.getTraitSet(),
+        newInputs.get(0),
+        newInputs.get(1).accept(
+            new RexRewritingRelShuttle(
+                new RexCorrelVariableMapShuttle(correlate.getCorrelationId(), newInputs.get(0).getRowType(), mapping, rexBuilder)
+            )
+        ),
+        correlate.getCorrelationId(),
+        correlate.getRequiredColumns().permute(mapping),
+        correlate.getJoinType()
+    );
+
+    return result(newCorrelate, mapping);
+  }
+
+  public TrimResult trimFields(LogicalUnnest correlate,
+      ImmutableBitSet fieldsUsed,
+      Set<RelDataTypeField> extraFields)
+  {
+    if (!extraFields.isEmpty()) {
+      // bail out with generic trim
+      return trimFields((RelNode) correlate, fieldsUsed, extraFields);
+    }
+    RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(extraFields);
+
+    correlate.getUnnestExpr().accept(inputFinder);
+
+    ImmutableBitSet finderFields = inputFinder.build();
+
+    ImmutableBitSet inputFieldsUsed = ImmutableBitSet.builder()
+        .addAll(fieldsUsed.clear(correlate.getRowType().getFieldCount() - 1))
+        .addAll(finderFields)
+        .build();
+
+    RelNode input = correlate.getInput();
+
+    // Create input with trimmed columns.
+    TrimResult trimResult = trimChild(correlate, input, inputFieldsUsed, extraFields);
+
+    RelNode newInput = trimResult.left;
+    final Mapping inputMapping = trimResult.right;
+
+    if (newInput == input) {
+      return result(correlate, Mappings.createIdentity(correlate.getRowType().getFieldCount()));
+    }
+
+    Mapping mapping = makeMapping(ImmutableList.of(inputMapping, Mappings.createIdentity(1)));
+
+    final RexVisitor<RexNode> shuttle = new RexPermuteInputsShuttle(inputMapping, newInput);
+
+    RexNode newUnnestExpr = correlate.getUnnestExpr().accept(shuttle);
+    RexNode newFilterExpr = correlate.getFilter();
+
+    final LogicalUnnest newCorrelate = correlate.copy(
+        correlate.getTraitSet(),
+        newInput,
+        newUnnestExpr,
+        newFilterExpr
+    );
+
+    return result(newCorrelate, mapping);
+  }
+
+  /**
+   * Concatenates multiple mapping.
+   *
+   * <pre>
+   * [ 1:0, 2:1] // sourceCount:100
+   * [ 1:0, 2:1] // sourceCount:100
+   * output:
+   * [ 1:0, 2:1, 101:2, 102:3 ] ; sourceCount:200
+   * </pre>
+   */
+  private Mapping makeMapping(List<Mapping> inputMappings)
+  {
+    int fieldCount = 0;
+    int newFieldCount = 0;
+    for (Mapping mapping : inputMappings) {
+      fieldCount += mapping.getSourceCount();
+      newFieldCount += mapping.getTargetCount();
+    }
+
+    Mapping mapping = Mappings.create(
+        MappingType.INVERSE_SURJECTION,
+        fieldCount,
+        newFieldCount
+    );
+    int offset = 0;
+    int newOffset = 0;
+    for (int i = 0; i < inputMappings.size(); i++) {
+      Mapping inputMapping = inputMappings.get(i);
+      for (IntPair pair : inputMapping) {
+        mapping.set(pair.source + offset, pair.target + newOffset);
+      }
+      offset += inputMapping.getSourceCount();
+      newOffset += inputMapping.getTargetCount();
+    }
+    return mapping;
+  }
+
+  static class RexCorrelVariableMapShuttle extends RexShuttle
+  {
+    private final CorrelationId correlationId;
+    private final Mapping mapping;
+    private final RelDataType newCorrelRowType;
+    private final RexBuilder rexBuilder;
+
+    public RexCorrelVariableMapShuttle(final CorrelationId correlationId, RelDataType newCorrelRowType, Mapping mapping, RexBuilder rexBuilder)
+    {
+      this.correlationId = correlationId;
+      this.newCorrelRowType = newCorrelRowType;
+      this.mapping = mapping;
+      this.rexBuilder = rexBuilder;
+    }
+
+    @Override
+    public RexNode visitFieldAccess(final RexFieldAccess fieldAccess)
+    {
+      if (fieldAccess.getReferenceExpr() instanceof RexCorrelVariable) {
+        RexCorrelVariable referenceExpr = (RexCorrelVariable) fieldAccess.getReferenceExpr();
+        final RexCorrelVariable encounteredCorrelationId = referenceExpr;
+        if (encounteredCorrelationId.id.equals(correlationId)) {
+          int sourceIndex = fieldAccess.getField().getIndex();
+          return rexBuilder.makeFieldAccess(map(referenceExpr), mapping.getTarget(sourceIndex));
+        }
+      }
+      return super.visitFieldAccess(fieldAccess);
+    }
+
+    private RexNode map(RexCorrelVariable referenceExpr)
+    {
+      return rexBuilder.makeCorrel(newCorrelRowType, referenceExpr.id);
+    }
+  }
+
+  static class RexRewritingRelShuttle extends RelHomogeneousShuttle
+  {
+    private final RexShuttle rexVisitor;
+
+    RexRewritingRelShuttle(RexShuttle rexVisitor)
+    {
+      this.rexVisitor = rexVisitor;
+    }
+
+    @Override
+    public RelNode visit(RelNode other)
+    {
+      RelNode next = super.visit(other);
+      return next.accept(rexVisitor);
+    }
+  }
+}