may-complete-n

Author: Caideyipi

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

@@ -27,7 +27,7 @@
 
 /**
  * Used to collect rows generated by {@link TabletInsertionEvent#processRowByRow(BiConsumer)},
- * {@link TabletInsertionEvent#processTablet(BiConsumer)}.
+ * {@link TabletInsertionEvent#processTabletWithTabletCollection(BiConsumer)}.
  */
 public interface RowCollector {
 

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 RowCollector}
+   */
+  Iterable<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer);
 }

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/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

@@ -26,8 +26,6 @@
 
 import org.apache.tsfile.write.record.Tablet;
 
-import java.io.IOException;
-
 public class PipeTabletCollector extends PipeRawTabletEventConverter implements TabletCollector {
 
   public PipeTabletCollector(PipeTaskMeta pipeTaskMeta, EnrichedEvent sourceEvent) {

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

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

@@ -27,6 +27,7 @@
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode;
 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.write.UnSupportedDataTypeException;
@@ -113,6 +114,15 @@ public List<TabletInsertionEvent> processTablet(BiConsumer<Tablet, RowCollector>
     return Collections.emptyList();
   }
 
+  @Override
+  public List<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer) {
+    if (LOGGER.isDebugEnabled()) {
+      LOGGER.warn("TablePatternParser does not support tablet processing");
+    }
+    return Collections.emptyList();
+  }
+
   ////////////////////////////  convertToTablet  ////////////////////////////
 
   @Override

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

@@ -25,11 +25,13 @@
 import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.db.pipe.event.common.row.PipeRow;
 import org.apache.iotdb.db.pipe.event.common.row.PipeRowCollector;
+import org.apache.iotdb.db.pipe.event.common.tablet.PipeTabletCollector;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode;
 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.write.UnSupportedDataTypeException;
@@ -154,6 +156,14 @@ public List<TabletInsertionEvent> processTablet(final BiConsumer<Tablet, RowColl
     return rowCollector.convertToTabletInsertionEvents(shouldReport);
   }
 
+  @Override
+  public List<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer) {
+    final PipeTabletCollector tabletCollector = new PipeTabletCollector(pipeTaskMeta, sourceEvent);
+    consumer.accept(convertToTablet(), tabletCollector);
+    return tabletCollector.convertToTabletInsertionEvents(shouldReport);
+  }
+
   ////////////////////////////  convertToTablet  ////////////////////////////
 
   @Override