Use consistent prepareStatement to resolve the issue

Signed-off-by: Songkan Tang <songkant@amazon.com>

Author: songkant-aws

core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java

@@ -63,6 +63,7 @@
 import org.apache.calcite.plan.RelOptSchema;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelOptTable.ViewExpander;
+import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.plan.hep.HepPlanner;
 import org.apache.calcite.plan.hep.HepProgram;
 import org.apache.calcite.plan.hep.HepProgramBuilder;
@@ -367,6 +368,36 @@ protected SqlToRelConverter getSqlToRelConverter(
       return new OpenSearchSqlToRelConverter(
           this, validator, catalogReader, this.cluster, convertletTable, config);
     }
+
+    @Override
+    protected RelRoot trimUnusedFields(RelRoot root) {
+      final SqlToRelConverter.Config config =
+          SqlToRelConverter.config()
+              .withTrimUnusedFields(shouldTrim(root.rel))
+              .withExpand(THREAD_EXPAND.get())
+              .withInSubQueryThreshold(requireNonNull(THREAD_INSUBQUERY_THRESHOLD.get()));
+      // PPL analyzes into a pre-built RelNode before prepareStatement(rel). Reuse the incoming
+      // RelNode's cluster here so prepare-time trimming does not create replacement nodes under a
+      // different planner than the rest of the tree.
+      final SqlToRelConverter converter =
+          new OpenSearchSqlToRelConverter(
+              this,
+              getSqlValidator(),
+              catalogReader,
+              root.rel.getCluster(),
+              convertletTable,
+              config);
+      final boolean ordered = !root.collation.getFieldCollations().isEmpty();
+      final boolean dml = SqlKind.DML.contains(root.kind);
+      return root.withRel(converter.trimUnusedFields(dml || ordered, root.rel));
+    }
+
+    private static boolean shouldTrim(RelNode rootRel) {
+      // For now, don't trim if there are more than 3 joins. The projects
+      // near the leaves created by trim migrate past joins and seem to
+      // prevent join-reordering.
+      return THREAD_TRIM.get() || RelOptUtil.countJoins(rootRel) < 2;
+    }
   }
 
   public static class OpenSearchSqlToRelConverter extends SqlToRelConverter {

core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchRelFieldTrimmer.java

@@ -5,31 +5,21 @@
 
 package org.opensearch.sql.calcite.utils;
 
-import java.util.ArrayList;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
-import org.apache.calcite.linq4j.Ord;
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.rel.RelNode;
-import org.apache.calcite.rel.core.Aggregate;
-import org.apache.calcite.rel.core.CorrelationId;
-import org.apache.calcite.rel.core.Project;
-import org.apache.calcite.rel.core.Values;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeField;
-import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexPermuteInputsShuttle;
-import org.apache.calcite.rex.RexSubQuery;
-import org.apache.calcite.rex.RexUtil;
 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.Mapping;
-import org.apache.calcite.util.mapping.MappingType;
 import org.apache.calcite.util.mapping.Mappings;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.opensearch.sql.calcite.plan.rel.Dedup;
@@ -40,94 +30,9 @@
  * <p>This class extends Calcite's RelFieldTrimmer to support trimming customized operators.
  */
 public class OpenSearchRelFieldTrimmer extends RelFieldTrimmer {
-  private final RelBuilder openSearchRelBuilder;
 
   public OpenSearchRelFieldTrimmer(@Nullable SqlValidator validator, RelBuilder relBuilder) {
     super(validator, relBuilder);
-    this.openSearchRelBuilder = relBuilder;
-  }
-
-  @Override
-  public TrimResult trimFields(
-      Project project, ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
-    final RelDataType rowType = project.getRowType();
-    final int fieldCount = rowType.getFieldCount();
-    final RelNode input = project.getInput();
-
-    final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>(extraFields);
-    RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
-    for (Ord<RexNode> ord : Ord.zip(project.getProjects())) {
-      if (fieldsUsed.get(ord.i)) {
-        ord.e.accept(inputFinder);
-      }
-    }
-
-    List<RexSubQuery> subQueries = RexUtil.SubQueryCollector.collect(project);
-    Set<CorrelationId> correlationIds = RelOptUtil.getVariablesUsed(subQueries);
-    ImmutableBitSet requiredColumns = ImmutableBitSet.of();
-    if (!correlationIds.isEmpty()) {
-      assert correlationIds.size() == 1;
-      requiredColumns = RelOptUtil.correlationColumns(correlationIds.iterator().next(), project);
-    }
-
-    ImmutableBitSet finderFields = inputFinder.build();
-    ImmutableBitSet inputFieldsUsed =
-        ImmutableBitSet.builder().addAll(requiredColumns).addAll(finderFields).build();
-
-    TrimResult trimResult = trimChild(project, input, inputFieldsUsed, inputExtraFields);
-    RelNode newInput = trimResult.left;
-    final Mapping inputMapping = trimResult.right;
-
-    if (newInput == input && fieldsUsed.cardinality() == fieldCount) {
-      return result(project, Mappings.createIdentity(fieldCount));
-    }
-
-    if (fieldsUsed.cardinality() == 0) {
-      return dummyProject(fieldCount, newInput, project);
-    }
-
-    final List<RexNode> newProjects = new ArrayList<>();
-    final RexVisitor<RexNode> shuttle;
-    if (!correlationIds.isEmpty()) {
-      assert correlationIds.size() == 1;
-      shuttle =
-          new RexPermuteInputsShuttle(inputMapping, newInput) {
-            @Override
-            public RexNode visitSubQuery(RexSubQuery subQuery) {
-              subQuery = (RexSubQuery) super.visitSubQuery(subQuery);
-              return RelOptUtil.remapCorrelatesInSuqQuery(
-                  openSearchRelBuilder.getRexBuilder(),
-                  subQuery,
-                  correlationIds.iterator().next(),
-                  newInput.getRowType(),
-                  inputMapping);
-            }
-          };
-    } else {
-      shuttle = new RexPermuteInputsShuttle(inputMapping, newInput);
-    }
-
-    final Mapping mapping =
-        Mappings.create(MappingType.INVERSE_SURJECTION, fieldCount, fieldsUsed.cardinality());
-    for (Ord<RexNode> ord : Ord.zip(project.getProjects())) {
-      if (fieldsUsed.get(ord.i)) {
-        mapping.set(ord.i, newProjects.size());
-        RexNode newProjectExpr = ord.e.accept(shuttle);
-        newProjects.add(newProjectExpr);
-      }
-    }
-
-    final RelDataType newRowType =
-        RelOptUtil.permute(project.getCluster().getTypeFactory(), rowType, mapping);
-
-    if (shouldAvoidSimplifyValues(newProjects, newInput)) {
-      return result(
-          project.copy(project.getTraitSet(), newInput, newProjects, newRowType), mapping, project);
-    }
-
-    openSearchRelBuilder.push(newInput);
-    openSearchRelBuilder.project(newProjects, newRowType.getFieldNames(), false, correlationIds);
-    return result(openSearchRelBuilder.build(), mapping, project);
   }
 
   public TrimResult trimFields(
@@ -162,19 +67,4 @@ public TrimResult trimFields(
     // needs them for its condition.
     return result(dedup.copy(newInput, newDedupFields), inputMapping);
   }
-
-  private boolean shouldAvoidSimplifyValues(List<RexNode> projects, RelNode input) {
-    return projects.stream().allMatch(RexLiteral.class::isInstance) && isFixedRowCount(input);
-  }
-
-  private boolean isFixedRowCount(RelNode input) {
-    if (input instanceof Values) {
-      return true;
-    }
-    if (input instanceof Aggregate aggregate) {
-      return aggregate.getGroupSet().isEmpty()
-          && aggregate.getGroupType() == Aggregate.Group.SIMPLE;
-    }
-    return false;
-  }
 }