From 2325844a389ef3b917eecd4a6e89799b6f367555 Mon Sep 17 00:00:00 2001
From: Dayue Gao <dygao66@gmail.com>
Date: Tue, 28 Aug 2018 07:28:29 +0800
Subject: [PATCH] fix incorrect check of maxSemiJoinRowsInMemory (#6242)

---
 .../druid/sql/calcite/rel/DruidSemiJoin.java  | 10 ++++---
 .../druid/sql/calcite/CalciteQueryTest.java   | 26 +++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/sql/src/main/java/io/druid/sql/calcite/rel/DruidSemiJoin.java b/sql/src/main/java/io/druid/sql/calcite/rel/DruidSemiJoin.java
index f54fd51c68e..da50c04c8fb 100644
--- a/sql/src/main/java/io/druid/sql/calcite/rel/DruidSemiJoin.java
+++ b/sql/src/main/java/io/druid/sql/calcite/rel/DruidSemiJoin.java
@@ -288,6 +288,8 @@ public class DruidSemiJoin extends DruidRel<DruidSemiJoin>
         new ArrayList<>(),
         new Accumulator<List<RexNode>, Object[]>()
         {
+          int numRows;
+
           @Override
           public List<RexNode> accumulate(final List<RexNode> theConditions, final Object[] row)
           {
@@ -301,14 +303,14 @@ public class DruidSemiJoin extends DruidRel<DruidSemiJoin>
               }
               final String stringValue = DimensionHandlerUtils.convertObjectToString(value);
               values.add(stringValue);
-              if (values.size() > maxSemiJoinRowsInMemory) {
+            }
+
+            if (valuess.add(values)) {
+              if (++numRows > maxSemiJoinRowsInMemory) {
                 throw new ResourceLimitExceededException(
                     StringUtils.format("maxSemiJoinRowsInMemory[%,d] exceeded", maxSemiJoinRowsInMemory)
                 );
               }
-            }
-
-            if (valuess.add(values)) {
               final List<RexNode> subConditions = new ArrayList<>();
 
               for (int i = 0; i < values.size(); i++) {
diff --git a/sql/src/test/java/io/druid/sql/calcite/CalciteQueryTest.java b/sql/src/test/java/io/druid/sql/calcite/CalciteQueryTest.java
index d8b01567ba6..d8cdd14603a 100644
--- a/sql/src/test/java/io/druid/sql/calcite/CalciteQueryTest.java
+++ b/sql/src/test/java/io/druid/sql/calcite/CalciteQueryTest.java
@@ -39,6 +39,7 @@ import io.druid.query.Query;
 import io.druid.query.QueryContexts;
 import io.druid.query.QueryDataSource;
 import io.druid.query.QueryRunnerFactoryConglomerate;
+import io.druid.query.ResourceLimitExceededException;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.CountAggregatorFactory;
 import io.druid.query.aggregation.DoubleMaxAggregatorFactory;
@@ -187,6 +188,13 @@ public class CalciteQueryTest extends CalciteTestBase
       return DateTimes.inferTzfromString("America/Los_Angeles");
     }
   };
+  private static final PlannerConfig PLANNER_CONFIG_SEMI_JOIN_ROWS_LIMIT = new PlannerConfig() {
+    @Override
+    public int getMaxSemiJoinRowsInMemory()
+    {
+      return 2;
+    }
+  };
 
   private static final String LOS_ANGELES = "America/Los_Angeles";
 
@@ -4696,6 +4704,24 @@ public class CalciteQueryTest extends CalciteTestBase
     );
   }
 
+  @Test
+  public void testMaxSemiJoinRowsInMemory() throws Exception
+  {
+    expectedException.expect(ResourceLimitExceededException.class);
+    expectedException.expectMessage("maxSemiJoinRowsInMemory[2] exceeded");
+    testQuery(
+        PLANNER_CONFIG_SEMI_JOIN_ROWS_LIMIT,
+        "SELECT COUNT(*)\n"
+        + "FROM druid.foo\n"
+        + "WHERE SUBSTRING(dim2, 1, 1) IN (\n"
+        + "  SELECT SUBSTRING(dim1, 1, 1) FROM druid.foo WHERE dim1 <> ''\n"
+        + ")\n",
+        CalciteTests.REGULAR_USER_AUTH_RESULT,
+        ImmutableList.of(),
+        ImmutableList.of()
+    );
+  }
+
   @Test
   public void testExplainExactCountDistinctOfSemiJoinResult() throws Exception
   {