[CALCITE-6199] Trim unused fields for SNAPSHOT and SAMPLE if table ha…

…s VIRTUAL column
sap-contributions · Jan 27, 2024 · 75511b8 · 75511b8
1 parent dc4ce8b
commit 75511b8
Show file tree

Hide file tree

Showing 3 changed files with 161 additions and 0 deletions.
diff --git a/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java b/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
@@ -34,7 +34,9 @@
 import org.apache.calcite.rel.core.JoinRelType;
 import org.apache.calcite.rel.core.Project;
 import org.apache.calcite.rel.core.RelFactories;
+import org.apache.calcite.rel.core.Sample;
 import org.apache.calcite.rel.core.SetOp;
+import org.apache.calcite.rel.core.Snapshot;
 import org.apache.calcite.rel.core.Sort;
 import org.apache.calcite.rel.core.SortExchange;
 import org.apache.calcite.rel.core.TableScan;
@@ -1236,6 +1238,69 @@ public TrimResult trimFields(
     return result(newValues, mapping, values);
   }
 
+  /**
+   * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
+   * {@link org.apache.calcite.rel.core.Sample}.
+   */
+  public TrimResult trimFields(
+      final Sample sample,
+      ImmutableBitSet fieldsUsed,
+      Set<RelDataTypeField> extraFields) {
+    final RelDataType rowType = sample.getRowType();
+    final int fieldCount = rowType.getFieldCount();
+    final RelNode input = sample.getInput();
+
+    // Create input with trimmed columns.
+    final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
+    TrimResult trimResult =
+        trimChild(sample, input, fieldsUsed, inputExtraFields);
+    final RelNode newInput = trimResult.left;
+    final Mapping inputMapping = trimResult.right;
+
+    // If the input is unchanged, and we need to project all columns,
+    // there's nothing we can do.
+    if (newInput == input
+        && fieldsUsed.cardinality() == fieldCount) {
+      return result(sample, Mappings.createIdentity(fieldCount));
+    }
+
+    final RelNode newSample =
+        sample.copy(sample.getTraitSet(), ImmutableList.of(newInput));
+    return result(newSample, inputMapping, sample);
+  }
+
+  /**
+   * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
+   * {@link org.apache.calcite.rel.core.Snapshot}.
+   */
+  public TrimResult trimFields(
+      final Snapshot snapshot,
+      ImmutableBitSet fieldsUsed,
+      Set<RelDataTypeField> extraFields) {
+    final RelDataType rowType = snapshot.getRowType();
+    final int fieldCount = rowType.getFieldCount();
+    final RelNode input = snapshot.getInput();
+
+    // Create input with trimmed columns.
+    final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
+    TrimResult trimResult =
+        trimChild(snapshot, input, fieldsUsed, inputExtraFields);
+    final RelNode newInput = trimResult.left;
+    final Mapping inputMapping = trimResult.right;
+
+    // If the input is unchanged, and we need to project all columns,
+    // there's nothing we can do.
+    if (newInput == input
+        && fieldsUsed.cardinality() == fieldCount) {
+      return result(snapshot, Mappings.createIdentity(fieldCount));
+    }
+
+    final Snapshot newSnapshot =
+        snapshot.copy(snapshot.getTraitSet(), newInput,
+        snapshot.getPeriod());
+    return result(newSnapshot, inputMapping, snapshot);
+  }
+
   protected Mapping createMapping(ImmutableBitSet fieldsUsed, int fieldCount) {
     final Mapping mapping =
         Mappings.create(

diff --git a/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java b/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java
@@ -4955,6 +4955,54 @@ void checkUserDefinedOrderByOver(NullCollation nullCollation) {
     sql(sql).withTrim(true).ok();
   }
 
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6199">[CALCITE-6199]
+   * Trim unused fields for SNAPSHOT and SAMPLE if table has VIRTUAL column</a>.
+   */
+  @Test void testTrimSnapshotOnTemporalTable1() {
+    // Test temporal table with virtual columns.
+    final String sql = "select D, E from VIRTUALCOLUMNS.VC_T1 "
+        + "for system_time as of TIMESTAMP '2011-01-02 00:00:00'";
+    sql(sql).withExtendedTester().withTrim(true).ok();
+  }
+
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6199">[CALCITE-6199]
+   * Trim unused fields for SNAPSHOT and SAMPLE if table has VIRTUAL column</a>.
+   */
+  @Test void testTrimSnapshotOnTemporalTable2() {
+    // Test temporal table with virtual columns.
+    final String sql = "select * from VIRTUALCOLUMNS.VC_T1 "
+        + "for system_time as of TIMESTAMP '2011-01-02 00:00:00'";
+    sql(sql).withExtendedTester().withTrim(true).ok();
+  }
+
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6199">[CALCITE-6199]
+   * Trim unused fields for SNAPSHOT and SAMPLE if table has VIRTUAL column</a>.
+   */
+  @Test void testTrimSampleOnTemporalTable1() {
+    // Test temporal table with virtual columns.
+    final String sql = "select D, E from VIRTUALCOLUMNS.VC_T1 "
+        + " tablesample bernoulli(50)";
+    sql(sql).withExtendedTester().withTrim(true).ok();
+  }
+
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6199">[CALCITE-6199]
+   * Trim unused fields for SNAPSHOT and SAMPLE if table has VIRTUAL column</a>.
+   */
+  @Test void testTrimSampleOnTemporalTable2() {
+    // Test temporal table with virtual columns.
+    final String sql = "select * from VIRTUALCOLUMNS.VC_T1 "
+        + " tablesample bernoulli(50)";
+    sql(sql).withExtendedTester().withTrim(true).ok();
+  }
+
   @Test void testJoinWithOnConditionQuery() {
     String sql = ""
         + "SELECT emp.deptno, emp.sal\n"

diff --git a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
@@ -7906,6 +7906,54 @@ LogicalProject(DEPTNO=[$1])
       LogicalTableScan(table=[[CATALOG, SALES, EMP]])
     LogicalProject(NAME=[$1], DEPTNO=[$0])
       LogicalTableScan(table=[[CATALOG, SALES, DEPT]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testTrimSampleOnTemporalTable1">
+    <Resource name="sql">
+      <![CDATA[select D, E from VIRTUALCOLUMNS.VC_T1  tablesample bernoulli(50)]]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+Sample(mode=[bernoulli], rate=[0.5], repeatableSeed=[-])
+  LogicalProject(D=[$3], $f4=[+($0, 1)])
+    LogicalTableScan(table=[[CATALOG, VIRTUALCOLUMNS, VC_T1]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testTrimSampleOnTemporalTable2">
+    <Resource name="sql">
+      <![CDATA[select * from VIRTUALCOLUMNS.VC_T1  tablesample bernoulli(50)]]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+Sample(mode=[bernoulli], rate=[0.5], repeatableSeed=[-])
+  LogicalProject(A=[$0], B=[$1], C=[$2], D=[$3], $f4=[+($0, 1)])
+    LogicalTableScan(table=[[CATALOG, VIRTUALCOLUMNS, VC_T1]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testTrimSnapshotOnTemporalTable1">
+    <Resource name="sql">
+      <![CDATA[select D, E from VIRTUALCOLUMNS.VC_T1 for system_time as of TIMESTAMP '2011-01-02 00:00:00']]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+LogicalSnapshot(period=[2011-01-02 00:00:00])
+  LogicalProject(D=[$3], $f4=[+($0, 1)])
+    LogicalTableScan(table=[[CATALOG, VIRTUALCOLUMNS, VC_T1]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testTrimSnapshotOnTemporalTable2">
+    <Resource name="sql">
+      <![CDATA[select * from VIRTUALCOLUMNS.VC_T1 for system_time as of TIMESTAMP '2011-01-02 00:00:00']]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+LogicalSnapshot(period=[2011-01-02 00:00:00])
+  LogicalProject(A=[$0], B=[$1], C=[$2], D=[$3], $f4=[+($0, 1)])
+    LogicalTableScan(table=[[CATALOG, VIRTUALCOLUMNS, VC_T1]])
 ]]>
     </Resource>
   </TestCase>