[CALCITE-5984] Allow disabling field trimming in Prepare via `SqlTo…

…RelConverter.Config#isTrimUnusedFields` Close #3436
apache · Nov 1, 2023 · 8af1f11 · 8af1f11
1 parent 620058f
commit 8af1f11
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 7 deletions.
diff --git a/core/src/main/java/org/apache/calcite/prepare/Prepare.java b/core/src/main/java/org/apache/calcite/prepare/Prepare.java
@@ -293,10 +293,12 @@ public PreparedResult prepareSql(
       root = root.withRel(decorrelate(sqlToRelConverter, sqlQuery, root.rel));
     }
 
-    // Trim unused fields.
-    root = trimUnusedFields(root);
+    if (configHolder.get().isTrimUnusedFields()) {
+      // Trim unused fields.
+      root = trimUnusedFields(root);
 
-    Hook.TRIMMED.run(root.rel);
+      Hook.TRIMMED.run(root.rel);
+    }
 
     // Display physical plan after decorrelation.
     if (sqlExplain != null) {

diff --git a/core/src/main/java/org/apache/calcite/tools/Programs.java b/core/src/main/java/org/apache/calcite/tools/Programs.java
@@ -250,11 +250,17 @@ public static Program getProgram() {
 
   /** Returns the standard program used by Prepare. */
   public static Program standard() {
-    return standard(DefaultRelMetadataProvider.INSTANCE);
+    return standard(DefaultRelMetadataProvider.INSTANCE, true);
   }
 
   /** Returns the standard program with user metadata provider. */
   public static Program standard(RelMetadataProvider metadataProvider) {
+    return standard(metadataProvider, true);
+  }
+
+  /** Returns the standard program with user metadata provider and enableFieldTrimming config. */
+  public static Program standard(RelMetadataProvider metadataProvider,
+      boolean enableFieldTrimming) {
     final Program program1 =
         (planner, rel, requiredOutputTraits, materializations, lattices) -> {
           for (RelOptMaterialization materialization : materializations) {
@@ -278,14 +284,19 @@ public static Program standard(RelMetadataProvider metadataProvider) {
           return rootRel3;
         };
 
-    return sequence(subQuery(metadataProvider),
+    List<Program> programs =
+        Lists.newArrayList(subQuery(metadataProvider),
         new DecorrelateProgram(),
         new TrimFieldsProgram(),
         program1,
 
         // Second planner pass to do physical "tweaks". This the first time
         // that EnumerableCalcRel is introduced.
         calc(metadataProvider));
+
+    programs.removeIf(program -> !enableFieldTrimming && program instanceof TrimFieldsProgram);
+
+    return new SequenceProgram(ImmutableList.copyOf(programs));
   }
 
   /** Program backed by a {@link RuleSet}. */

diff --git a/core/src/test/java/org/apache/calcite/test/JdbcTest.java b/core/src/test/java/org/apache/calcite/test/JdbcTest.java
@@ -51,6 +51,7 @@
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.prepare.CalcitePrepareImpl;
 import org.apache.calcite.prepare.Prepare;
+import org.apache.calcite.rel.metadata.DefaultRelMetadataProvider;
 import org.apache.calcite.rel.rules.CoreRules;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
@@ -85,6 +86,8 @@
 import org.apache.calcite.test.schemata.hr.Department;
 import org.apache.calcite.test.schemata.hr.Employee;
 import org.apache.calcite.test.schemata.hr.HrSchema;
+import org.apache.calcite.tools.Program;
+import org.apache.calcite.tools.Programs;
 import org.apache.calcite.util.Bug;
 import org.apache.calcite.util.Holder;
 import org.apache.calcite.util.JsonBuilder;
@@ -106,6 +109,7 @@
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 
 import java.io.File;
@@ -168,6 +172,7 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
+import static org.junit.jupiter.params.provider.Arguments.arguments;
 
 /**
  * Tests for using Calcite via JDBC.
@@ -248,6 +253,45 @@ static Stream<String> explainFormats() {
     return Stream.of("text", "dot");
   }
 
+  static Stream<Arguments> disableTrimmingConfigsTestArguments() {
+    /** enableTrimmingByConfig, enableTrimmingByProgram, expectedLogicalPlan. */
+    return Stream.of(
+        arguments(true, true,
+            ""
+                + "LogicalProject(name=[$1])\n"
+                + "  LogicalJoin(condition=[=($0, $2)], joinType=[inner])\n"
+                + "    LogicalProject(deptno=[$0], name=[$1])\n"
+                + "      LogicalTableScan(table=[[hr, depts]])\n"
+                + "    LogicalProject(deptno=[$1])\n"
+                + "      LogicalTableScan(table=[[hr, emps]])\n"
+                + ""),
+        arguments(true, false,
+            ""
+                + "LogicalProject(name=[$1])\n"
+                + "  LogicalJoin(condition=[=($0, $2)], joinType=[inner])\n"
+                + "    LogicalProject(deptno=[$0], name=[$1])\n"
+                + "      LogicalTableScan(table=[[hr, depts]])\n"
+                + "    LogicalProject(deptno=[$1])\n"
+                + "      LogicalTableScan(table=[[hr, emps]])\n"
+                + ""),
+        arguments(false, true,
+            ""
+                + "LogicalProject(name=[$1])\n"
+                + "  LogicalJoin(condition=[=($0, $2)], joinType=[inner])\n"
+                + "    LogicalProject(deptno=[$0], name=[$1])\n"
+                + "      LogicalTableScan(table=[[hr, depts]])\n"
+                + "    LogicalProject(deptno=[$1])\n"
+                + "      LogicalTableScan(table=[[hr, emps]])\n"
+                + ""),
+        arguments(false, false,
+            ""
+                + "LogicalProject(name=[$1])\n"
+                + "  LogicalJoin(condition=[=($0, $5)], joinType=[inner])\n"
+                + "    LogicalTableScan(table=[[hr, depts]])\n"
+                + "    LogicalTableScan(table=[[hr, emps]])\n"
+                + ""));
+  }
+
   /** Runs a task (such as a test) with and without expansion. */
   static void forEachExpand(Runnable r) {
     try (TryThreadLocal.Memo ignored = Prepare.THREAD_EXPAND.push(false)) {
@@ -3119,6 +3163,28 @@ void testInnerJoinValues(String format) {
             "deptno=10; name=Sales; employees=[{100, 10, Bill, 10000.0, 1000}, {150, 10, Sebastian, 7000.0, null}]; location={-122, 38}");
   }
 
+  /** Test cases for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-5984">[CALCITE-5984]</a>
+   * Disabling trimming of unused fields via config and program. */
+  @ParameterizedTest
+  @MethodSource("disableTrimmingConfigsTestArguments")
+  void testJoinWithTrimmingConfigs(boolean enableTrimmingByConfig,
+      boolean enableTrimmingByProgram,
+      String expectedLogicalPlan) {
+    CalciteAssert.hr().query("select \"d\".\"name\" from \"hr\".\"depts\" as \"d\" \n"
+                + "  join \"hr\".\"emps\" as \"e\" on \"d\".\"deptno\" = \"e\".\"deptno\" \n")
+        .withHook(Hook.SQL2REL_CONVERTER_CONFIG_BUILDER,
+            (Consumer<Holder<Config>>) configHolder ->
+            configHolder.set(configHolder.get().withTrimUnusedFields(enableTrimmingByConfig)))
+        .withHook(Hook.PROGRAM,
+            (Consumer<Holder<Program>>)
+            programHolder -> programHolder
+                    .set(
+                        Programs.standard(
+                        DefaultRelMetadataProvider.INSTANCE, enableTrimmingByProgram)))
+        .convertContains(expectedLogicalPlan);
+  }
+
   /** A difficult query: an IN list so large that the planner promotes it
    * to a semi-join against a VALUES relation. */
   @Disabled

diff --git a/testkit/src/main/java/org/apache/calcite/test/CalciteAssert.java b/testkit/src/main/java/org/apache/calcite/test/CalciteAssert.java
@@ -681,6 +681,17 @@ static void assertPrepare(
       boolean materializationsEnabled,
       final Consumer<RelNode> convertChecker,
       final Consumer<RelNode> substitutionChecker) {
+    assertPrepare(connection, sql, materializationsEnabled, ImmutableList.of(),
+        convertChecker, substitutionChecker);
+  }
+
+  static void assertPrepare(
+      Connection connection,
+      String sql,
+      boolean materializationsEnabled,
+      List<Pair<Hook, Consumer>> hooks,
+      final Consumer<RelNode> convertChecker,
+      final Consumer<RelNode> substitutionChecker) {
     try (Closer closer = new Closer()) {
       if (convertChecker != null) {
         closer.add(
@@ -690,6 +701,9 @@ static void assertPrepare(
         closer.add(
             Hook.SUB.addThread(substitutionChecker));
       }
+      for (Pair<Hook, Consumer> hook : hooks) {
+        closer.add(hook.left.addThread(hook.right));
+      }
       ((CalciteConnection) connection).getProperties().setProperty(
           CalciteConnectionProperty.MATERIALIZATIONS_ENABLED.camelName(),
           Boolean.toString(materializationsEnabled));
@@ -1559,14 +1573,14 @@ public AssertQuery consumesPreparedStatement(
 
     public AssertQuery convertMatches(final Consumer<RelNode> checker) {
       return withConnection(connection ->
-        assertPrepare(connection, sql, this.materializationsEnabled,
+        assertPrepare(connection, sql, this.materializationsEnabled, hooks,
             checker, null));
     }
 
     public AssertQuery substitutionMatches(
         final Consumer<RelNode> checker) {
       return withConnection(connection ->
-        assertPrepare(connection, sql, materializationsEnabled, null, checker));
+        assertPrepare(connection, sql, materializationsEnabled, hooks, null, checker));
     }
 
     public AssertQuery explainContains(String expected) {