[To dev/1.3] Pipe: Implemented the processTabletWithCollect api with tablet collector (#16941)

* Pipe: Implemented the processTabletWithCollect api with tablet collector (#16930)

* patch

* fix

* fix

* fix

Author: Caideyipi

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

@@ -21,6 +21,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.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;
@@ -36,17 +37,11 @@
 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
 
   public PipeRowCollector(PipeTaskMeta pipeTaskMeta, EnrichedEvent sourceEvent) {
-    this.pipeTaskMeta = pipeTaskMeta;
-    this.sourceEvent = sourceEvent;
+    super(pipeTaskMeta, sourceEvent);
   }
 
   @Override
@@ -113,14 +108,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

@@ -40,6 +40,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;
 
@@ -309,6 +310,17 @@ public Iterable<TabletInsertionEvent> processTablet(
         .collect(Collectors.toList());
   }
 
+  @Override
+  public Iterable<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer) {
+    return initDataContainers().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,50 @@
+/*
+ * 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.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
+
+  public PipeRawTabletEventConverter(PipeTaskMeta pipeTaskMeta, EnrichedEvent sourceEvent) {
+    this.pipeTaskMeta = pipeTaskMeta;
+    this.sourceEvent = sourceEvent;
+  }
+
+  @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

@@ -34,6 +34,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;
@@ -291,6 +292,16 @@ public Iterable<TabletInsertionEvent> processTablet(
     return dataContainer.processTablet(consumer);
   }
 
+  @Override
+  public Iterable<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer) {
+    if (dataContainer == null) {
+      dataContainer =
+          new TabletInsertionDataContainer(pipeTaskMeta, this, tablet, isAligned, pipePattern);
+    }
+    return dataContainer.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,46 @@
+/*
+ * 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.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);
+  }
+
+  @Override
+  public void collectTablet(final Tablet tablet) {
+    tabletInsertionEventList.add(
+        new PipeRawTabletInsertionEvent(
+            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/TabletInsertionDataContainer.java

@@ -30,6 +30,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.TSDataType;
@@ -689,6 +690,13 @@ public List<TabletInsertionEvent> processTablet(final BiConsumer<Tablet, RowColl
     return rowCollector.convertToTabletInsertionEvents(shouldReport);
   }
 
+  public List<TabletInsertionEvent> processTabletWithCollect(
+      BiConsumer<Tablet, TabletCollector> consumer) {
+    final PipeTabletCollector tabletCollector = new PipeTabletCollector(pipeTaskMeta, sourceEvent);
+    consumer.accept(convertToTablet(), tabletCollector);
+    return tabletCollector.convertToTabletInsertionEvents(shouldReport);
+  }
+
   ////////////////////////////  convertToTablet  ////////////////////////////
 
   public Tablet convertToTablet() {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java

@@ -22,7 +22,6 @@
 import org.apache.iotdb.commons.pipe.datastructure.pattern.IoTDBPipePattern;
 import org.apache.iotdb.commons.pipe.datastructure.pattern.PipePattern;
 import org.apache.iotdb.commons.pipe.datastructure.pattern.PrefixPipePattern;
-import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent;
 import org.apache.iotdb.db.pipe.event.common.tsfile.container.TsFileInsertionDataContainer;
 import org.apache.iotdb.db.pipe.event.common.tsfile.container.query.TsFileInsertionQueryDataContainer;
 import org.apache.iotdb.db.pipe.event.common.tsfile.container.scan.TsFileInsertionScanDataContainer;
@@ -593,18 +592,15 @@ private void testTsFilePointNum(
                                       })
                                   .forEach(
                                       tabletInsertionEvent2 ->
-                                          tabletInsertionEvent2.processTablet(
-                                              (tablet, rowCollector) ->
-                                                  new PipeRawTabletInsertionEvent(tablet, false)
-                                                      .processRowByRow(
-                                                          (row, collector) -> {
-                                                            try {
-                                                              rowCollector.collectRow(row);
-                                                              count3.addAndGet(getNonNullSize(row));
-                                                            } catch (final IOException e) {
-                                                              throw new RuntimeException(e);
-                                                            }
-                                                          })))));
+                                          tabletInsertionEvent2.processTabletWithCollect(
+                                              (tablet, collector) -> {
+                                                try {
+                                                  collector.collectTablet(tablet);
+                                                  count3.addAndGet(getNonNullSize(tablet));
+                                                } catch (final IOException e) {
+                                                  throw new RuntimeException(e);
+                                                }
+                                              }))));
 
       Assert.assertEquals(expectedCount, count1.get());
       Assert.assertEquals(expectedCount, count2.get());
@@ -624,4 +620,16 @@ private int getNonNullSize(final Row row) {
     }
     return count;
   }
+
+  private int getNonNullSize(final Tablet tablet) {
+    int count = 0;
+    for (int i = 0; i < tablet.rowSize; ++i) {
+      for (int j = 0; j < tablet.getSchemas().size(); ++j) {
+        if (tablet.bitMaps == null || tablet.bitMaps[j] == null || !tablet.bitMaps[j].isMarked(i)) {
+          ++count;
+        }
+      }
+    }
+    return count;
+  }
 }