Pipe: Implemented the processTabletWithCollect api with tablet collector

iotdb-api/pipe-api/src/main/java/org/apache/iotdb/pipe/api/collector/DataCollector.java

@@ -0,0 +1,37 @@
+/*
+ * 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.pipe.api.collector;
+
+import org.apache.iotdb.pipe.api.access.Row;
+import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
+
+import java.util.List;
+
+/** Transform data to {@link TabletInsertionEvent}. */
+public interface DataCollector {
+
+  /**
+   * Transform data to {@link TabletInsertionEvent}.
+   *
+   * @param shouldReport Whether to report progress for generated events
+   * @see Row
+   */
+  List<TabletInsertionEvent> convertToTabletInsertionEvents(final boolean shouldReport);
+}

iotdb-api/pipe-api/src/main/java/org/apache/iotdb/pipe/api/collector/TabletCollector.java

@@ -0,0 +1,44 @@
+/*
+ * 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.pipe.api.collector;
+
+import org.apache.iotdb.pipe.api.access.Row;
+import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
+
+import org.apache.tsfile.write.record.Tablet;
+
+import java.io.IOException;
+import java.util.function.BiConsumer;
+
+/**
+ * Used to collect rows generated by {@link
+ * TabletInsertionEvent#processTabletWithCollect(BiConsumer)}.
+ */
+public interface TabletCollector {
+
+  /**
+   * Collects a tablet.
+   *
+   * @param tablet Tablet to be collected
+   * @throws IOException if any I/O errors occur
+   * @see Row
+   */
+  void collectTablet(Tablet tablet) throws IOException;
+}

iotdb-api/pipe-api/src/main/java/org/apache/iotdb/pipe/api/event/dml/insertion/TabletInsertionEvent.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.pipe.api.access.Row;
 import org.apache.iotdb.pipe.api.collector.RowCollector;
+import org.apache.iotdb.pipe.api.collector.TabletCollector;
 import org.apache.iotdb.pipe.api.event.Event;
 
 import org.apache.tsfile.write.record.Tablet;
@@ -45,4 +46,14 @@ public interface TabletInsertionEvent extends Event {
    *     contains the results collected by the {@link RowCollector}
    */
   Iterable<TabletInsertionEvent> processTablet(BiConsumer<Tablet, RowCollector> consumer);
+
+  /**
+   * The consumer processes the Tablet directly and collects the results by {@link
+   * org.apache.iotdb.pipe.api.collector.TabletCollector}.
+   *
+   * @return {@code Iterable<TabletInsertionEvent>} a list of new {@link TabletInsertionEvent}
+   *     contains the results collected by the {@link TabletCollector}
+   */
+  Iterable<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer);
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/row/PipeRowCollector.java

@@ -22,6 +22,7 @@
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
 import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
 import org.apache.iotdb.db.pipe.event.common.PipeInsertionEvent;
+import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletEventConverter;
 import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent;
 import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryWeightUtil;
 import org.apache.iotdb.pipe.api.access.Row;
@@ -38,38 +39,19 @@
 import java.util.Arrays;
 import java.util.List;
 
-public class PipeRowCollector implements RowCollector {
-
-  private final List<TabletInsertionEvent> tabletInsertionEventList = new ArrayList<>();
+public class PipeRowCollector extends PipeRawTabletEventConverter implements RowCollector {
   private Tablet tablet = null;
-  private boolean isAligned = false;
-  private final PipeTaskMeta pipeTaskMeta; // Used to report progress
-  private final EnrichedEvent sourceEvent; // Used to report progress
-  private final String sourceEventDataBaseName;
-  private final Boolean isTableModel;
 
   public PipeRowCollector(PipeTaskMeta pipeTaskMeta, EnrichedEvent sourceEvent) {
-    this.pipeTaskMeta = pipeTaskMeta;
-    this.sourceEvent = sourceEvent;
-    if (sourceEvent instanceof PipeInsertionEvent) {
-      sourceEventDataBaseName =
-          ((PipeInsertionEvent) sourceEvent).getSourceDatabaseNameFromDataRegion();
-      isTableModel = ((PipeInsertionEvent) sourceEvent).getRawIsTableModelEvent();
-    } else {
-      sourceEventDataBaseName = null;
-      isTableModel = null;
-    }
+    super(pipeTaskMeta, sourceEvent);
   }
 
   public PipeRowCollector(
       PipeTaskMeta pipeTaskMeta,
       EnrichedEvent sourceEvent,
       String sourceEventDataBase,
       Boolean isTableModel) {
-    this.pipeTaskMeta = pipeTaskMeta;
-    this.sourceEvent = sourceEvent;
-    this.sourceEventDataBaseName = sourceEventDataBase;
-    this.isTableModel = isTableModel;
+    super(pipeTaskMeta, sourceEvent, sourceEventDataBase, isTableModel);
   }
 
   @Override
@@ -142,14 +124,9 @@ private void collectTabletInsertionEvent() {
     this.tablet = null;
   }
 
+  @Override
   public List<TabletInsertionEvent> convertToTabletInsertionEvents(final boolean shouldReport) {
     collectTabletInsertionEvent();
-
-    final int eventListSize = tabletInsertionEventList.size();
-    if (eventListSize > 0 && shouldReport) { // The last event should report progress
-      ((PipeRawTabletInsertionEvent) tabletInsertionEventList.get(eventListSize - 1))
-          .markAsNeedToReport();
-    }
-    return tabletInsertionEventList;
+    return super.convertToTabletInsertionEvents(shouldReport);
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tablet/PipeInsertNodeTabletInsertionEvent.java

@@ -52,6 +52,7 @@
 import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALPipeException;
 import org.apache.iotdb.pipe.api.access.Row;
 import org.apache.iotdb.pipe.api.collector.RowCollector;
+import org.apache.iotdb.pipe.api.collector.TabletCollector;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
 import org.apache.iotdb.pipe.api.exception.PipeException;
 
@@ -425,6 +426,17 @@ public Iterable<TabletInsertionEvent> processTablet(
         .collect(Collectors.toList());
   }
 
+  @Override
+  public Iterable<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer) {
+    return initEventParsers().stream()
+        .map(
+            tabletInsertionEventParser ->
+                tabletInsertionEventParser.processTabletWithCollect(consumer))
+        .flatMap(Collection::stream)
+        .collect(Collectors.toList());
+  }
+
   /////////////////////////// convertToTablet ///////////////////////////
 
   public boolean isAligned(final int i) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tablet/PipeRawTabletEventConverter.java

@@ -0,0 +1,72 @@
+/*
+ * 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.pipe.event.common.tablet;
+
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
+import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
+import org.apache.iotdb.db.pipe.event.common.PipeInsertionEvent;
+import org.apache.iotdb.pipe.api.collector.DataCollector;
+import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public abstract class PipeRawTabletEventConverter implements DataCollector {
+  protected final List<TabletInsertionEvent> tabletInsertionEventList = new ArrayList<>();
+  protected boolean isAligned = false;
+  protected final PipeTaskMeta pipeTaskMeta; // Used to report progress
+  protected final EnrichedEvent sourceEvent; // Used to report progress
+  protected final String sourceEventDataBaseName;
+  protected final Boolean isTableModel;
+
+  public PipeRawTabletEventConverter(PipeTaskMeta pipeTaskMeta, EnrichedEvent sourceEvent) {
+    this.pipeTaskMeta = pipeTaskMeta;
+    this.sourceEvent = sourceEvent;
+    if (sourceEvent instanceof PipeInsertionEvent) {
+      sourceEventDataBaseName =
+          ((PipeInsertionEvent) sourceEvent).getSourceDatabaseNameFromDataRegion();
+      isTableModel = ((PipeInsertionEvent) sourceEvent).getRawIsTableModelEvent();
+    } else {
+      sourceEventDataBaseName = null;
+      isTableModel = null;
+    }
+  }
+
+  public PipeRawTabletEventConverter(
+      PipeTaskMeta pipeTaskMeta,
+      EnrichedEvent sourceEvent,
+      String sourceEventDataBase,
+      Boolean isTableModel) {
+    this.pipeTaskMeta = pipeTaskMeta;
+    this.sourceEvent = sourceEvent;
+    this.sourceEventDataBaseName = sourceEventDataBase;
+    this.isTableModel = isTableModel;
+  }
+
+  @Override
+  public List<TabletInsertionEvent> convertToTabletInsertionEvents(final boolean shouldReport) {
+    final int eventListSize = tabletInsertionEventList.size();
+    if (eventListSize > 0 && shouldReport) { // The last event should report progress
+      ((PipeRawTabletInsertionEvent) tabletInsertionEventList.get(eventListSize - 1))
+          .markAsNeedToReport();
+    }
+    return tabletInsertionEventList;
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tablet/PipeRawTabletInsertionEvent.java

@@ -39,6 +39,7 @@
 import org.apache.iotdb.db.pipe.resource.memory.PipeTabletMemoryBlock;
 import org.apache.iotdb.pipe.api.access.Row;
 import org.apache.iotdb.pipe.api.collector.RowCollector;
+import org.apache.iotdb.pipe.api.collector.TabletCollector;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
 
 import org.apache.tsfile.utils.RamUsageEstimator;
@@ -415,6 +416,12 @@ public Iterable<TabletInsertionEvent> processTablet(
     return initEventParser().processTablet(consumer);
   }
 
+  @Override
+  public Iterable<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer) {
+    return initEventParser().processTabletWithCollect(consumer);
+  }
+
   /////////////////////////// convertToTablet ///////////////////////////
 
   public boolean isAligned() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tablet/PipeTabletCollector.java

@@ -0,0 +1,61 @@
+/*
+ * 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.pipe.event.common.tablet;
+
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
+import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
+import org.apache.iotdb.db.pipe.event.common.PipeInsertionEvent;
+import org.apache.iotdb.pipe.api.collector.TabletCollector;
+
+import org.apache.tsfile.write.record.Tablet;
+
+public class PipeTabletCollector extends PipeRawTabletEventConverter implements TabletCollector {
+
+  public PipeTabletCollector(PipeTaskMeta pipeTaskMeta, EnrichedEvent sourceEvent) {
+    super(pipeTaskMeta, sourceEvent);
+  }
+
+  public PipeTabletCollector(
+      PipeTaskMeta pipeTaskMeta,
+      EnrichedEvent sourceEvent,
+      String sourceEventDataBase,
+      Boolean isTableModel) {
+    super(pipeTaskMeta, sourceEvent, sourceEventDataBase, isTableModel);
+  }
+
+  @Override
+  public void collectTablet(final Tablet tablet) {
+    final PipeInsertionEvent pipeInsertionEvent =
+        sourceEvent instanceof PipeInsertionEvent ? ((PipeInsertionEvent) sourceEvent) : null;
+    tabletInsertionEventList.add(
+        new PipeRawTabletInsertionEvent(
+            isTableModel,
+            sourceEventDataBaseName,
+            pipeInsertionEvent == null ? null : pipeInsertionEvent.getRawTableModelDataBase(),
+            pipeInsertionEvent == null ? null : pipeInsertionEvent.getRawTreeModelDataBase(),
+            tablet,
+            isAligned,
+            sourceEvent == null ? null : sourceEvent.getPipeName(),
+            sourceEvent == null ? 0 : sourceEvent.getCreationTime(),
+            pipeTaskMeta,
+            sourceEvent,
+            false));
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tablet/parser/TabletInsertionEventParser.java

@@ -26,6 +26,7 @@
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertTabletNode;
 import org.apache.iotdb.pipe.api.access.Row;
 import org.apache.iotdb.pipe.api.collector.RowCollector;
+import org.apache.iotdb.pipe.api.collector.TabletCollector;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
 
 import org.apache.tsfile.enums.ColumnCategory;
@@ -636,5 +637,8 @@ public abstract List<TabletInsertionEvent> processRowByRow(
   public abstract List<TabletInsertionEvent> processTablet(
       final BiConsumer<Tablet, RowCollector> consumer);
 
+  public abstract List<TabletInsertionEvent> processTabletWithCollect(
+      final BiConsumer<Tablet, TabletCollector> consumer);
+
   public abstract Tablet convertToTablet();
 }