From dcbc493bf699d961427952c5efc047b76d859096 Mon Sep 17 00:00:00 2001
From: IceMimosa <chk19940609@gmail.com>
Date: Wed, 23 Feb 2022 18:06:10 +0800
Subject: [PATCH] [CALCITE-5019] Avoid multiple scans when table is
 ProjectableFilterableTable and projections and filters act on different
 columns

---
 .../calcite/interpreter/TableScanNode.java    |  4 +--
 .../calcite/test/ScannableTableTest.java      | 29 +++++++++++++++++--
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/interpreter/TableScanNode.java b/core/src/main/java/org/apache/calcite/interpreter/TableScanNode.java
index 2f552f1374a..ece2008dd07 100644
--- a/core/src/main/java/org/apache/calcite/interpreter/TableScanNode.java
+++ b/core/src/main/java/org/apache/calcite/interpreter/TableScanNode.java
@@ -200,8 +200,6 @@ private static TableScanNode createProjectableFilterable(Compiler compiler,
       } else {
         projectInts = projects.toIntArray();
       }
-      final Enumerable<@Nullable Object[]> enumerable1 =
-          pfTable.scan(root, mutableFilters, projectInts);
       for (RexNode filter : mutableFilters) {
         if (!filters.contains(filter)) {
           throw RESOURCE.filterableTableInventedFilter(filter.toString())
@@ -227,6 +225,8 @@ private static TableScanNode createProjectableFilterable(Compiler compiler,
           continue;
         }
       }
+      final Enumerable<@Nullable Object[]> enumerable1 =
+          pfTable.scan(root, mutableFilters, projectInts);
       final Enumerable<Row> rowEnumerable = Enumerables.toRow(enumerable1);
       final ImmutableIntList rejectedProjects;
       if (originalProjects == null || originalProjects.equals(projects)) {
diff --git a/core/src/test/java/org/apache/calcite/test/ScannableTableTest.java b/core/src/test/java/org/apache/calcite/test/ScannableTableTest.java
index 2a64d97b2f4..0339cd45fc6 100644
--- a/core/src/test/java/org/apache/calcite/test/ScannableTableTest.java
+++ b/core/src/test/java/org/apache/calcite/test/ScannableTableTest.java
@@ -155,7 +155,7 @@ public class ScannableTableTest {
             "j=Paul");
     // Only 2 rows came out of the table. If the value is 4, it means that the
     // planner did not pass the filter down.
-    assertThat(buf.toString(), is("returnCount=2, filter=<0, 4>, projects=[1]"));
+    assertThat(buf.toString(), is("returnCount=2, filter=<0, 4>, projects=[1, 0]"));
   }
 
   @Test void testProjectableFilterableNonCooperative() throws Exception {
@@ -188,7 +188,7 @@ public class ScannableTableTest {
         .returnsUnordered("k=1940; j=John",
             "k=1942; j=Paul");
     assertThat(buf.toString(),
-        is("returnCount=2, filter=<0, 4>, projects=[2, 1]"));
+        is("returnCount=2, filter=<0, 4>, projects=[2, 1, 0]"));
   }
 
   /** A filter on a {@link org.apache.calcite.schema.ProjectableFilterableTable}
@@ -396,6 +396,25 @@ private static Pair<Integer, Object> getFilter(boolean cooperative, List<RexNode
             .explainContains(explain);
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-5019">[CALCITE-5019]
+   * Avoid multiple scans when table is ProjectableFilterableTable</a>.*/
+  @Test void testProjectableFilterableWithScanCounter() throws Exception {
+    final StringBuilder buf = new StringBuilder();
+    final BeatlesProjectableFilterableTable table =
+        new BeatlesProjectableFilterableTable(buf, false);
+    final String explain = "PLAN="
+        + "EnumerableInterpreter\n"
+        + "  BindableTableScan(table=[[s, beatles]], filters=[[=($0, 4)]], projects=[[1]]";
+    CalciteAssert.that()
+        .with(newSchema("s", Pair.of("beatles", table)))
+        .query("select \"j\" from \"s\".\"beatles\" where \"i\" = 4")
+        .explainContains(explain)
+        .returnsUnordered("j=John", "j=Paul");
+    assertThat(table.getScanCount(), is(1));
+    assertThat(buf.toString(), is("returnCount=4, projects=[1, 0]"));
+  }
+
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-1031">[CALCITE-1031]
    * In prepared statement, CsvScannableTable.scan is called twice</a>. */
@@ -558,6 +577,7 @@ public Enumerator<Object[]> enumerator() {
    * interface. */
   public static class BeatlesProjectableFilterableTable
       extends AbstractTable implements ProjectableFilterableTable {
+    private final AtomicInteger scanCounter = new AtomicInteger();
     private final StringBuilder buf;
     private final boolean cooperative;
 
@@ -577,6 +597,7 @@ public RelDataType getRowType(RelDataTypeFactory typeFactory) {
 
     public Enumerable<@Nullable Object[]> scan(DataContext root, List<RexNode> filters,
         final int @Nullable [] projects) {
+      scanCounter.incrementAndGet();
       final Pair<Integer, Object> filter = getFilter(cooperative, filters);
       return new AbstractEnumerable<Object[]>() {
         public Enumerator<Object[]> enumerator() {
@@ -584,6 +605,10 @@ public Enumerator<Object[]> enumerator() {
         }
       };
     }
+
+    public int getScanCount() {
+      return this.scanCounter.get();
+    }
   }
 
   private static Enumerator<Object[]> tens() {