Pipe: Handle hybrid meta progress indexes

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

@@ -52,6 +52,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;
@@ -182,10 +183,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/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java

@@ -434,6 +434,8 @@ public final class DataNodePipeMessages {
       "Failed to send request {} (watermark = {}) to {}";
   public static final String FAILED_TO_TRIGGER_COMBINE_WATERMARK_COUNT_PROGRESSINDEX =
       "Failed to trigger combine. watermark={}, count={}, progressIndex={}";
+  public static final String EXCEPTION_FAILED_TO_INITIALIZE_STATEPROGRESSINDEX_FROM_PROGRESS_INDEX_ARG_E95617F9 =
+      "Failed to initialize StateProgressIndex from progress index %s.";
   public static final String FAILURE_OCCURRED_WHEN_TRYING_TO_COMMIT_PROGRESS =
       "Failure occurred when trying to commit progress index. timestamp={}, count={}, "
           + "progressIndex={}";

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java

@@ -412,6 +412,8 @@ public final class DataNodePipeMessages {
       "发送 request {}（watermark = {}）到 {} 失败";
   public static final String FAILED_TO_TRIGGER_COMBINE_WATERMARK_COUNT_PROGRESSINDEX =
       "触发合并失败。watermark={}, count={}, progressIndex={}";
+  public static final String EXCEPTION_FAILED_TO_INITIALIZE_STATEPROGRESSINDEX_FROM_PROGRESS_INDEX_ARG_E95617F9 =
+      "无法从进度索引 %s 初始化 StateProgressIndex。";
   public static final String FAILURE_OCCURRED_WHEN_TRYING_TO_COMMIT_PROGRESS =
       "尝试提交进度索引时发生失败。timestamp={}, count={}, "
           + "progressIndex={}";

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

@@ -351,19 +351,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(
-              DataNodePipeMessages
-                  .PIPE_EXCEPTION_THE_AGGREGATE_PROCESSOR_DOES_NOT_SUPPORT_PROGRESSINDEXTYPE_35351D27,
-              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

@@ -156,13 +156,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(
@@ -191,6 +185,27 @@ static PartialPath parseOutputSeries(final PipeParameters parameters)
         parameters.getStringByKeys(PROCESSOR_OUTPUT_SERIES_KEY, _PROCESSOR_OUTPUT_SERIES_KEY));
   }
 
+  static StateProgressIndex initializeStateProgressIndex(final PipeTaskMeta pipeTaskMeta) {
+    final ProgressIndex progressIndex = pipeTaskMeta.getProgressIndex();
+    if (progressIndex instanceof StateProgressIndex stateProgressIndex) {
+      return stateProgressIndex;
+    }
+
+    final ProgressIndex updatedProgressIndex =
+        pipeTaskMeta.updateProgressIndex(
+            new StateProgressIndex(
+                Long.MIN_VALUE, Collections.emptyMap(), MinimumProgressIndex.INSTANCE));
+    return updatedProgressIndex
+        .getProgressIndexByType(StateProgressIndex.class)
+        .orElseThrow(
+            () ->
+                new PipeException(
+                    String.format(
+                        DataNodePipeMessages
+                            .EXCEPTION_FAILED_TO_INITIALIZE_STATEPROGRESSINDEX_FROM_PROGRESS_INDEX_ARG_E95617F9,
+                        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/PipeHistoricalDataRegionTsFileAndDeletionSource.java

@@ -696,7 +696,7 @@ private void sortExtractedResources(final List<PersistentResource> resourceList)
 
     resourceList.sort(
         (o1, o2) ->
-            startIndex instanceof TimeWindowStateProgressIndex
+            Objects.nonNull(getTimeWindowStateProgressIndex(startIndex))
                 ? Long.compare(o1.getFileStartTime(), o2.getFileStartTime())
                 : comparePersistentResourcesByProgressIndex(o1, o2));
   }
@@ -898,11 +898,11 @@ private boolean shouldExtractTsFileResource(
   }
 
   private boolean mayTsFileContainUnprocessedData(final TsFileResource resource) {
-    final ProgressIndex innerStartIndex = getInnerProgressIndex(startIndex);
-    if (innerStartIndex 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) innerStartIndex).getMinTime()
-          <= resource.getFileEndTime();
+      return timeWindowStateProgressIndex.getMinTime() <= resource.getFileEndTime();
     }
 
     if (pipeName.startsWith(PipeStaticMeta.CONSENSUS_PIPE_PREFIX)) {
@@ -941,6 +941,13 @@ private ProgressIndex getInnerProgressIndex(final ProgressIndex progressIndex) {
         : Objects.isNull(progressIndex) ? MinimumProgressIndex.INSTANCE : progressIndex;
   }
 
+  private TimeWindowStateProgressIndex getTimeWindowStateProgressIndex(
+      final ProgressIndex progressIndex) {
+    return Objects.isNull(progressIndex)
+        ? null
+        : progressIndex.getProgressIndexByType(TimeWindowStateProgressIndex.class).orElse(null);
+  }
+
   private boolean isProgressIndexCoveredByTimePartitionProgressIndex(
       final PersistentResource resource,
       final ProgressIndex progressIndex,

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.PipeParameters;
 
 import org.junit.Assert;
@@ -37,6 +44,54 @@ public void testOutputSeriesSupportsNewAndLegacyKeys() throws Exception {
         "root.db.d.s2", parseOutputSeries("processor.output-series", "root.db.d.s2").getFullPath());
   }
 
+  @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;
 
@@ -217,6 +218,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/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java

@@ -957,6 +957,8 @@ private PipeMessages() {}
   public static final String LOG_ORIGIN_REQUEST_TYPE_MISMATCH_EXPECTED_ARG_ACTUAL_ARG_D96D10AE = "Origin request type mismatch: expected {}, actual {}";
   public static final String LOG_ORIGIN_BODY_SIZE_MISMATCH_EXPECTED_ARG_ACTUAL_ARG_5D410B75 = "Origin body size mismatch: expected {}, actual {}";
   public static final String LOG_INVALID_SLICE_INDEX_EXPECTED_ARG_ACTUAL_ARG_2AC41628 = "Invalid slice index: expected {}, actual {}";
+  public static final String LOG_PIPE_ARG_ARG_ENCOUNTERED_AN_UNEXPECTED_HYBRIDPROGRESSINDEX_IN_ARG_PROGRESS_INDEX_ARG_7C578B17 =
+      "Pipe {}@{} encountered an unexpected HybridProgressIndex in {}. Progress index: {}.";
   public static final String EXCEPTION_DECOMPRESSED_LENGTH_SHOULD_BETWEEN_0_ARG_BUT_GOT_ARG_488B3073 = "Decompressed length should be between 0 and %d, but got %d.";
   public static final String EXCEPTION_COMMA_50AD1C01 = ", ";
   public static final String MESSAGE_NO_DATAPARTITIONTABLE_GENERATION_TASK_FOUND_4414BE55 = "No DataPartitionTable generation task found";

iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java

@@ -923,6 +923,8 @@ private PipeMessages() {}
   public static final String LOG_ORIGIN_REQUEST_TYPE_MISMATCH_EXPECTED_ARG_ACTUAL_ARG_D96D10AE = "Origin request type 不匹配：期望 {}，实际 {}";
   public static final String LOG_ORIGIN_BODY_SIZE_MISMATCH_EXPECTED_ARG_ACTUAL_ARG_5D410B75 = "Origin body size 不匹配：期望 {}，实际 {}";
   public static final String LOG_INVALID_SLICE_INDEX_EXPECTED_ARG_ACTUAL_ARG_2AC41628 = "无效的 slice index：期望 {}，实际 {}";
+  public static final String LOG_PIPE_ARG_ARG_ENCOUNTERED_AN_UNEXPECTED_HYBRIDPROGRESSINDEX_IN_ARG_PROGRESS_INDEX_ARG_7C578B17 =
+      "Pipe {}@{} 在 {} 中遇到了非预期的 HybridProgressIndex。进度索引：{}。";
   public static final String EXCEPTION_DECOMPRESSED_LENGTH_SHOULD_BETWEEN_0_ARG_BUT_GOT_ARG_488B3073 = "解压后长度应介于 0 和 %d 之间，但实际为 %d。";
   public static final String EXCEPTION_COMMA_50AD1C01 = ", ";
   public static final String MESSAGE_NO_DATAPARTITIONTABLE_GENERATION_TASK_FOUND_4414BE55 = "未找到 DataPartitionTable 生成任务";

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();