This is an automated email from the ASF dual-hosted git repository. jt2594838 pushed a commit to branch insert_numeric_row in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit f0bed1c20aeaf6fdc84087a24a044ccf3e7a5cbc Author: Tian Jiang <[email protected]> AuthorDate: Wed Sep 16 18:13:04 2026 +0800 tmp --- .../isession/NumericRowPreparedStatement.java | 132 +++++++++++ .../session/PreparedNumericRowBufferTest.java | 141 ++++++++++++ .../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 22 ++ .../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 22 ++ .../protocol/session/PreparedNumericRowInfo.java | 242 +++++++++++++++++++++ 5 files changed, 559 insertions(+) diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/NumericRowPreparedStatement.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/NumericRowPreparedStatement.java new file mode 100644 index 00000000000..8852d05a8a8 --- /dev/null +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/NumericRowPreparedStatement.java @@ -0,0 +1,132 @@ +/* + * 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. + */ +package org.apache.iotdb.isession; + +import org.apache.iotdb.isession.i18n.ISessionMessages; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; + +import java.nio.ByteBuffer; +import java.util.List; + +/** + * Immutable metadata and reusable serialization buffers for repeated numeric table row inserts. + * + * <p>This statement is not thread-safe. An execution must return before the statement is reused or + * closed, because its buffers are retained until the synchronous RPC completes. + */ +public final class NumericRowPreparedStatement implements AutoCloseable { + private final String tableName; + private final List<String> columnNames; + private final List<TSDataType> dataTypes; + private final List<ColumnCategory> columnCategories; + private final int tagCount; + private final int valueBytes; + private ByteBuffer tagBuffer; + private ByteBuffer valueBuffer; + private long preparedId = -1; + private Object owner; + private Runnable closeAction = () -> {}; + private boolean closed; + + public NumericRowPreparedStatement( + String tableName, + List<String> columnNames, + List<TSDataType> dataTypes, + List<ColumnCategory> columnCategories) { + this.tableName = tableName; + this.columnNames = List.copyOf(columnNames); + this.dataTypes = List.copyOf(dataTypes); + this.columnCategories = List.copyOf(columnCategories); + // Column categories are immutable, so repeated inserts can reuse this count. + this.tagCount = + (int) this.columnCategories.stream().filter(c -> c == ColumnCategory.TAG).count(); + // FIELD encodings have fixed widths for the lifetime of the prepared schema. + int bytes = 0; + for (int i = 0; i < this.dataTypes.size(); i++) { + if (this.columnCategories.get(i) == ColumnCategory.FIELD) { + bytes += + switch (this.dataTypes.get(i)) { + case BOOLEAN -> 1; + case INT32, DATE, FLOAT -> Integer.BYTES; + case INT64, TIMESTAMP, DOUBLE -> Long.BYTES; + default -> + throw new IllegalArgumentException( + ISessionMessages.EXCEPTION_UNSUPPORTED_NUMERIC_TYPE_06DD3044); + }; + } + } + this.valueBytes = bytes; + } + + public String getTableName() { + return tableName; + } + + public List<String> getColumnNames() { + return columnNames; + } + + public List<TSDataType> getDataTypes() { + return dataTypes; + } + + public List<ColumnCategory> getColumnCategories() { + return columnCategories; + } + + public int getTagCount() { + return tagCount; + } + + /** + * Returns the TAG buffer cleared for writing, growing it when the encoded tags need more space. + */ + public ByteBuffer prepareTagBuffer(int requiredCapacity) { + if (tagBuffer == null || tagBuffer.capacity() < requiredCapacity) { + tagBuffer = ByteBuffer.allocate(requiredCapacity); + } + tagBuffer.clear(); + return tagBuffer; + } + + /** Returns the FIELD buffer cleared for writing; a prepared schema has a fixed encoded size. */ + public ByteBuffer prepareValueBuffer() { + if (valueBuffer == null) { + valueBuffer = ByteBuffer.allocate(valueBytes); + } + valueBuffer.clear(); + return valueBuffer; + } + + public long getPreparedId() { + return preparedId; + } + + public Object getOwner() { + return owner; + } + + public void bind(long id, Object owner, Runnable closeAction) { + this.preparedId = id; + this.owner = owner; + this.closeAction = closeAction; + this.closed = false; + } + + public boolean isClosed() { + return closed; + } + + @Override + public void close() { + closeAction.run(); + closeAction = () -> {}; + closed = true; + } +} diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowBufferTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowBufferTest.java new file mode 100644 index 00000000000..a409d6844f0 --- /dev/null +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowBufferTest.java @@ -0,0 +1,141 @@ +/* + * 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.session; + +import org.apache.iotdb.isession.NumericRowPreparedStatement; +import org.apache.iotdb.rpc.IoTDBConnectionException; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.utils.ReadWriteIOUtils; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertSame; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class PreparedNumericRowBufferTest { + @Test + public void testBuffersAcrossTagGrowthShrinkAndRpcFailure() throws Exception { + Session session = mock(Session.class); + SessionConnection connection = mock(SessionConnection.class); + when(session.getDefaultSessionConnection()).thenReturn(connection); + TableSession tableSession = new TableSession(session); + NumericRowPreparedStatement prepared = + new NumericRowPreparedStatement( + "t", + List.of("tag", "flag", "count", "value"), + List.of(TSDataType.STRING, TSDataType.BOOLEAN, TSDataType.INT32, TSDataType.DOUBLE), + List.of( + ColumnCategory.TAG, + ColumnCategory.FIELD, + ColumnCategory.FIELD, + ColumnCategory.FIELD)); + prepared.bind(1, connection, () -> {}); + List<ByteBuffer> tagBuffers = new ArrayList<>(); + List<ByteBuffer> valueBuffers = new ArrayList<>(); + String[] tags = {"x", "long-tag-\u4e2d\u6587", "", null, "retry"}; + doAnswer( + invocation -> { + int index = tagBuffers.size(); + ByteBuffer values = invocation.getArgument(2); + ByteBuffer encodedTags = invocation.getArgument(3); + tagBuffers.add(encodedTags); + valueBuffers.add(values); + // Consume both buffers as a transport may do, and verify no previous row bytes + // remain. + assertEquals(tags[index], ReadWriteIOUtils.readString(encodedTags)); + assertFalse(encodedTags.hasRemaining()); + assertEquals(13, values.remaining()); + assertEquals((byte) (index % 2), values.get()); + assertEquals(index, values.getInt()); + assertEquals(Double.doubleToRawLongBits(index + 0.5), values.getLong()); + assertFalse(values.hasRemaining()); + if (index == 3) { + throw new IoTDBConnectionException("simulated transport failure"); + } + return null; + }) + .when(connection) + .executePreparedNumericRow(eq(1L), anyLong(), any(ByteBuffer.class), any(ByteBuffer.class)); + for (int i = 0; i < tags.length; i++) { + try { + tableSession.executePreparedInsertNumericRow( + prepared, + i, + new long[] {i % 2, i, Double.doubleToRawLongBits(i + 0.5)}, + new String[] {tags[i]}); + assertNotEquals(3, i); + } catch (IoTDBConnectionException e) { + assertEquals(3, i); + } + } + assertEquals(tags.length, tagBuffers.size()); + assertNotSame(tagBuffers.get(0), tagBuffers.get(1)); + for (int i = 1; i < tags.length; i++) { + assertSame(valueBuffers.get(0), valueBuffers.get(i)); + assertSame(tagBuffers.get(1), tagBuffers.get(i)); + } + } + + @Test + public void testBuffersAreOwnedByEachStatementWithoutTags() throws Exception { + Session session = mock(Session.class); + SessionConnection connection = mock(SessionConnection.class); + when(session.getDefaultSessionConnection()).thenReturn(connection); + TableSession tableSession = new TableSession(session); + List<ByteBuffer> buffers = new ArrayList<>(); + doAnswer( + invocation -> { + ByteBuffer values = invocation.getArgument(2); + ByteBuffer tags = invocation.getArgument(3); + assertEquals(0, tags.remaining()); + assertEquals(42L, values.getLong()); + assertFalse(values.hasRemaining()); + buffers.add(values); + buffers.add(tags); + return null; + }) + .when(connection) + .executePreparedNumericRow( + anyLong(), anyLong(), any(ByteBuffer.class), any(ByteBuffer.class)); + for (int i = 0; i < 2; i++) { + NumericRowPreparedStatement prepared = + new NumericRowPreparedStatement( + "t", List.of("v"), List.of(TSDataType.INT64), List.of(ColumnCategory.FIELD)); + prepared.bind(i, connection, () -> {}); + tableSession.executePreparedInsertNumericRow(prepared, 0, new long[] {42}, new String[0]); + } + assertNotSame(buffers.get(0), buffers.get(2)); + assertNotSame(buffers.get(1), buffers.get(3)); + } +} diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index d4ce7f97ba8..d7e2081e760 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java @@ -1490,5 +1490,27 @@ public final class DataNodeMiscMessages { "table_query_device_entry_batch_size_in_bytes (%d) exceeds the maximum RPC payload (dn_thrift_max_frame_size %d minus 1024 bytes); using %d as the effective value"; public static final String EXCEPTION_INVALID_PREPARED_NUMERIC_ROW_METADATA_OR_PAYLOAD_AECCCEAA = "Invalid prepared numeric row metadata or payload."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_REQUEST_IS_NULL_5A9F9BA9 = "Prepared numeric row request is null."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_DATABASE_IS_NULL_F2A2E0AC = "Prepared numeric row database is null."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TABLE_NAME_IS_NULL_OR_EMPTY_F2D38BE5 = "Prepared numeric row table name is null or empty."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_MEASUREMENTS_ARE_NULL_830DE646 = "Prepared numeric row measurements are null."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPES_ARE_NULL_7C15F477 = "Prepared numeric row data types are null."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_CATEGORIES_ARE_NULL_6F2EE1BE = "Prepared numeric row column categories are null."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_HAS_NO_COLUMNS_D4E40365 = "Prepared numeric row has no columns."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_METADATA_SIZE_MISMATCH_MEASUREMENTS_ARG_DATA_TYPES_ARG_CATEGORIES_ARG_23DAB34B = "Prepared numeric row column metadata size mismatch: measurements %d, data types %d, categories %d."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_NAME_IS_NULL_OR_EMPTY_AT_INDEX_ARG_DEBDCFF8 = "Prepared numeric row column name is null or empty at index %d."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_NAME_IS_DUPLICATED_AT_INDEX_ARG_ARG_73AAFECC = "Prepared numeric row column name is duplicated at index %d: %s."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPE_IS_NULL_AT_INDEX_ARG_FB2A1FD6 = "Prepared numeric row data type is null at index %d."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPE_ORDINAL_IS_INVALID_AT_INDEX_ARG_ARG_BB5059BB = "Prepared numeric row data type ordinal is invalid at index %d: %d."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_CATEGORY_IS_NULL_AT_INDEX_ARG_3ED8315E = "Prepared numeric row column category is null at index %d."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_IS_INVALID_AT_INDEX_ARG_CATEGORY_ARG_AND_TYPE_ARG_9DAEEA10 = "Prepared numeric row column is invalid at index %d: category %d and type %s."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_HAS_NO_FIELD_COLUMNS_A76209AE = "Prepared numeric row has no FIELD columns."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TAG_BUFFER_IS_NULL_3FF27BF7 = "Prepared numeric row tag buffer is null."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_BUFFER_IS_NULL_A928F919 = "Prepared numeric row value buffer is null."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_BUFFER_SIZE_MISMATCH_EXPECTED_ARG_BYTES_BUT_GOT_ARG_B75D48C8 = "Prepared numeric row value buffer size mismatch: expected %d bytes but got %d."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TAG_LENGTH_PREFIX_IS_INCOMPLETE_AT_INDEX_ARG_ARG_BYTES_REMAINING_6D342F3A = "Prepared numeric row tag length prefix is incomplete at index %d: %d bytes remaining."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TAG_LENGTH_IS_INVALID_AT_INDEX_ARG_LENGTH_ARG_ARG_BYTES_REMAINING_DEDD1E8F = "Prepared numeric row tag length is invalid at index %d: length %d, %d bytes remaining."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TAG_BUFFER_HAS_ARG_TRAILING_BYTES_042CF4D2 = "Prepared numeric row tag buffer has %d trailing bytes."; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_PAYLOAD_DECODING_FAILED_ARG_F87689DD = "Prepared numeric row payload decoding failed: %s."; public static final String EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_DOES_NOT_BELONG_TO_THIS_SESSION_5CC5C50A = "Prepared numeric row is closed or does not belong to this session."; } diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index 2d19c74be8b..c8ecb24f9e8 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java @@ -1471,6 +1471,28 @@ public final class DataNodeMiscMessages { public static final String EXCEPTION_INVALID_PREPARED_NUMERIC_ROW_METADATA_OR_PAYLOAD_AECCCEAA = "预备数值行元数据或数据载荷无效"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_REQUEST_IS_NULL_5A9F9BA9 = "预备数值行请求为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_DATABASE_IS_NULL_F2A2E0AC = "预备数值行数据库为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TABLE_NAME_IS_NULL_OR_EMPTY_F2D38BE5 = "预备数值行表名为 null 或为空"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_MEASUREMENTS_ARE_NULL_830DE646 = "预备数值行 measurements 为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPES_ARE_NULL_7C15F477 = "预备数值行数据类型为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_CATEGORIES_ARE_NULL_6F2EE1BE = "预备数值行列类别为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_HAS_NO_COLUMNS_D4E40365 = "预备数值行没有列"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_METADATA_SIZE_MISMATCH_MEASUREMENTS_ARG_DATA_TYPES_ARG_CATEGORIES_ARG_23DAB34B = "预备数值行列元数据数量不匹配：measurements %d，数据类型 %d，列类别 %d"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_NAME_IS_NULL_OR_EMPTY_AT_INDEX_ARG_DEBDCFF8 = "预备数值行列名在下标 %d 处为 null 或为空"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_NAME_IS_DUPLICATED_AT_INDEX_ARG_ARG_73AAFECC = "预备数值行列名在下标 %d 处重复：%s"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPE_IS_NULL_AT_INDEX_ARG_FB2A1FD6 = "预备数值行数据类型在下标 %d 处为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPE_ORDINAL_IS_INVALID_AT_INDEX_ARG_ARG_BB5059BB = "预备数值行数据类型序号在下标 %d 处无效：%d"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_CATEGORY_IS_NULL_AT_INDEX_ARG_3ED8315E = "预备数值行列类别在下标 %d 处为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_IS_INVALID_AT_INDEX_ARG_CATEGORY_ARG_AND_TYPE_ARG_9DAEEA10 = "预备数值行列在下标 %d 处无效：类别 %d，类型 %s"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_HAS_NO_FIELD_COLUMNS_A76209AE = "预备数值行没有 FIELD 列"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TAG_BUFFER_IS_NULL_3FF27BF7 = "预备数值行 TAG 缓冲区为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_BUFFER_IS_NULL_A928F919 = "预备数值行 FIELD 缓冲区为 null"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_BUFFER_SIZE_MISMATCH_EXPECTED_ARG_BYTES_BUT_GOT_ARG_B75D48C8 = "预备数值行 FIELD 缓冲区大小不匹配：期望 %d 字节，实际 %d 字节"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TAG_LENGTH_PREFIX_IS_INCOMPLETE_AT_INDEX_ARG_ARG_BYTES_REMAINING_6D342F3A = "预备数值行 TAG 长度前缀在下标 %d 处不完整：剩余 %d 字节"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TAG_LENGTH_IS_INVALID_AT_INDEX_ARG_LENGTH_ARG_ARG_BYTES_REMAINING_DEDD1E8F = "预备数值行 TAG 长度在下标 %d 处无效：长度 %d，剩余 %d 字节"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_TAG_BUFFER_HAS_ARG_TRAILING_BYTES_042CF4D2 = "预备数值行 TAG 缓冲区有 %d 个尾随字节"; + public static final String EXCEPTION_PREPARED_NUMERIC_ROW_PAYLOAD_DECODING_FAILED_ARG_F87689DD = "预备数值行数据载荷解码失败：%s"; public static final String EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_DOES_NOT_BELONG_TO_THIS_SESSION_5CC5C50A = "预备数值行已关闭或不属于当前会话"; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/PreparedNumericRowInfo.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/PreparedNumericRowInfo.java new file mode 100644 index 00000000000..a1a67e45100 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/PreparedNumericRowInfo.java @@ -0,0 +1,242 @@ +/* + * 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.protocol.session; + +import org.apache.iotdb.commons.exception.SemanticException; +import org.apache.iotdb.commons.path.PartialPath; +import org.apache.iotdb.commons.queryengine.utils.TimestampPrecisionUtils; +import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory; +import org.apache.iotdb.db.i18n.DataNodeMiscMessages; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertNumericRowStatement; +import org.apache.iotdb.db.utils.NumericRowValues; +import org.apache.iotdb.service.rpc.thrift.TSPrepareNumericRowReq; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.utils.RamUsageEstimator; +import org.apache.tsfile.utils.ReadWriteIOUtils; + +import java.nio.BufferUnderflowException; +import java.nio.ByteBuffer; +import java.util.HashSet; +import java.util.Locale; +import java.util.Set; + +/** + * Session-owned metadata. Each execution gets fresh arrays because schema analysis mutates them. + */ +public final class PreparedNumericRowInfo { + private final String database; + private final String table; + private final String[] measurements; + private final TSDataType[] types; + private final TsTableColumnCategory[] categories; + private final int tagCount; + private final int valueBytes; + + public PreparedNumericRowInfo(String database, TSPrepareNumericRowReq req) { + if (req == null) { + throw invalid(DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_REQUEST_IS_NULL_5A9F9BA9); + } + if (database == null) { + throw invalid(DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_DATABASE_IS_NULL_F2A2E0AC); + } + if (req.getTableName() == null || req.getTableName().isEmpty()) { + throw invalid( + DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_TABLE_NAME_IS_NULL_OR_EMPTY_F2D38BE5); + } + if (req.getMeasurements() == null) { + throw invalid( + DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_MEASUREMENTS_ARE_NULL_830DE646); + } + if (req.getDataTypes() == null) { + throw invalid( + DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPES_ARE_NULL_7C15F477); + } + if (req.getColumnCategories() == null) { + throw invalid( + DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_CATEGORIES_ARE_NULL_6F2EE1BE); + } + if (req.getMeasurementsSize() == 0) { + throw invalid(DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_HAS_NO_COLUMNS_D4E40365); + } + if (req.getMeasurementsSize() != req.getDataTypesSize() + || req.getMeasurementsSize() != req.getColumnCategoriesSize()) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_METADATA_SIZE_MISMATCH_MEASUREMENTS_ARG_DATA_TYPES_ARG_CATEGORIES_ARG_23DAB34B, + req.getMeasurementsSize(), + req.getDataTypesSize(), + req.getColumnCategoriesSize()); + } + this.database = database; + table = req.getTableName(); + measurements = req.getMeasurements().toArray(new String[0]); + types = new TSDataType[measurements.length]; + categories = new TsTableColumnCategory[measurements.length]; + int tags = 0; + int bytes = 0; + Set<String> names = new HashSet<>(); + for (int i = 0; i < measurements.length; i++) { + Integer type = req.getDataTypes().get(i); + Byte category = req.getColumnCategories().get(i); + if (measurements[i] == null || measurements[i].isEmpty()) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_NAME_IS_NULL_OR_EMPTY_AT_INDEX_ARG_DEBDCFF8, + i); + } + if (!names.add(measurements[i].toLowerCase(Locale.ROOT))) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_NAME_IS_DUPLICATED_AT_INDEX_ARG_ARG_73AAFECC, + i, + measurements[i]); + } + if (type == null) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPE_IS_NULL_AT_INDEX_ARG_FB2A1FD6, + i); + } + if (type < 0 || type >= TSDataType.values().length) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_DATA_TYPE_ORDINAL_IS_INVALID_AT_INDEX_ARG_ARG_BB5059BB, + i, + type); + } + if (category == null) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_CATEGORY_IS_NULL_AT_INDEX_ARG_3ED8315E, + i); + } + types[i] = TSDataType.values()[type]; + if (category == ColumnCategory.TAG.ordinal() && types[i] == TSDataType.STRING && tags == i) { + tags++; + categories[i] = TsTableColumnCategory.TAG; + } else if (category == ColumnCategory.FIELD.ordinal() + && NumericRowValues.supports(types[i])) { + categories[i] = TsTableColumnCategory.FIELD; + bytes = Math.addExact(bytes, types[i].getDataTypeSize()); + } else { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_COLUMN_IS_INVALID_AT_INDEX_ARG_CATEGORY_ARG_AND_TYPE_ARG_9DAEEA10, + i, + category, + types[i]); + } + } + tagCount = tags; + valueBytes = bytes; + if (tagCount == measurements.length) { + throw invalid( + DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_HAS_NO_FIELD_COLUMNS_A76209AE); + } + } + + public InsertNumericRowStatement bind(long timestamp, ByteBuffer tags, ByteBuffer values) { + if (tags == null) { + throw invalid( + DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_TAG_BUFFER_IS_NULL_3FF27BF7); + } + if (values == null) { + throw invalid( + DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_BUFFER_IS_NULL_A928F919); + } + if (values.remaining() != valueBytes) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_BUFFER_SIZE_MISMATCH_EXPECTED_ARG_BYTES_BUT_GOT_ARG_B75D48C8, + valueBytes, + values.remaining()); + } + TimestampPrecisionUtils.checkTimestampPrecision(timestamp); + NumericRowValues row = new NumericRowValues(tagCount, types.length - tagCount); + // Duplicate cursors so validation never changes a caller's buffers, even on failure. + tags = tags.duplicate(); + values = values.duplicate(); + try { + for (int i = 0; i < tagCount; i++) { + if (tags.remaining() < Integer.BYTES) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_TAG_LENGTH_PREFIX_IS_INCOMPLETE_AT_INDEX_ARG_ARG_BYTES_REMAINING_6D342F3A, + i, + tags.remaining()); + } + int length = tags.getInt(tags.position()); + if (length < -1 || length > tags.remaining() - Integer.BYTES) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_TAG_LENGTH_IS_INVALID_AT_INDEX_ARG_LENGTH_ARG_ARG_BYTES_REMAINING_DEDD1E8F, + i, + length, + tags.remaining()); + } + row.setTag(i, ReadWriteIOUtils.readString(tags)); + } + if (tags.hasRemaining()) { + throw invalid( + DataNodeMiscMessages + .EXCEPTION_PREPARED_NUMERIC_ROW_TAG_BUFFER_HAS_ARG_TRAILING_BYTES_042CF4D2, + tags.remaining()); + } + for (int i = tagCount; i < types.length; i++) { + row.read(i, types[i], values); + } + } catch (BufferUnderflowException e) { + throw invalid( + DataNodeMiscMessages.EXCEPTION_PREPARED_NUMERIC_ROW_PAYLOAD_DECODING_FAILED_ARG_F87689DD, + e.getMessage()); + } + InsertNumericRowStatement statement = new InsertNumericRowStatement(); + statement.setDevicePath(new PartialPath(table, false)); + statement.setDatabaseName(database); + statement.setMeasurements(measurements.clone()); + statement.setDataTypes(types.clone()); + statement.setColumnCategories(categories.clone()); + statement.setTime(timestamp); + statement.setNumericValues(row); + return statement; + } + + public long ramBytesUsed() { + return RamUsageEstimator.shallowSizeOfInstance(PreparedNumericRowInfo.class) + + RamUsageEstimator.sizeOf(database) + + RamUsageEstimator.sizeOf(table) + + RamUsageEstimator.sizeOf(measurements) + + RamUsageEstimator.shallowSizeOf(types) + + RamUsageEstimator.shallowSizeOf(categories) + // Account for the map entry and boxed ID retained by the session. + + 64; + } + + private static SemanticException invalid() { + return new SemanticException( + DataNodeMiscMessages.EXCEPTION_INVALID_PREPARED_NUMERIC_ROW_METADATA_OR_PAYLOAD_AECCCEAA); + } + + private static SemanticException invalid(String message, Object... args) { + return new SemanticException(String.format(message, args)); + } +}
