Fix operator tests for transient null TsBlocks

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/AlignedSeriesAggregationScanOperatorTest.java

@@ -69,6 +69,7 @@
 import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
 import static org.apache.iotdb.db.queryengine.execution.operator.AggregationOperatorTest.TEST_TIME_SLICE;
 import static org.apache.iotdb.db.queryengine.execution.operator.AggregationUtil.initTimeRangeIterator;
+import static org.apache.iotdb.db.queryengine.execution.operator.OperatorTestUtils.nextNonNullOrEmpty;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -122,7 +123,7 @@ public void testAggregationWithoutTimeFilter() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, true, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       for (int i = 0; i < measurementSchemas.size(); i++) {
         assertEquals(500, resultTsBlock.getColumn(i).getLong(0));
       }
@@ -153,7 +154,7 @@ public void testAggregationWithoutTimeFilterOrderByTimeDesc() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, false, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       for (int i = 0; i < measurementSchemas.size(); i++) {
         assertEquals(500, resultTsBlock.getColumn(i).getLong(0));
       }
@@ -189,7 +190,7 @@ public void testMultiAggregationFuncWithoutTimeFilter1() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, true, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       assertEquals(500, resultTsBlock.getColumn(0).getLong(0));
       assertEquals(6524750.0, resultTsBlock.getColumn(1).getDouble(0), 0.0001);
       count++;
@@ -227,7 +228,7 @@ public void testMultiAggregationFuncWithoutTimeFilter2() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, true, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       assertTrue(resultTsBlock.getColumn(0).getBoolean(0));
       assertEquals(10499, resultTsBlock.getColumn(1).getInt(0));
       assertEquals(20199, resultTsBlock.getColumn(2).getLong(0));
@@ -269,7 +270,7 @@ public void testMultiAggregationFuncWithoutTimeFilterOrderByTimeDesc() throws Ex
         initAlignedSeriesAggregationScanOperator(aggregators, null, false, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       assertTrue(resultTsBlock.getColumn(0).getBoolean(0));
       assertEquals(10499, resultTsBlock.getColumn(1).getInt(0));
       assertEquals(20199, resultTsBlock.getColumn(2).getLong(0));
@@ -304,7 +305,7 @@ public void testAggregationWithTimeFilter1() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, timeFilter, true, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       for (int i = 0; i < measurementSchemas.size(); i++) {
         assertEquals(resultTsBlock.getColumn(i).getLong(0), 380);
       }
@@ -337,7 +338,7 @@ public void testAggregationWithTimeFilter2() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, timeFilter, true, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       for (int i = 0; i < measurementSchemas.size(); i++) {
         assertEquals(resultTsBlock.getColumn(i).getLong(0), 380);
       }
@@ -370,7 +371,7 @@ public void testAggregationWithTimeFilter3() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, timeFilter, true, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       for (int i = 0; i < measurementSchemas.size(); i++) {
         assertEquals(resultTsBlock.getColumn(i).getLong(0), 300);
       }
@@ -410,7 +411,7 @@ public void testMultiAggregationWithTimeFilter() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, timeFilter, true, null);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       assertTrue(resultTsBlock.getColumn(0).getBoolean(0));
       assertEquals(399, resultTsBlock.getColumn(1).getInt(0));
       assertEquals(20199, resultTsBlock.getColumn(2).getLong(0));
@@ -448,7 +449,7 @@ public void testGroupByWithoutGlobalTimeFilter() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, true, groupByTimeParameter);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       int positionCount = resultTsBlock.getPositionCount();
       for (int pos = 0; pos < positionCount; pos++) {
         assertEquals(100 * count, resultTsBlock.getTimeColumn().getLong(pos));
@@ -489,7 +490,7 @@ public void testGroupByWithGlobalTimeFilter() throws Exception {
             aggregators, timeFilter, true, groupByTimeParameter);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       int positionCount = resultTsBlock.getPositionCount();
       for (int pos = 0; pos < positionCount; pos++) {
         assertEquals(100 * count, resultTsBlock.getTimeColumn().getLong(pos));
@@ -534,7 +535,7 @@ public void testGroupByWithMultiFunction() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, true, groupByTimeParameter);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       int positionCount = resultTsBlock.getPositionCount();
       for (int pos = 0; pos < positionCount; pos++) {
         assertEquals(100 * count, resultTsBlock.getTimeColumn().getLong(pos));
@@ -580,7 +581,7 @@ public void testGroupByWithMultiFunctionOrderByTimeDesc() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, false, groupByTimeParameter);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       int positionCount = resultTsBlock.getPositionCount();
       for (int pos = 0; pos < positionCount; pos++) {
         assertEquals(100 * (3 - count), resultTsBlock.getTimeColumn().getLong(pos));
@@ -616,7 +617,7 @@ public void testGroupBySlidingTimeWindow() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, true, groupByTimeParameter);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       int positionCount = resultTsBlock.getPositionCount();
       for (int pos = 0; pos < positionCount; pos++) {
         assertEquals(50 * count, resultTsBlock.getTimeColumn().getLong(pos));
@@ -649,7 +650,7 @@ public void testGroupBySlidingTimeWindow2() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, true, groupByTimeParameter);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       int positionCount = resultTsBlock.getPositionCount();
       for (int pos = 0; pos < positionCount; pos++) {
         assertEquals(timeColumn[count], resultTsBlock.getTimeColumn().getLong(pos));
@@ -693,7 +694,7 @@ public void testGroupBySlidingWindowWithMultiFunction() throws Exception {
         initAlignedSeriesAggregationScanOperator(aggregators, null, true, groupByTimeParameter);
     int count = 0;
     while (seriesAggregationScanOperator.hasNext()) {
-      TsBlock resultTsBlock = seriesAggregationScanOperator.next();
+      TsBlock resultTsBlock = nextNonNullOrEmpty(seriesAggregationScanOperator);
       int positionCount = resultTsBlock.getPositionCount();
       for (int pos = 0; pos < positionCount; pos++) {
         assertEquals(timeColumn[count], resultTsBlock.getTimeColumn().getLong(pos));

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/DeviceViewIntoOperatorTest.java

@@ -72,6 +72,7 @@
 
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
+import static org.apache.iotdb.db.queryengine.execution.operator.OperatorTestUtils.lastNonNullOrEmpty;
 import static org.apache.iotdb.rpc.RpcUtils.SUCCESS_STATUS;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -258,10 +259,7 @@ public void testSingleDeviceSmallData() throws Exception {
     prepareDeviceData("device0", 2);
     operator = createAndInitOperatorForSingleDevices(2);
 
-    TsBlock result = null;
-    while (operator.isBlocked().isDone() && operator.hasNext()) {
-      result = operator.next();
-    }
+    TsBlock result = lastNonNullOrEmpty(operator);
     assertNotNull(result);
     assertEquals(2, result.getPositionCount());
 
@@ -276,10 +274,7 @@ public void testSingleDeviceExceedsMaxTsBlockSize() throws Exception {
     prepareDeviceData("device0", 10);
     operator = createAndInitOperatorForSingleDevices(10);
 
-    TsBlock result = null;
-    while (operator.isBlocked().isDone() && operator.hasNext()) {
-      result = operator.next();
-    }
+    TsBlock result = lastNonNullOrEmpty(operator);
     assertNotNull(result);
     assertEquals(10, result.getPositionCount());
 
@@ -295,10 +290,7 @@ public void testMultipleDeviceSmallData() throws Exception {
     prepareDeviceData("device1", 1);
     operator = createAndInitOperatorForMultipleDevices(2, 1);
 
-    TsBlock result = null;
-    while (operator.isBlocked().isDone() && operator.hasNext()) {
-      result = operator.next();
-    }
+    TsBlock result = lastNonNullOrEmpty(operator);
     assertNotNull(result);
     assertEquals(2, result.getPositionCount());
 

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/FillOperatorTest.java

@@ -43,6 +43,7 @@
 import java.util.concurrent.ExecutorService;
 
 import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
+import static org.apache.iotdb.db.queryengine.execution.operator.OperatorTestUtils.nextNonNullOrEmpty;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -210,7 +211,7 @@ public long ramBytesUsed() {
             }
           };
       while (fillOperator.hasNext()) {
-        TsBlock block = fillOperator.next();
+        TsBlock block = nextNonNullOrEmpty(fillOperator);
         for (int i = 0; i < block.getPositionCount(); i++) {
           long expectedTime = i + 1 + count * 10000L;
           assertEquals(expectedTime, block.getTimeByIndex(i));
@@ -383,7 +384,7 @@ public long ramBytesUsed() {
             }
           };
       while (fillOperator.hasNext()) {
-        TsBlock block = fillOperator.next();
+        TsBlock block = nextNonNullOrEmpty(fillOperator);
         for (int i = 0; i < block.getPositionCount(); i++) {
           long expectedTime = i + 1 + count * 10000L;
           assertEquals(expectedTime, block.getTimeByIndex(i));
@@ -560,7 +561,7 @@ public long ramBytesUsed() {
             }
           };
       while (fillOperator.hasNext()) {
-        TsBlock block = fillOperator.next();
+        TsBlock block = nextNonNullOrEmpty(fillOperator);
         for (int i = 0; i < block.getPositionCount(); i++) {
           long expectedTime = i + 1 + count * 10000L;
           assertEquals(expectedTime, block.getTimeByIndex(i));

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/HorizontallyConcatOperatorTest.java

@@ -68,6 +68,7 @@
 
 import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
 import static org.apache.iotdb.db.queryengine.execution.operator.AggregationUtil.initTimeRangeIterator;
+import static org.apache.iotdb.db.queryengine.execution.operator.OperatorTestUtils.nextNonNullOrEmpty;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
 
@@ -203,7 +204,7 @@ public void batchTest1() throws Exception {
       int count = 0;
       while (horizontallyConcatOperator.isBlocked().isDone()
           && horizontallyConcatOperator.hasNext()) {
-        TsBlock tsBlock = horizontallyConcatOperator.next();
+        TsBlock tsBlock = nextNonNullOrEmpty(horizontallyConcatOperator);
         assertEquals(6, tsBlock.getValueColumnCount());
         for (int i = 0; i < tsBlock.getPositionCount(); i++, count++) {
           assertEquals(count, tsBlock.getTimeByIndex(i));

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java

@@ -93,6 +93,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
+import static org.apache.iotdb.db.queryengine.execution.operator.OperatorTestUtils.nextNonNullOrEmpty;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -1794,7 +1795,7 @@ public void mergeSortTest() throws Exception {
 
     int index = 0;
     while (treeMergeSortOperator.isBlocked().isDone() && treeMergeSortOperator.hasNext()) {
-      TsBlock result = treeMergeSortOperator.next();
+      TsBlock result = nextNonNullOrEmpty(treeMergeSortOperator);
       for (int i = 0; i < result.getPositionCount(); i++) {
         long time = result.getTimeByIndex(i);
         assertEquals(time, ans[index++]);

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/OffsetOperatorTest.java

@@ -18,6 +18,7 @@
  */
 package org.apache.iotdb.db.queryengine.execution.operator;
 
+import org.apache.iotdb.calc.execution.operator.Operator;
 import org.apache.iotdb.calc.execution.operator.process.LimitOperator;
 import org.apache.iotdb.calc.execution.operator.process.OffsetOperator;
 import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
@@ -65,6 +66,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
+import static org.apache.iotdb.db.queryengine.execution.operator.OperatorTestUtils.nextNonNullOrEmpty;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -168,7 +170,7 @@ public void batchTest1() throws Exception {
           new LimitOperator(driverContext.getOperatorContexts().get(4), 250, offsetOperator);
       int count = 100;
       while (limitOperator.isBlocked().isDone() && limitOperator.hasNext()) {
-        TsBlock tsBlock = limitOperator.next();
+        TsBlock tsBlock = nextNonNullOrEmpty(limitOperator);
         assertEquals(2, tsBlock.getValueColumnCount());
         assertTrue(tsBlock.getColumn(0) instanceof IntColumn);
         assertTrue(tsBlock.getColumn(1) instanceof IntColumn);
@@ -276,7 +278,7 @@ public void batchTest2() throws Exception {
 
       int count = 0;
       while (offsetOperator.isBlocked().isDone() && offsetOperator.hasNext()) {
-        TsBlock tsBlock = offsetOperator.next();
+        TsBlock tsBlock = nextNonNullOrEmpty(offsetOperator);
         assertEquals(2, tsBlock.getValueColumnCount());
         assertTrue(tsBlock.getColumn(0) instanceof IntColumn);
         assertTrue(tsBlock.getColumn(1) instanceof IntColumn);
@@ -382,7 +384,7 @@ public void batchTest3() throws Exception {
           new OffsetOperator(driverContext.getOperatorContexts().get(3), 500, timeJoinOperator);
 
       while (offsetOperator.isBlocked().isDone() && offsetOperator.hasNext()) {
-        TsBlock tsBlock = offsetOperator.next();
+        TsBlock tsBlock = nextNonNull(offsetOperator);
         assertEquals(2, tsBlock.getValueColumnCount());
         assertTrue(tsBlock.getColumn(0) instanceof IntColumn);
         assertTrue(tsBlock.getColumn(1) instanceof IntColumn);
@@ -474,7 +476,7 @@ public void batchTest4() throws Exception {
               driverContext.getOperatorContexts().get(3), 98_784_247_808L, timeJoinOperator);
 
       while (offsetOperator.isBlocked().isDone() && offsetOperator.hasNext()) {
-        TsBlock tsBlock = offsetOperator.next();
+        TsBlock tsBlock = nextNonNull(offsetOperator);
         assertEquals(2, tsBlock.getValueColumnCount());
         assertTrue(tsBlock.getColumn(0) instanceof IntColumn);
         assertTrue(tsBlock.getColumn(1) instanceof IntColumn);
@@ -487,4 +489,14 @@ public void batchTest4() throws Exception {
       instanceNotificationExecutor.shutdown();
     }
   }
+
+  private static TsBlock nextNonNull(Operator operator) throws Exception {
+    while (operator.hasNext()) {
+      TsBlock result = operator.next();
+      if (result != null) {
+        return result;
+      }
+    }
+    throw new AssertionError("Expected a non-null TsBlock from operator");
+  }
 }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/OperatorTestUtils.java

@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iotdb.db.queryengine.execution.operator;
+
+import org.apache.iotdb.calc.execution.operator.Operator;
+
+import org.apache.tsfile.read.common.block.TsBlock;
+
+public final class OperatorTestUtils {
+
+  private OperatorTestUtils() {
+    // Utility class.
+  }
+
+  public static TsBlock nextNonNullOrEmpty(Operator operator) throws Exception {
+    while (operator.hasNext()) {
+      TsBlock result = operator.next();
+      if (!isNullOrEmpty(result)) {
+        return result;
+      }
+    }
+    throw new AssertionError("Expected a non-null and non-empty TsBlock from operator");
+  }
+
+  public static TsBlock lastNonNullOrEmpty(Operator operator) throws Exception {
+    TsBlock result = null;
+    while (operator.isBlocked().isDone() && operator.hasNext()) {
+      TsBlock nextResult = operator.next();
+      if (!isNullOrEmpty(nextResult)) {
+        result = nextResult;
+      }
+    }
+    return result;
+  }
+
+  private static boolean isNullOrEmpty(TsBlock tsBlock) {
+    return tsBlock == null || tsBlock.getPositionCount() == 0;
+  }
+}