Add size control of result for SimpleNestedLoopCrossJoinOperator

Signed-off-by: Weihao Li <[email protected]>

Author: Wei-hao-Li

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/join/SimpleNestedLoopCrossJoinOperator.java

@@ -120,7 +120,8 @@ public TsBlock next() throws Exception {
       cachedProbeBlock = null;
       return null;
     }
-    while (probeIndex < cachedProbeBlock.getPositionCount()
+    while (!resultBuilder.isFull()
+        && probeIndex < cachedProbeBlock.getPositionCount()
         && System.nanoTime() - start < maxRuntime) {
       for (TsBlock buildBlock : buildBlocks) {
         appendValueToResult(probeIndex, buildBlock);

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/JoinTest.java

@@ -19,9 +19,21 @@
 
 package org.apache.iotdb.db.queryengine.plan.relational.analyzer;
 
+import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
+import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
+import org.apache.iotdb.db.queryengine.common.PlanFragmentId;
+import org.apache.iotdb.db.queryengine.common.QueryId;
+import org.apache.iotdb.db.queryengine.execution.driver.DriverContext;
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine;
+import org.apache.iotdb.db.queryengine.execution.operator.Operator;
+import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext;
+import org.apache.iotdb.db.queryengine.execution.operator.process.join.SimpleNestedLoopCrossJoinOperator;
+import org.apache.iotdb.db.queryengine.execution.operator.source.relational.TableScanOperator;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.DistributedQueryPlan;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.IdentitySinkNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.PlanTester;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.Symbol;
@@ -47,13 +59,21 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.tsfile.block.column.ColumnBuilder;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.tsfile.read.common.block.column.RunLengthEncodedColumn;
 import org.junit.Ignore;
 import org.junit.Test;
 
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 
+import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
+import static org.apache.iotdb.db.queryengine.execution.operator.source.relational.AbstractTableScanOperator.TIME_COLUMN_TEMPLATE;
 import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.AnalyzerTest.analyzeSQL;
 import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.ALL_DEVICE_ENTRIES;
 import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.BEIJING_A1_DEVICE_ENTRY;
@@ -87,6 +107,7 @@
 import static org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression.Operator.GREATER_THAN;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 public class JoinTest {
   Analysis analysis;
@@ -811,4 +832,119 @@ public void joinSortEliminationTest() {
 
     assertPlan(planTester.getFragmentPlan(8), tableScan("testdb.table1"));
   }
+
+  @Test
+  // case: lines of result are more than Integer.MAX_VALUE
+  public void crossJoinLargeDataTest() {
+    // lines = 46341
+    int lines = ((int) Math.sqrt(Integer.MAX_VALUE)) + 1;
+    try (SimpleNestedLoopCrossJoinOperator aggregationOperator =
+        genSimpleNestedLoopCrossJoinOperator(lines)) {
+      ListenableFuture<?> listenableFuture = aggregationOperator.isBlocked();
+      listenableFuture.get();
+      long count = 0;
+      while (!aggregationOperator.isFinished() && aggregationOperator.hasNext()) {
+        TsBlock tsBlock = aggregationOperator.next();
+        if (tsBlock != null) {
+          count += tsBlock.getPositionCount();
+        }
+        listenableFuture = aggregationOperator.isBlocked();
+        listenableFuture.get();
+      }
+      assertEquals((long) lines * lines, count);
+    } catch (Exception e) {
+      e.printStackTrace();
+      fail(e.getMessage());
+    }
+  }
+
+  private SimpleNestedLoopCrossJoinOperator genSimpleNestedLoopCrossJoinOperator(int lines) {
+
+    // Construct operator tree
+    QueryId queryId = new QueryId("stub_query");
+
+    FragmentInstanceId instanceId =
+        new FragmentInstanceId(new PlanFragmentId(queryId, 0), "stub-instance");
+    FragmentInstanceStateMachine stateMachine =
+        new FragmentInstanceStateMachine(
+            instanceId,
+            IoTDBThreadPoolFactory.newFixedThreadPool(
+                1, "SimpleNestedLoopCrossJoinOperator-test-instance-notification"));
+    FragmentInstanceContext fragmentInstanceContext =
+        createFragmentInstanceContext(instanceId, stateMachine);
+    DriverContext driverContext = new DriverContext(fragmentInstanceContext, 0);
+    driverContext.addOperatorContext(
+        1, new PlanNodeId("1"), TableScanOperator.class.getSimpleName());
+    driverContext.addOperatorContext(
+        2, new PlanNodeId("2"), TableScanOperator.class.getSimpleName());
+    driverContext.addOperatorContext(
+        3, new PlanNodeId("3"), SimpleNestedLoopCrossJoinOperator.class.getSimpleName());
+    return new SimpleNestedLoopCrossJoinOperator(
+        driverContext.getOperatorContexts().get(2),
+        genChildOperator(driverContext.getOperatorContexts().get(0), lines),
+        genChildOperator(driverContext.getOperatorContexts().get(1), lines),
+        new int[0],
+        new int[1],
+        Collections.singletonList(TSDataType.TIMESTAMP));
+  }
+
+  private Operator genChildOperator(OperatorContext operatorContext, int lines) {
+    return new Operator() {
+      boolean finished = false;
+
+      @Override
+      public OperatorContext getOperatorContext() {
+        return operatorContext;
+      }
+
+      @Override
+      public TsBlock next() {
+        TsBlockBuilder builder =
+            new TsBlockBuilder(Collections.singletonList(TSDataType.TIMESTAMP));
+        ColumnBuilder columnBuilder = builder.getValueColumnBuilders()[0];
+        for (int i = 0; i < lines; i++) {
+          columnBuilder.writeLong(1);
+        }
+        builder.declarePositions(lines);
+        TsBlock result =
+            builder.build(
+                new RunLengthEncodedColumn(TIME_COLUMN_TEMPLATE, builder.getPositionCount()));
+        finished = true;
+        return result;
+      }
+
+      @Override
+      public boolean hasNext() throws Exception {
+        return !finished;
+      }
+
+      @Override
+      public void close() throws Exception {}
+
+      @Override
+      public boolean isFinished() throws Exception {
+        return finished;
+      }
+
+      @Override
+      public long calculateMaxPeekMemory() {
+        return 0;
+      }
+
+      @Override
+      public long calculateMaxReturnSize() {
+        return 0;
+      }
+
+      @Override
+      public long calculateRetainedSizeAfterCallingNext() {
+        return 0;
+      }
+
+      @Override
+      public long ramBytesUsed() {
+        return 0;
+      }
+    };
+  }
 }