[drill] 02/04: DRILL-6553: Fix TopN for unnest operator

[ar...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

arina pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/drill.git

commit 5c9303d7e3797377191ff9a57103ef033e023db1
Author: Volodymyr Vysotskyi <vv...@gmail.com>
AuthorDate: Thu Jun 28 20:20:38 2018 +0300

    DRILL-6553: Fix TopN for unnest operator
    
    closes #1353
---
 .../planner/common/DrillLateralJoinRelBase.java    | 15 ++++++----
 .../impl/lateraljoin/TestE2EUnnestAndLateral.java  |  9 ++++--
 .../impl/lateraljoin/TestLateralPlans.java         | 33 ++++++++++++++++++++++
 3 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/common/DrillLateralJoinRelBase.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/common/DrillLateralJoinRelBase.java
index 2f895e2..5a2b40e 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/common/DrillLateralJoinRelBase.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/common/DrillLateralJoinRelBase.java
@@ -50,15 +50,15 @@ public abstract class DrillLateralJoinRelBase extends Correlate implements Drill
     this.excludeCorrelateColumn = excludeCorrelateCol;
   }
 
-  @Override public RelOptCost computeSelfCost(RelOptPlanner planner,
-                                              RelMetadataQuery mq) {
+  @Override
+  public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) {
     DrillCostBase.DrillCostFactory costFactory = (DrillCostBase.DrillCostFactory) planner.getCostFactory();
 
-    double rowCount = mq.getRowCount(this.getLeft());
+    double rowCount = estimateRowCount(mq);
     long fieldWidth = PrelUtil.getPlannerSettings(planner).getOptions()
-        .getOption(ExecConstants.AVERAGE_FIELD_WIDTH_KEY).num_val;
+        .getLong(ExecConstants.AVERAGE_FIELD_WIDTH_KEY);
 
-    double rowSize = (this.getLeft().getRowType().getFieldList().size()) * fieldWidth;
+    double rowSize = left.getRowType().getFieldList().size() * fieldWidth;
 
     double cpuCost = rowCount * rowSize * DrillCostBase.BASE_CPU_COST;
     double memCost = !excludeCorrelateColumn ? CORRELATE_MEM_COPY_COST : 0.0;
@@ -117,4 +117,9 @@ public abstract class DrillLateralJoinRelBase extends Correlate implements Drill
     }
     return inputRowType;
   }
+
+  @Override
+  public double estimateRowCount(RelMetadataQuery mq) {
+    return mq.getRowCount(left);
+  }
 }
diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestE2EUnnestAndLateral.java b/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestE2EUnnestAndLateral.java
index 98f051e..c57093c 100644
--- a/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestE2EUnnestAndLateral.java
+++ b/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestE2EUnnestAndLateral.java
@@ -18,8 +18,10 @@
 package org.apache.drill.exec.physical.impl.lateraljoin;
 
 import org.apache.drill.categories.OperatorTest;
+import org.apache.drill.exec.ExecConstants;
 import org.apache.drill.exec.planner.physical.PlannerSettings;
 import org.apache.drill.test.ClusterFixture;
+import org.apache.drill.test.ClusterFixtureBuilder;
 import org.apache.drill.test.ClusterTest;
 import org.apache.drill.test.TestBuilder;
 import org.junit.BeforeClass;
@@ -43,8 +45,10 @@ public class TestE2EUnnestAndLateral extends ClusterTest {
   public static void setupTestFiles() throws Exception {
     dirTestWatcher.copyResourceToRoot(Paths.get("lateraljoin", "multipleFiles", regularTestFile_1));
     dirTestWatcher.copyResourceToRoot(Paths.get("lateraljoin", "multipleFiles", regularTestFile_2));
-    startCluster(ClusterFixture.builder(dirTestWatcher).maxParallelization(1));
-    test("alter session set `planner.enable_unnest_lateral`=%s", true);
+    ClusterFixtureBuilder builder = ClusterFixture.builder(dirTestWatcher)
+        .sessionOption(ExecConstants.ENABLE_UNNEST_LATERAL_KEY, true)
+        .maxParallelization(1);
+    startCluster(builder);
   }
 
   /***********************************************************************************************
@@ -97,7 +101,6 @@ public class TestE2EUnnestAndLateral extends ClusterTest {
   /**
    * Test which disables the TopN operator from planner settings before running query using SORT and LIMIT in
    * subquery. The same query as in above test is executed and same result is expected.
-   * @throws Exception
    */
   @Test
   public void testLateral_WithSortAndLimitInSubQuery() throws Exception {
diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestLateralPlans.java b/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestLateralPlans.java
index 8ff164f..77d245f 100644
--- a/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestLateralPlans.java
+++ b/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestLateralPlans.java
@@ -17,8 +17,11 @@
  */
 package org.apache.drill.exec.physical.impl.lateraljoin;
 
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.core.IsNot.not;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
 import org.apache.drill.PlanTestBase;
@@ -499,4 +502,34 @@ public class TestLateralPlans extends BaseTestQuery {
       assertTrue(matcher.find());
     }
   }
+
+  @Test
+  public void testUnnestTopN() throws Exception {
+    String query =
+        "select customer.c_custkey," +
+                "customer.c_name," +
+                "t.o.o_orderkey," +
+                "t.o.o_totalprice\n" +
+        "from dfs.`lateraljoin/multipleFiles` customer," +
+              "unnest(customer.c_orders) t(o)\n" +
+        "order by customer.c_custkey," +
+                  "t.o.o_orderkey," +
+                  "t.o.o_totalprice\n" +
+        "limit 50";
+
+    ClusterFixtureBuilder builder = ClusterFixture.builder(dirTestWatcher)
+        .setOptionDefault(ExecConstants.ENABLE_UNNEST_LATERAL_KEY, true);
+
+    try (ClusterFixture cluster = builder.build();
+         ClientFixture client = cluster.clientFixture()) {
+      String plan = client.queryBuilder()
+          .sql(query)
+          .explainText();
+
+      assertThat("Query plan doesn't contain TopN operator",
+          plan, containsString("TopN(limit=[50])"));
+      assertThat("Query plan shouldn't contain Sort operator",
+          plan, not(containsString("Sort")));
+    }
+  }
 }