[Pipe] Optimize memory usage (#17770)

Author: Caideyipi

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

@@ -152,7 +152,10 @@ public Type getDataType(final int columnIndex) {
 
   @Override
   public boolean isNull(final int columnIndex) {
-    return bitMaps[columnIndex].isMarked(rowIndex);
+    return bitMaps != null
+        && columnIndex < bitMaps.length
+        && bitMaps[columnIndex] != null
+        && bitMaps[columnIndex].isMarked(rowIndex);
   }
 
   @Override

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

@@ -23,6 +23,7 @@
 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.event.common.tablet.PipeTabletUtils;
 import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryWeightUtil;
 import org.apache.iotdb.pipe.api.access.Row;
 import org.apache.iotdb.pipe.api.collector.RowCollector;
@@ -66,27 +67,26 @@ public void collectRow(Row row) {
       Pair<Integer, Integer> rowCountAndMemorySize =
           PipeMemoryWeightUtil.calculateTabletRowCountAndMemory(pipeRow);
       tablet = new Tablet(deviceId, measurementSchemaList, rowCountAndMemorySize.getLeft());
-      tablet.initBitMaps();
       isAligned = pipeRow.isAligned();
     }
 
     final int rowIndex = tablet.rowSize;
-    tablet.addTimestamp(rowIndex, row.getTime());
+    PipeTabletUtils.putTimestamp(tablet, rowIndex, row.getTime());
     for (int i = 0; i < row.size(); i++) {
       final Object value = row.getObject(i);
-      if (value instanceof org.apache.iotdb.pipe.api.type.Binary) {
-        tablet.addValue(
-            measurementSchemaArray[i].getMeasurementId(),
-            rowIndex,
-            PipeBinaryTransformer.transformToBinary((org.apache.iotdb.pipe.api.type.Binary) value));
-      } else {
-        tablet.addValue(measurementSchemaArray[i].getMeasurementId(), rowIndex, value);
-      }
+      PipeTabletUtils.putValue(
+          tablet,
+          rowIndex,
+          i,
+          measurementSchemaArray[i].getType(),
+          value instanceof org.apache.iotdb.pipe.api.type.Binary
+              ? PipeBinaryTransformer.transformToBinary(
+                  (org.apache.iotdb.pipe.api.type.Binary) value)
+              : value);
       if (row.isNull(i)) {
-        tablet.bitMaps[i].mark(rowIndex);
+        PipeTabletUtils.markNullValue(tablet, rowIndex, i);
       }
     }
-    tablet.rowSize++;
 
     if (tablet.rowSize == tablet.getMaxRowNumber()) {
       collectTabletInsertionEvent();
@@ -95,6 +95,7 @@ public void collectRow(Row row) {
 
   private void collectTabletInsertionEvent() {
     if (tablet != null) {
+      PipeTabletUtils.compactBitMaps(tablet);
       tabletInsertionEventList.add(
           new PipeRawTabletInsertionEvent(
               tablet,

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

@@ -323,6 +323,7 @@ public boolean isAligned() {
 
   public Tablet convertToTablet() {
     if (!shouldParseTimeOrPattern()) {
+      PipeTabletUtils.compactBitMaps(tablet);
       return tablet;
     }
 

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

@@ -0,0 +1,245 @@
+/*
+ * 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.db.utils.BitMapUtils;
+
+import org.apache.tsfile.common.conf.TSFileConfig;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.utils.Binary;
+import org.apache.tsfile.utils.BitMap;
+import org.apache.tsfile.write.UnSupportedDataTypeException;
+import org.apache.tsfile.write.record.Tablet;
+import org.apache.tsfile.write.schema.MeasurementSchema;
+
+import java.time.LocalDate;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+public final class PipeTabletUtils {
+
+  private PipeTabletUtils() {}
+
+  public static final class TabletStringInternPool {
+
+    private final Map<String, String> internedStrings = new HashMap<>();
+
+    public String intern(final String value) {
+      if (Objects.isNull(value)) {
+        return null;
+      }
+
+      final String internedValue = internedStrings.get(value);
+      if (Objects.nonNull(internedValue)) {
+        return internedValue;
+      }
+
+      internedStrings.put(value, value);
+      return value;
+    }
+
+    public void intern(final String[] values) {
+      if (Objects.isNull(values)) {
+        return;
+      }
+
+      for (int i = 0; i < values.length; ++i) {
+        values[i] = intern(values[i]);
+      }
+    }
+
+    public void intern(final List<String> values) {
+      if (Objects.isNull(values)) {
+        return;
+      }
+
+      for (int i = 0; i < values.size(); ++i) {
+        values.set(i, intern(values.get(i)));
+      }
+    }
+
+    public Tablet intern(final Tablet tablet) {
+      if (Objects.isNull(tablet)) {
+        return null;
+      }
+
+      tablet.setDeviceId(intern(tablet.deviceId));
+      internMeasurementSchemas(tablet.getSchemas());
+      return tablet;
+    }
+
+    public void internMeasurementSchemas(final List<MeasurementSchema> schemas) {
+      if (Objects.isNull(schemas)) {
+        return;
+      }
+
+      for (final MeasurementSchema schema : schemas) {
+        intern(schema);
+      }
+    }
+
+    public MeasurementSchema intern(final MeasurementSchema schema) {
+      if (Objects.isNull(schema)) {
+        return null;
+      }
+
+      schema.setMeasurementId(intern(schema.getMeasurementId()));
+      schema.setProps(intern(schema.getProps()));
+      return schema;
+    }
+
+    private Map<String, String> intern(final Map<String, String> props) {
+      if (Objects.isNull(props) || props.isEmpty()) {
+        return props;
+      }
+
+      final Map<String, String> internedProps = new HashMap<>(props.size());
+      for (final Map.Entry<String, String> entry : props.entrySet()) {
+        internedProps.put(intern(entry.getKey()), intern(entry.getValue()));
+      }
+      return internedProps;
+    }
+  }
+
+  public static Tablet internTablet(
+      final Tablet tablet, final TabletStringInternPool tabletStringInternPool) {
+    return Objects.nonNull(tabletStringInternPool) ? tabletStringInternPool.intern(tablet) : tablet;
+  }
+
+  public static void compactBitMaps(final Tablet tablet) {
+    if (Objects.isNull(tablet)) {
+      return;
+    }
+    tablet.bitMaps = compactBitMaps(tablet.bitMaps, tablet.rowSize);
+  }
+
+  public static BitMap[] compactBitMaps(final BitMap[] bitMaps, final int rowCount) {
+    return BitMapUtils.compactBitMaps(bitMaps, rowCount);
+  }
+
+  public static BitMap[] copyBitMapsOrCreateEmpty(final Tablet tablet) {
+    final BitMap[] bitMaps = tablet.bitMaps;
+    return Objects.nonNull(bitMaps)
+        ? Arrays.copyOf(bitMaps, bitMaps.length)
+        : new BitMap[getColumnCount(tablet)];
+  }
+
+  public static void markNullValue(final Tablet tablet, final int rowIndex, final int columnIndex) {
+    final BitMap[] bitMaps = ensureBitMaps(tablet, columnIndex + 1);
+    if (Objects.isNull(bitMaps[columnIndex])) {
+      bitMaps[columnIndex] = new BitMap(tablet.getMaxRowNumber());
+    }
+    bitMaps[columnIndex].mark(rowIndex);
+  }
+
+  public static void putTimestamp(final Tablet tablet, final int rowIndex, final long timestamp) {
+    tablet.timestamps[rowIndex] = timestamp;
+    tablet.rowSize = Math.max(tablet.rowSize, rowIndex + 1);
+  }
+
+  public static void putValue(
+      final Tablet tablet,
+      final int rowIndex,
+      final int columnIndex,
+      final TSDataType dataType,
+      final Object value) {
+    switch (dataType) {
+      case BOOLEAN:
+        ((boolean[]) tablet.values[columnIndex])[rowIndex] = (Boolean) value;
+        break;
+      case INT32:
+        ((int[]) tablet.values[columnIndex])[rowIndex] = (Integer) value;
+        break;
+      case DATE:
+        ((LocalDate[]) tablet.values[columnIndex])[rowIndex] = (LocalDate) value;
+        break;
+      case INT64:
+      case TIMESTAMP:
+        ((long[]) tablet.values[columnIndex])[rowIndex] = (Long) value;
+        break;
+      case FLOAT:
+        ((float[]) tablet.values[columnIndex])[rowIndex] = (Float) value;
+        break;
+      case DOUBLE:
+        ((double[]) tablet.values[columnIndex])[rowIndex] = (Double) value;
+        break;
+      case TEXT:
+      case BLOB:
+      case STRING:
+        ((Binary[]) tablet.values[columnIndex])[rowIndex] = toBinary(value);
+        break;
+      default:
+        throw new UnSupportedDataTypeException("Unsupported data type: " + dataType);
+    }
+    unmarkNullValue(tablet, rowIndex, columnIndex);
+  }
+
+  private static void unmarkNullValue(
+      final Tablet tablet, final int rowIndex, final int columnIndex) {
+    final BitMap[] bitMaps = tablet.bitMaps;
+    if (Objects.nonNull(bitMaps)
+        && columnIndex < bitMaps.length
+        && Objects.nonNull(bitMaps[columnIndex])) {
+      bitMaps[columnIndex].unmark(rowIndex);
+    }
+  }
+
+  private static BitMap[] ensureBitMaps(final Tablet tablet, final int minColumnCount) {
+    final int columnCount = Math.max(getColumnCount(tablet), minColumnCount);
+    BitMap[] bitMaps = tablet.bitMaps;
+    if (Objects.isNull(bitMaps)) {
+      bitMaps = new BitMap[columnCount];
+      tablet.bitMaps = bitMaps;
+    } else if (bitMaps.length < columnCount) {
+      final BitMap[] expandedBitMaps = new BitMap[columnCount];
+      System.arraycopy(bitMaps, 0, expandedBitMaps, 0, bitMaps.length);
+      bitMaps = expandedBitMaps;
+      tablet.bitMaps = bitMaps;
+    }
+    return bitMaps;
+  }
+
+  private static int getColumnCount(final Tablet tablet) {
+    if (Objects.nonNull(tablet.getSchemas())) {
+      return tablet.getSchemas().size();
+    }
+    return Objects.nonNull(tablet.values) ? tablet.values.length : 0;
+  }
+
+  private static Binary toBinary(final Object value) {
+    if (Objects.isNull(value)) {
+      return Binary.EMPTY_VALUE;
+    }
+    if (value instanceof Binary) {
+      return (Binary) value;
+    }
+    if (value instanceof byte[]) {
+      return new Binary((byte[]) value);
+    }
+    if (value instanceof String) {
+      return new Binary(((String) value).getBytes(TSFileConfig.STRING_CHARSET));
+    }
+    throw new IllegalArgumentException(
+        String.format("Expected Binary, byte[] or String, but was %s.", value.getClass()));
+  }
+}