Support aggregation push-down optimize for table view of non aligned series

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/AbstractAggTableScanOperator.java

@@ -34,6 +34,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.AlignedDeviceEntry;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.ColumnSchema;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.Symbol;
 import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
 import org.apache.iotdb.db.storageengine.dataregion.read.IQueryDataSource;
 import org.apache.iotdb.db.storageengine.dataregion.read.QueryDataSource;
@@ -71,7 +72,7 @@
 public abstract class AbstractAggTableScanOperator extends AbstractDataSourceOperator {
 
   private boolean finished = false;
-  private TsBlock inputTsBlock;
+  protected TsBlock inputTsBlock;
 
   protected List<TableAggregator> tableAggregators;
   protected final List<ColumnSchema> groupingKeySchemas;
@@ -104,11 +105,11 @@
   // e.g. for aggregation `last(s1), count(s2), count(s1)`, the inputChannels should be [0, 1, 0]
   protected List<Integer> aggregatorInputChannels;
 
-  private QueryDataSource queryDataSource;
+  protected QueryDataSource queryDataSource;
 
   protected ITableTimeRangeIterator timeIterator;
 
-  private boolean allAggregatorsHasFinalResult = false;
+  protected boolean allAggregatorsHasFinalResult = false;
 
   protected AbstractAggTableScanOperator(AbstractAggTableScanOperatorParameter parameter) {
 
@@ -193,7 +194,7 @@
   }
 
   /** Return true if we have the result of this timeRange. */
-  protected Optional<Boolean> calculateAggregationResultForCurrentTimeRange() {
+  protected Optional<Boolean> calculateAggregationResultForCurrentTimeRange() throws Exception {
     try {
       if (calcFromCachedData()) {
         updateResultTsBlock();
@@ -706,7 +707,7 @@
     return true;
   }
 
-  private void checkIfAllAggregatorHasFinalResult() {
+  protected void checkIfAllAggregatorHasFinalResult() throws Exception {
     if (allAggregatorsHasFinalResult
         && (timeIterator.getType() == ITableTimeRangeIterator.TimeIteratorType.SINGLE_TIME_ITERATOR
             || tableAggregators.isEmpty())) {
@@ -729,7 +730,7 @@
     }
   }
 
-  private void nextDevice() {
+  protected void nextDevice() throws Exception {
     currentDeviceIndex++;
     this.operatorContext.recordSpecifiedInfo(
         CURRENT_DEVICE_INDEX_STRING, Integer.toString(currentDeviceIndex));
@@ -812,6 +813,8 @@
     protected List<DeviceEntry> deviceEntries;
     protected int deviceCount;
 
+    private List<Symbol> outputSymbols;
+
     public AbstractAggTableScanOperatorParameter(
         PlanNodeId sourceId,
         OperatorContext context,
@@ -830,7 +833,8 @@
         boolean ascending,
         boolean canUseStatistics,
         List<Integer> aggregatorInputChannels,
-        String timeColumnName) {
+        String timeColumnName,
+        List<Symbol> outputSymbols) {
       this.sourceId = sourceId;
       this.context = context;
       this.aggColumnSchemas = aggColumnSchemas;
@@ -849,6 +853,11 @@
       this.canUseStatistics = canUseStatistics;
       this.aggregatorInputChannels = aggregatorInputChannels;
       this.timeColumnName = timeColumnName;
+      this.outputSymbols = outputSymbols;
+    }
+
+    public List<Symbol> getOutputSymbols() {
+      return outputSymbols;
     }
 
     public OperatorContext getOperatorContext() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/DeviceIteratorScanOperator.java

@@ -113,7 +113,7 @@ private void constructCurrentDeviceOperatorTree() {
     }
     DeviceEntry deviceEntry = this.deviceEntries.get(this.currentDeviceIndex);
 
-    deviceChildOperatorTreeGenerator.generateCurrentDeviceOperatorTree(deviceEntry);
+    deviceChildOperatorTreeGenerator.generateCurrentDeviceOperatorTree(deviceEntry, true);
     currentDeviceRootOperator = deviceChildOperatorTreeGenerator.getCurrentDeviceRootOperator();
     dataSourceOperators = deviceChildOperatorTreeGenerator.getCurrentDeviceDataSourceOperators();
     currentDeviceInit = false;
@@ -188,6 +188,10 @@ public long ramBytesUsed() {
         + RamUsageEstimator.sizeOfCollection(deviceEntries);
   }
 
+  public DeviceChildOperatorTreeGenerator getDeviceChildOperatorTreeGenerator() {
+    return deviceChildOperatorTreeGenerator;
+  }
+
   public static class TreeNonAlignedDeviceViewScanParameters {
     public final OperatorContext context;
     public final List<DeviceEntry> deviceEntries;
@@ -217,7 +221,7 @@ public interface DeviceChildOperatorTreeGenerator {
     boolean keepOffsetAndLimitOperatorAfterDeviceIterator();
 
     // Generate the following operator subtree based on the current deviceEntry
-    void generateCurrentDeviceOperatorTree(DeviceEntry deviceEntry);
+    void generateCurrentDeviceOperatorTree(DeviceEntry deviceEntry, boolean needAdaptor);
 
     // Returns the root operator of the subtree
     Operator getCurrentDeviceRootOperator();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/LastQueryAggTableScanOperator.java

@@ -156,7 +156,7 @@
   }
 
   /** Main process logic, calc the last aggregation results of current device. */
-  private void processCurrentDevice() {
+  private void processCurrentDevice() throws Exception {
     if (currentHitCacheIndex < hitCachesIndexes.size()
         && outputDeviceIndex == hitCachesIndexes.get(currentHitCacheIndex)) {
       currentDeviceEntry = cachedDeviceEntries.get(currentHitCacheIndex);
@@ -175,7 +175,7 @@
     }
   }
 
   private void buildResultUseLastRowCache() {
     appendGroupKeysToResult(cachedDeviceEntries, currentHitCacheIndex);
     Pair<OptionalLong, TsPrimitiveType[]> currentHitResult =
         lastRowCacheResults.get(currentHitCacheIndex);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/TreeNonAlignedDeviceViewAggregationScanOperator.java

@@ -0,0 +1,219 @@
+/*
+ * 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.source.relational;
+
+import org.apache.iotdb.db.queryengine.execution.MemoryEstimationHelper;
+import org.apache.iotdb.db.queryengine.execution.aggregation.timerangeiterator.ITableTimeRangeIterator;
+import org.apache.iotdb.db.queryengine.execution.operator.Operator;
+import org.apache.iotdb.db.queryengine.execution.operator.source.AbstractDataSourceOperator;
+import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry;
+import org.apache.iotdb.db.storageengine.dataregion.read.IQueryDataSource;
+import org.apache.iotdb.db.storageengine.dataregion.read.QueryDataSource;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.tsfile.file.metadata.IDeviceID;
+import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.read.common.block.TsBlockBuilder;
+import org.apache.tsfile.utils.RamUsageEstimator;
+
+import java.util.List;
+import java.util.Optional;
+
+import static org.apache.iotdb.db.queryengine.execution.operator.source.relational.AbstractTableScanOperator.CURRENT_DEVICE_INDEX_STRING;
+
+public class TreeNonAlignedDeviceViewAggregationScanOperator
+    extends AbstractDefaultAggTableScanOperator {
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(
+          TreeNonAlignedDeviceViewAggregationScanOperator.class);
+
+  private final IDeviceID.TreeDeviceIdColumnValueExtractor extractor;
+  private final DeviceIteratorScanOperator.DeviceChildOperatorTreeGenerator childOperatorGenerator;
+
+  private Operator child;
+  private List<Operator> dataSourceOperators;
+
+  public TreeNonAlignedDeviceViewAggregationScanOperator(
+      AbstractAggTableScanOperatorParameter parameter,
+      IDeviceID.TreeDeviceIdColumnValueExtractor extractor,
+      DeviceIteratorScanOperator.DeviceChildOperatorTreeGenerator childOperatorGenerator) {
+    super(parameter);
+    this.extractor = extractor;
+    this.childOperatorGenerator = childOperatorGenerator;
+    constructCurrentDeviceOperatorTree();
+  }
+
+  @Override
+  public ListenableFuture<?> isBlocked() {
+    return child.isBlocked();
+  }
+
+  @Override
+  String getNthIdColumnValue(DeviceEntry deviceEntry, int idColumnIndex) {
+    return (String) extractor.extract(deviceEntry.getDeviceID(), idColumnIndex);
+  }
+
+  @Override
+  protected Optional<Boolean> calculateAggregationResultForCurrentTimeRange() {
+    try {
+      // First try to calculate from cached data
+      if (calcFromCachedData()) {
+        updateResultTsBlock();
+        checkIfAllAggregatorHasFinalResult();
+        return Optional.of(true);
+      }
+
+      // Read from child operator
+      if (readAndCalcFromChild()) {
+        updateResultTsBlock();
+        checkIfAllAggregatorHasFinalResult();
+        return Optional.of(true);
+      }
+
+      // No more data from child, finish the current device
+      if (!child.hasNext()) {
+        updateResultTsBlock();
+        timeIterator.resetCurTimeRange();
+        nextDevice();
+
+        if (currentDeviceIndex >= deviceCount) {
+          // All devices consumed
+          timeIterator.setFinished();
+          return Optional.of(true);
+        } else {
+          // More devices to process, child should provide next device's data
+          return Optional.of(false);
+        }
+      }
+
+      return Optional.of(false);
+    } catch (Exception e) {
+      throw new RuntimeException("Error while processing aggregation from child operator", e);
+    }
+  }
+
+  /** Read data from child operator and calculate aggregation. */
+  private boolean readAndCalcFromChild() throws Exception {
+    long start = System.nanoTime();
+
+    while (System.nanoTime() - start < leftRuntimeOfOneNextCall && child.hasNext()) {
+      // Get next TsBlock from child
+      TsBlock tsBlock = child.nextWithTimer();
+      if (tsBlock == null || tsBlock.isEmpty()) {
+        continue;
+      }
+
+      // Calculate aggregation from raw data
+      if (calcUsingRawData(tsBlock)) {
+        return true;
+      }
+
+      // If not finished, continue reading from  child
+    }
+
+    return false;
+  }
+
+  @Override
+  protected void nextDevice() throws Exception {
+    currentDeviceIndex++;
+    childOperatorGenerator.getCurrentDeviceStartCloseOperator().close();
+    if (currentDeviceIndex >= deviceEntries.size()) {
+      return;
+    }
+    constructCurrentDeviceOperatorTree();
+    queryDataSource.reset();
+    initQueryDataSource(queryDataSource);
+    this.operatorContext.recordSpecifiedInfo(
+        CURRENT_DEVICE_INDEX_STRING, Integer.toString(currentDeviceIndex));
+  }
+
+  private void constructCurrentDeviceOperatorTree() {
+    if (this.deviceEntries.isEmpty()) {
+      return;
+    }
+    if (this.deviceEntries.get(this.currentDeviceIndex) == null) {
+      throw new IllegalStateException(
+          "Device entries of index " + this.currentDeviceIndex + " is empty");
+    }
+    DeviceEntry deviceEntry = this.deviceEntries.get(this.currentDeviceIndex);
+
+    childOperatorGenerator.generateCurrentDeviceOperatorTree(deviceEntry, false);
+    child = childOperatorGenerator.getCurrentDeviceRootOperator();
+    dataSourceOperators = childOperatorGenerator.getCurrentDeviceDataSourceOperators();
+  }
+
+  /** same with {@link DeviceIteratorScanOperator#initQueryDataSource(IQueryDataSource)} */
+  @Override
+  public void initQueryDataSource(IQueryDataSource dataSource) {
+
+    this.queryDataSource = (QueryDataSource) dataSource;
+    this.resultTsBlockBuilder = new TsBlockBuilder(getResultDataTypes());
+    if (dataSourceOperators == null || dataSourceOperators.isEmpty()) {
+      return;
+    }
+    for (Operator operator : dataSourceOperators) {
+      ((AbstractDataSourceOperator) operator).initQueryDataSource(dataSource);
+    }
+  }
+
+  @Override
+  protected void checkIfAllAggregatorHasFinalResult() throws Exception {
+    if (allAggregatorsHasFinalResult
+        && (timeIterator.getType() == ITableTimeRangeIterator.TimeIteratorType.SINGLE_TIME_ITERATOR
+            || tableAggregators.isEmpty())) {
+      nextDevice();
+      inputTsBlock = null;
+
+      if (currentDeviceIndex >= deviceCount) {
+        // all devices have been consumed
+        timeIterator.setFinished();
+      }
+
+      allAggregatorsHasFinalResult = false;
+    }
+  }
+
+  @Override
+  public long calculateMaxPeekMemory() {
+    return child.calculateMaxPeekMemory();
+  }
+
+  @Override
+  public long calculateMaxReturnSize() {
+    return child.calculateMaxReturnSize();
+  }
+
+  @Override
+  public long calculateRetainedSizeAfterCallingNext() {
+    return child.calculateRetainedSizeAfterCallingNext();
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    return INSTANCE_SIZE
+        + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(seriesScanUtil)
+        + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext)
+        + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(sourceId)
+        + (resultTsBlockBuilder == null ? 0 : resultTsBlockBuilder.getRetainedSizeInBytes())
+        + RamUsageEstimator.sizeOfCollection(deviceEntries)
+        + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(child);
+  }
+}