luoluoyuyu commented on code in PR #16844: URL: https://github.com/apache/iotdb/pull/16844#discussion_r2601031982
########## iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/TabletStatementConverter.java: ########## @@ -0,0 +1,676 @@ +/* + * 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.sink.util; + +import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.exception.MetadataException; +import org.apache.iotdb.commons.path.PartialPath; +import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory; +import org.apache.iotdb.db.pipe.resource.memory.InsertNodeMemoryEstimator; +import org.apache.iotdb.db.queryengine.plan.analyze.cache.schema.DataNodeDevicePathCache; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertTabletStatement; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.utils.Binary; +import org.apache.tsfile.utils.BitMap; +import org.apache.tsfile.utils.BytesUtils; +import org.apache.tsfile.utils.DateUtils; +import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.utils.ReadWriteIOUtils; +import org.apache.tsfile.write.UnSupportedDataTypeException; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; + +import java.nio.ByteBuffer; +import java.time.LocalDate; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +/** + * Utility class for converting between InsertTabletStatement and Tablet format ByteBuffer. This + * avoids creating intermediate Tablet objects and directly converts between formats with only the + * fields needed. + */ +public class TabletStatementConverter { + + // Memory calculation constants - extracted from RamUsageEstimator for better performance + private static final long NUM_BYTES_ARRAY_HEADER = + org.apache.tsfile.utils.RamUsageEstimator.NUM_BYTES_ARRAY_HEADER; + private static final long NUM_BYTES_OBJECT_REF = + org.apache.tsfile.utils.RamUsageEstimator.NUM_BYTES_OBJECT_REF; + private static final long NUM_BYTES_OBJECT_HEADER = + org.apache.tsfile.utils.RamUsageEstimator.NUM_BYTES_OBJECT_HEADER; + private static final long SIZE_OF_ARRAYLIST = + org.apache.tsfile.utils.RamUsageEstimator.shallowSizeOfInstance(java.util.ArrayList.class); + private static final long SIZE_OF_BITMAP = + org.apache.tsfile.utils.RamUsageEstimator.shallowSizeOfInstance( + org.apache.tsfile.utils.BitMap.class); + + private TabletStatementConverter() { + // Utility class, no instantiation + } + + /** + * Convert Tablet to InsertTabletStatement. + * + * @param tablet Tablet object to convert + * @param isAligned whether the tablet is aligned + * @param databaseName database name (optional, for table model) + * @return InsertTabletStatement + * @throws MetadataException if conversion fails + */ + private static InsertTabletStatement convertTabletToStatement( + final Tablet tablet, final boolean isAligned, final String databaseName) + throws MetadataException { + try { + final String deviceId = tablet.getDeviceId(); + final int rowSize = tablet.getRowSize(); + final List<IMeasurementSchema> schemas = tablet.getSchemas(); + final int schemaSize = schemas.size(); + final long[] times = tablet.getTimestamps(); + final Object[] values = tablet.getValues(); + final BitMap[] bitMaps = tablet.getBitMaps(); + final List<ColumnCategory> columnCategories = tablet.getColumnTypes(); + + // Construct InsertTabletStatement + final InsertTabletStatement statement = new InsertTabletStatement(); + + // Set device path based on databaseName + // For table model, use getTableName() which returns the table name + // For tree model, use getDeviceId() which returns the device path + if (databaseName != null) { + // For table model, deviceId is actually the table name + statement.setDevicePath(new PartialPath(deviceId, false)); + statement.setDatabaseName(databaseName); + statement.setWriteToTable(true); + } else { + statement.setDevicePath(new PartialPath(deviceId)); + statement.setWriteToTable(false); + } + + // Set measurements, dataTypes, and measurementSchemas + final String[] measurements = new String[schemaSize]; + final TSDataType[] statementDataTypes = new TSDataType[schemaSize]; + final MeasurementSchema[] measurementSchemas = new MeasurementSchema[schemaSize]; + for (int i = 0; i < schemaSize; i++) { + final IMeasurementSchema schema = schemas.get(i); + if (schema != null) { + measurements[i] = schema.getMeasurementName(); + statementDataTypes[i] = schema.getType(); + measurementSchemas[i] = (MeasurementSchema) schema; + } else { + measurements[i] = null; + statementDataTypes[i] = null; + measurementSchemas[i] = null; + } + } + statement.setMeasurements(measurements); + statement.setDataTypes(statementDataTypes); + statement.setMeasurementSchemas(measurementSchemas); + + // Set column categories if databaseName is provided + if (databaseName != null && columnCategories != null) { + final TsTableColumnCategory[] statementColumnCategories = + new TsTableColumnCategory[schemaSize]; + for (int i = 0; i < schemaSize; i++) { + if (i < columnCategories.size() && columnCategories.get(i) != null) { + statementColumnCategories[i] = + TsTableColumnCategory.fromTsFileColumnCategory(columnCategories.get(i)); + } else { + statementColumnCategories[i] = null; + } + } + statement.setColumnCategories(statementColumnCategories); + } + + // Set times, rowCount, columns, bitMaps + // Tablet.getTimestamps() returns an array of size maxRowNumber, but we only need rowSize + // elements + final long[] statementTimes; + if (times != null && times.length >= rowSize && rowSize > 0) { + statementTimes = new long[rowSize]; + System.arraycopy(times, 0, statementTimes, 0, rowSize); + } else { + // If times array is null or too small, create empty array + statementTimes = new long[0]; + } + statement.setTimes(statementTimes); + statement.setRowCount(rowSize); + statement.setColumns(values); + statement.setBitMaps(bitMaps); + statement.setAligned(isAligned); + + return statement; + } catch (final Exception e) { + throw new MetadataException("Failed to convert Tablet to InsertTabletStatement", e); + } + } + + /** + * Convert InsertTabletStatement to Tablet. This method constructs a Tablet object from + * InsertTabletStatement, converting all necessary fields. + * + * @param statement InsertTabletStatement to convert + * @return Tablet object + * @throws MetadataException if conversion fails + */ + public static Tablet convertStatementToTablet(final InsertTabletStatement statement) + throws MetadataException { + try { + // Get deviceId/tableName from devicePath + final String deviceIdOrTableName = + statement.getDevicePath() != null ? statement.getDevicePath().getFullPath() : ""; + + // Get schemas from measurementSchemas + final MeasurementSchema[] measurementSchemas = statement.getMeasurementSchemas(); + final String[] measurements = statement.getMeasurements(); + final TSDataType[] dataTypes = statement.getDataTypes(); + final int schemaSize = measurements != null ? measurements.length : 0; + + final List<IMeasurementSchema> schemas = new ArrayList<>(); + for (int i = 0; i < schemaSize; i++) { + if (measurementSchemas != null && measurementSchemas[i] != null) { + schemas.add(measurementSchemas[i]); + } else if (measurements[i] != null && dataTypes[i] != null) { + // Create MeasurementSchema if not present + schemas.add(new MeasurementSchema(measurements[i], dataTypes[i])); + } else { + schemas.add(null); + } + } + + // Get columnTypes (for table model) + final TsTableColumnCategory[] columnCategories = statement.getColumnCategories(); + final List<ColumnCategory> tabletColumnTypes = new ArrayList<>(); + if (columnCategories != null && columnCategories.length > 0) { + for (int i = 0; i < schemaSize; i++) { + if (columnCategories[i] != null) { + tabletColumnTypes.add(columnCategories[i].toTsFileColumnType()); + } else { + tabletColumnTypes.add(ColumnCategory.FIELD); + } + } + } else { + // Default to FIELD for all columns if not specified + for (int i = 0; i < schemaSize; i++) { + tabletColumnTypes.add(ColumnCategory.FIELD); + } + } + + // Get timestamps + final long[] times = statement.getTimes(); + final int rowSize = statement.getRowCount(); + final long[] timestamps; + if (times != null && times.length >= rowSize) { + timestamps = new long[rowSize]; + System.arraycopy(times, 0, timestamps, 0, rowSize); + } else { + timestamps = new long[0]; + } + + // Get values - convert Statement columns to Tablet format + final Object[] statementColumns = statement.getColumns(); + final Object[] tabletValues = new Object[schemaSize]; + if (statementColumns != null && statementColumns.length > 0) { + for (int i = 0; i < schemaSize; i++) { + if (statementColumns[i] != null && dataTypes[i] != null) { + tabletValues[i] = convertStatementColumnToTablet(statementColumns[i], dataTypes[i]); + } else { + tabletValues[i] = null; + } + } + } + + // Get bitMaps + final BitMap[] bitMaps = statement.getBitMaps(); + + // Create Tablet using the full constructor + // Tablet(String tableName, List<IMeasurementSchema> schemas, List<ColumnCategory> + // columnTypes, + // long[] timestamps, Object[] values, BitMap[] bitMaps, int rowSize) + final Tablet tablet = + new Tablet( + deviceIdOrTableName, + schemas, + tabletColumnTypes, + timestamps, + tabletValues, + bitMaps, + rowSize); + + // Note: isAligned is not a property of Tablet, it's handled separately during insertion + return tablet; + } catch (final Exception e) { + throw new MetadataException("Failed to convert InsertTabletStatement to Tablet", e); + } + } + + /** + * Convert a single column value from Statement format to Tablet format. Statement uses primitive + * arrays (e.g., int[], long[], float[]), while Tablet may need different format. + * + * @param columnValue column value from Statement (primitive array) + * @param dataType data type of the column + * @return column value in Tablet format + */ + private static Object convertStatementColumnToTablet( + final Object columnValue, final TSDataType dataType) { + + if (TSDataType.DATE.equals(dataType)) { + final int[] values = (int[]) columnValue; + final LocalDate[] localDateValue = new LocalDate[values.length]; + for (int i = 0; i < values.length; i++) { + localDateValue[i] = DateUtils.parseIntToLocalDate(values[i]); + } + + return localDateValue; + } + // For primitive arrays (boolean[], int[], long[], float[], double[]), return as-is + return columnValue; + } + + /** + * Deserialize InsertTabletStatement from Tablet format ByteBuffer. + * + * @param byteBuffer ByteBuffer containing serialized data + * @param readDatabaseName whether to read databaseName from buffer (for V2 format) + * @return InsertTabletStatement with all fields set, including devicePath + */ + public static InsertTabletStatement deserializeStatementFromTabletFormat( + final ByteBuffer byteBuffer, final boolean readDatabaseName) throws IllegalPathException { Review Comment: This function is only useful for Pipes because the CPU usage for 'ClientRpcService' is primarily used for decoding Tablet columns. <img width="2938" height="1364" alt="image" src="https://github.com/user-attachments/assets/0e693cb2-7f46-42a0-8447-126807bb95ae"; /> -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
