Enhance unnest support in decoupled mode

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java

@@ -2494,12 +2494,11 @@ public void testJoinUsesDifferentAlgorithm(String contextName, Map<String, Objec
                 new QueryDataSource(
                     newScanQueryBuilder()
                         .dataSource("foo")
-                        .virtualColumns(expressionVirtualColumn("v0", "0", ColumnType.LONG))
-                        .columns("v0")
+                        .columns("__time")
                         .columnTypes(ColumnType.LONG)
                         .context(defaultScanQueryContext(
                             queryContext,
-                            RowSignature.builder().add("v0", ColumnType.LONG).build()
+                            RowSignature.builder().add("__time", ColumnType.LONG).build()
                         ))
                         .intervals(querySegmentSpec(Intervals.ETERNITY))
                         .build()
@@ -2518,13 +2517,18 @@ public void testJoinUsesDifferentAlgorithm(String contextName, Map<String, Objec
                                         "a0",
                                         "1",
                                         ColumnType.LONG
+                                    ),
+                                    expressionPostAgg(
+                                        "p0",
+                                        "(CAST(floor(100), 'DOUBLE') == \"d0\")",
+                                        ColumnType.LONG
                                     )
                                 )
                                 .build()
 
                 ),
                 "j0.",
-                "(CAST(floor(100), 'DOUBLE') == \"j0.d0\")",
+                "(1 == \"j0.p0\")",
                 JoinType.LEFT,
                 null,
                 JoinAlgorithm.SORT_MERGE

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;
@@ -281,11 +283,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(),
+        Programs.of(cleanupRules.build(), true, DefaultRelMetadataProvider.INSTANCE)
+    );
   }
 
   /**
@@ -541,7 +555,24 @@ 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).trim(decorrelatedRel);
+      return ret;
+    }
+
+  }
+
+  /** Program that trims fields. */
+  private static class DruidTrimFieldsProgram implements Program
+  {
+    @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(rel);
+      return ret;
     }
   }
 }