[To dev/1.3] Pipe: Handle hybrid meta progress indexes (#18331)

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeTaskAgent.java

@@ -50,6 +50,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
@@ -177,10 +178,13 @@ private void clearConfigRegionListeningQueueIfNecessary(
         }
 
         final ProgressIndex progressIndex = groupId2TaskMetaMap.get(regionId).getProgressIndex();
-        if (progressIndex instanceof MetaProgressIndex) {
-          if (((MetaProgressIndex) progressIndex).getIndex() + 1
-              < listeningQueueNewFirstIndex.get()) {
-            listeningQueueNewFirstIndex.set(((MetaProgressIndex) progressIndex).getIndex() + 1);
+        final Optional<MetaProgressIndex> metaProgressIndex =
+            Objects.isNull(progressIndex)
+                ? Optional.empty()
+                : progressIndex.getProgressIndexByType(MetaProgressIndex.class);
+        if (metaProgressIndex.isPresent()) {
+          if (metaProgressIndex.get().getIndex() + 1 < listeningQueueNewFirstIndex.get()) {
+            listeningQueueNewFirstIndex.set(metaProgressIndex.get().getIndex() + 1);
           }
         } else {
           // Do not clear "minimumProgressIndex"s related queues to avoid clearing

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgent.java

@@ -240,11 +240,15 @@ private Set<Integer> clearSchemaRegionListeningQueueIfNecessary(
         }
 
         final ProgressIndex progressIndex = pipeTaskMeta.getProgressIndex();
-        if (progressIndex instanceof MetaProgressIndex) {
-          if (((MetaProgressIndex) progressIndex).getIndex() + 1
+        final Optional<MetaProgressIndex> metaProgressIndex =
+            Objects.isNull(progressIndex)
+                ? Optional.empty()
+                : progressIndex.getProgressIndexByType(MetaProgressIndex.class);
+        if (metaProgressIndex.isPresent()) {
+          if (metaProgressIndex.get().getIndex() + 1
               < schemaRegionId2ListeningQueueNewFirstIndex.getOrDefault(id, Long.MAX_VALUE)) {
             schemaRegionId2ListeningQueueNewFirstIndex.put(
-                id, ((MetaProgressIndex) progressIndex).getIndex() + 1);
+                id, metaProgressIndex.get().getIndex() + 1);
           }
         } else {
           // Do not clear "minimumProgressIndex"s related queues to avoid clearing

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/processor/aggregate/AggregateProcessor.java

@@ -331,17 +331,15 @@ public void customize(
 
     // Restore window state
     final ProgressIndex index = pipeTaskMeta.getProgressIndex();
-    if (index == MinimumProgressIndex.INSTANCE) {
+    if (Objects.isNull(index) || index == MinimumProgressIndex.INSTANCE) {
       return;
     }
-    if (!(index instanceof TimeWindowStateProgressIndex)) {
-      throw new PipeException(
-          String.format(
-              "The aggregate processor does not support progressIndexType %s", index.getType()));
-    }
-
     final TimeWindowStateProgressIndex timeWindowStateProgressIndex =
-        (TimeWindowStateProgressIndex) index;
+        index.getProgressIndexByType(TimeWindowStateProgressIndex.class).orElse(null);
+    // A pipe altered from another processor may not have window state yet.
+    if (Objects.isNull(timeWindowStateProgressIndex)) {
+      return;
+    }
     for (final Map.Entry<String, Pair<Long, ByteBuffer>> entry :
         timeWindowStateProgressIndex.getTimeSeries2TimestampWindowBufferPairMap().entrySet()) {
       final AtomicReference<TimeSeriesRuntimeState> stateReference =

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/processor/twostage/plugin/TwoStageCountProcessor.java

@@ -136,13 +136,7 @@ public void customize(PipeParameters parameters, PipeProcessorRuntimeConfigurati
     outputSeries = parseOutputSeries(parameters);
 
     if (Objects.nonNull(pipeTaskMeta) && Objects.nonNull(pipeTaskMeta.getProgressIndex())) {
-      if (pipeTaskMeta.getProgressIndex() instanceof MinimumProgressIndex) {
-        pipeTaskMeta.updateProgressIndex(
-            new StateProgressIndex(Long.MIN_VALUE, new HashMap<>(), MinimumProgressIndex.INSTANCE));
-      }
-
-      final StateProgressIndex stateProgressIndex =
-          (StateProgressIndex) pipeTaskMeta.getProgressIndex();
+      final StateProgressIndex stateProgressIndex = initializeStateProgressIndex(pipeTaskMeta);
       localCommitProgressIndex.set(stateProgressIndex.getInnerProgressIndex());
       final Binary localCountState = stateProgressIndex.getState().get(LOCAL_COUNT_STATE_KEY);
       localCount.set(
@@ -173,6 +167,26 @@ static PartialPath parseOutputSeries(final PipeParameters parameters)
             PipeProcessorConstant.PROCESSOR_OUTPUT_SERIES_KEY, LEGACY_PROCESSOR_OUTPUT_SERIES_KEY));
   }
 
+  static StateProgressIndex initializeStateProgressIndex(final PipeTaskMeta pipeTaskMeta) {
+    final ProgressIndex progressIndex = pipeTaskMeta.getProgressIndex();
+    if (progressIndex instanceof StateProgressIndex) {
+      return (StateProgressIndex) progressIndex;
+    }
+
+    final ProgressIndex updatedProgressIndex =
+        pipeTaskMeta.updateProgressIndex(
+            new StateProgressIndex(
+                Long.MIN_VALUE, Collections.emptyMap(), MinimumProgressIndex.INSTANCE));
+    return updatedProgressIndex
+        .getProgressIndexByType(StateProgressIndex.class)
+        .orElseThrow(
+            () ->
+                new PipeException(
+                    String.format(
+                        "Failed to initialize StateProgressIndex from progress index %s.",
+                        updatedProgressIndex)));
+  }
+
   @Override
   public void process(TabletInsertionEvent tabletInsertionEvent, EventCollector eventCollector)
       throws Exception {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/historical/PipeHistoricalDataRegionTsFileSource.java

@@ -398,7 +398,7 @@ public synchronized void start() {
 
         originalResourceList.sort(
             (o1, o2) ->
-                startIndex instanceof TimeWindowStateProgressIndex
+                Objects.nonNull(getTimeWindowStateProgressIndex(startIndex))
                     ? Long.compare(o1.getFileStartTime(), o2.getFileStartTime())
                     : o1.getMaxProgressIndex().topologicalCompareTo(o2.getMaxProgressIndex()));
         pendingQueue = new ArrayDeque<>(originalResourceList);
@@ -466,9 +466,11 @@ private boolean shouldExtractTsFileResource(final TsFileResource resource) {
   }
 
   private boolean mayTsFileContainUnprocessedData(final TsFileResource resource) {
-    if (startIndex instanceof TimeWindowStateProgressIndex) {
+    final TimeWindowStateProgressIndex timeWindowStateProgressIndex =
+        getTimeWindowStateProgressIndex(startIndex);
+    if (Objects.nonNull(timeWindowStateProgressIndex)) {
       // The resource is closed thus the TsFileResource#getFileEndTime() is safe to use
-      return ((TimeWindowStateProgressIndex) startIndex).getMinTime() <= resource.getFileEndTime();
+      return timeWindowStateProgressIndex.getMinTime() <= resource.getFileEndTime();
     }
 
     if (startIndex instanceof StateProgressIndex) {
@@ -488,6 +490,13 @@ private boolean mayTsFileContainUnprocessedData(final TsFileResource resource) {
     return false;
   }
 
+  private TimeWindowStateProgressIndex getTimeWindowStateProgressIndex(
+      final ProgressIndex progressIndex) {
+    return Objects.isNull(progressIndex)
+        ? null
+        : progressIndex.getProgressIndexByType(TimeWindowStateProgressIndex.class).orElse(null);
+  }
+
   private boolean mayTsFileResourceOverlappedWithPattern(final TsFileResource resource) {
     // Trimming to avoid unnecessary file device getter
     if (isDbNameCoveredByPattern) {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/processor/twostage/plugin/TwoStageCountProcessorTest.java

@@ -19,7 +19,14 @@
 
 package org.apache.iotdb.db.pipe.processor.twostage.plugin;
 
+import org.apache.iotdb.commons.consensus.index.ProgressIndex;
+import org.apache.iotdb.commons.consensus.index.impl.HybridProgressIndex;
+import org.apache.iotdb.commons.consensus.index.impl.MetaProgressIndex;
+import org.apache.iotdb.commons.consensus.index.impl.SimpleProgressIndex;
+import org.apache.iotdb.commons.consensus.index.impl.StateProgressIndex;
+import org.apache.iotdb.commons.consensus.index.impl.TimeWindowStateProgressIndex;
 import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
 import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameterValidator;
 import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters;
 
@@ -44,6 +51,54 @@ public void testValidateOutputSeriesSupportsNewAndLegacyKeys() throws Exception
     validateOutputSeries("processor.output-series", "root.db.d.s2");
   }
 
+  @Test
+  public void testInitializeStateProgressIndexFromHybridProgressIndex() {
+    final MetaProgressIndex metaProgressIndex = new MetaProgressIndex(10L);
+    final SimpleProgressIndex simpleProgressIndex = new SimpleProgressIndex(1, 2L);
+    final ProgressIndex hybridProgressIndex =
+        new HybridProgressIndex(metaProgressIndex)
+            .updateToMinimumEqualOrIsAfterProgressIndex(simpleProgressIndex);
+    final PipeTaskMeta pipeTaskMeta = new PipeTaskMeta(hybridProgressIndex, 0);
+
+    final StateProgressIndex stateProgressIndex =
+        TwoStageCountProcessor.initializeStateProgressIndex(pipeTaskMeta);
+
+    Assert.assertSame(stateProgressIndex, pipeTaskMeta.getProgressIndex());
+    Assert.assertEquals(
+        metaProgressIndex,
+        stateProgressIndex.getProgressIndexByType(MetaProgressIndex.class).orElse(null));
+    Assert.assertEquals(
+        simpleProgressIndex,
+        stateProgressIndex.getProgressIndexByType(SimpleProgressIndex.class).orElse(null));
+  }
+
+  @Test
+  public void testInitializeStateProgressIndexFromTimeWindowStateProgressIndex() {
+    final TimeWindowStateProgressIndex timeWindowStateProgressIndex =
+        new TimeWindowStateProgressIndex(Collections.emptyMap());
+    final PipeTaskMeta pipeTaskMeta = new PipeTaskMeta(timeWindowStateProgressIndex, 0);
+
+    final StateProgressIndex stateProgressIndex =
+        TwoStageCountProcessor.initializeStateProgressIndex(pipeTaskMeta);
+    Assert.assertEquals(
+        timeWindowStateProgressIndex,
+        stateProgressIndex.getProgressIndexByType(TimeWindowStateProgressIndex.class).orElse(null));
+
+    final SimpleProgressIndex simpleProgressIndex = new SimpleProgressIndex(1, 2L);
+    final ProgressIndex updatedProgressIndex =
+        pipeTaskMeta.updateProgressIndex(
+            new StateProgressIndex(1L, Collections.emptyMap(), simpleProgressIndex));
+    Assert.assertTrue(updatedProgressIndex instanceof StateProgressIndex);
+    Assert.assertEquals(
+        timeWindowStateProgressIndex,
+        updatedProgressIndex
+            .getProgressIndexByType(TimeWindowStateProgressIndex.class)
+            .orElse(null));
+    Assert.assertEquals(
+        simpleProgressIndex,
+        updatedProgressIndex.getProgressIndexByType(SimpleProgressIndex.class).orElse(null));
+  }
+
   private PartialPath parseOutputSeries(final String key, final String value) throws Exception {
     return TwoStageCountProcessor.parseOutputSeries(
         new PipeParameters(Collections.singletonMap(key, value)));

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/TsFileResourceProgressIndexTest.java

@@ -54,6 +54,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Random;
 import java.util.stream.IntStream;
 
@@ -213,6 +214,14 @@ public ProgressIndexType getType() {
       throw new UnsupportedOperationException("method not implemented.");
     }
 
+    @Override
+    public <T extends ProgressIndex> Optional<T> getProgressIndexByType(
+        final Class<T> progressIndexClass) {
+      return progressIndexClass.isInstance(this)
+          ? Optional.of(progressIndexClass.cast(this))
+          : Optional.empty();
+    }
+
     @Override
     public TotalOrderSumTuple getTotalOrderSumTuple() {
       return new TotalOrderSumTuple((long) val);

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/consensus/index/ProgressIndex.java

@@ -34,6 +34,7 @@
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -158,6 +159,15 @@ public abstract ProgressIndex updateToMinimumEqualOrIsAfterProgressIndex(
    */
   public abstract ProgressIndexType getType();
 
+  /**
+   * Extracts a progress index of the given type from this progress index.
+   *
+   * <p>{@link StateProgressIndex} and {@link HybridProgressIndex} are recursively unwrapped because
+   * they may contain progress indexes from other causal chains.
+   */
+  public abstract <T extends ProgressIndex> Optional<T> getProgressIndexByType(
+      Class<T> progressIndexClass);
+
   /**
    * Get the sum of the tuples of each total order relation of the {@link ProgressIndex}, which is
    * used for topological sorting of the {@link ProgressIndex}.

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/consensus/index/impl/HybridProgressIndex.java

@@ -36,6 +36,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.stream.Collectors;
 
@@ -225,6 +226,30 @@ public ProgressIndexType getType() {
     return ProgressIndexType.HYBRID_PROGRESS_INDEX;
   }
 
+  @Override
+  public <T extends ProgressIndex> Optional<T> getProgressIndexByType(
+      final Class<T> progressIndexClass) {
+    if (progressIndexClass.isInstance(this)) {
+      return Optional.of(progressIndexClass.cast(this));
+    }
+
+    final Map<Short, ProgressIndex> type2Index = getType2Index();
+    // Prefer a direct component over one nested in another composite progress index.
+    for (final ProgressIndex progressIndex : type2Index.values()) {
+      if (progressIndexClass.isInstance(progressIndex)) {
+        return Optional.of(progressIndexClass.cast(progressIndex));
+      }
+    }
+    for (final ProgressIndex progressIndex : type2Index.values()) {
+      final Optional<T> extractedProgressIndex =
+          progressIndex.getProgressIndexByType(progressIndexClass);
+      if (extractedProgressIndex.isPresent()) {
+        return extractedProgressIndex;
+      }
+    }
+    return Optional.empty();
+  }
+
   @Override
   public TotalOrderSumTuple getTotalOrderSumTuple() {
     lock.readLock().lock();

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/consensus/index/impl/IoTProgressIndex.java

@@ -35,6 +35,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 public class IoTProgressIndex extends ProgressIndex {
@@ -197,6 +198,14 @@ public ProgressIndexType getType() {
     return ProgressIndexType.IOT_PROGRESS_INDEX;
   }
 
+  @Override
+  public <T extends ProgressIndex> Optional<T> getProgressIndexByType(
+      final Class<T> progressIndexClass) {
+    return progressIndexClass.isInstance(this)
+        ? Optional.of(progressIndexClass.cast(this))
+        : Optional.empty();
+  }
+
   @Override
   public TotalOrderSumTuple getTotalOrderSumTuple() {
     lock.readLock().lock();

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/consensus/index/impl/MetaProgressIndex.java

@@ -32,6 +32,7 @@
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 public class MetaProgressIndex extends ProgressIndex {
@@ -153,6 +154,14 @@ public ProgressIndexType getType() {
     return ProgressIndexType.META_PROGRESS_INDEX;
   }
 
+  @Override
+  public <T extends ProgressIndex> Optional<T> getProgressIndexByType(
+      final Class<T> progressIndexClass) {
+    return progressIndexClass.isInstance(this)
+        ? Optional.of(progressIndexClass.cast(this))
+        : Optional.empty();
+  }
+
   @Override
   public TotalOrderSumTuple getTotalOrderSumTuple() {
     lock.readLock().lock();

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/consensus/index/impl/MinimumProgressIndex.java

@@ -28,6 +28,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.Optional;
 
 public class MinimumProgressIndex extends ProgressIndex {
 
@@ -82,6 +83,14 @@ public ProgressIndexType getType() {
     return ProgressIndexType.MINIMUM_PROGRESS_INDEX;
   }
 
+  @Override
+  public <T extends ProgressIndex> Optional<T> getProgressIndexByType(
+      final Class<T> progressIndexClass) {
+    return progressIndexClass.isInstance(this)
+        ? Optional.of(progressIndexClass.cast(this))
+        : Optional.empty();
+  }
+
   @Override
   public TotalOrderSumTuple getTotalOrderSumTuple() {
     return TOTAL_ORDER_SUM_TUPLE;

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/consensus/index/impl/RecoverProgressIndex.java

@@ -36,6 +36,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.stream.Collectors;
 
@@ -202,6 +203,14 @@ public ProgressIndexType getType() {
     return ProgressIndexType.RECOVER_PROGRESS_INDEX;
   }
 
+  @Override
+  public <T extends ProgressIndex> Optional<T> getProgressIndexByType(
+      final Class<T> progressIndexClass) {
+    return progressIndexClass.isInstance(this)
+        ? Optional.of(progressIndexClass.cast(this))
+        : Optional.empty();
+  }
+
   @Override
   public TotalOrderSumTuple getTotalOrderSumTuple() {
     lock.readLock().lock();