From a6901c30e39ad7b3e66888ab20d46cf9160d1cd1 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Mon, 31 Aug 2026 18:09:14 +0800 Subject: [PATCH 1/8] [LOAD] fix encoded chunk memory sizing --- .../consensus/config/IoTConsensusConfig.java | 18 +- .../iot/logdispatcher/LogDispatcher.java | 148 ++- .../iotdb/db/i18n/DataNodeQueryMessages.java | 10 + .../iotdb/db/i18n/StorageEngineMessages.java | 50 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 10 + .../iotdb/db/i18n/StorageEngineMessages.java | 48 + .../org/apache/iotdb/db/conf/IoTDBConfig.java | 2 +- .../node/load/LoadTsFileConsensusNode.java | 807 ++++++++++++ .../plan/node/load/LoadTsFileConsensusOp.java | 51 + .../plan/node/load/LoadTsFilePieceNode.java | 18 +- .../load/DataPartitionBatchFetcher.java | 114 ++ .../scheduler/load/DataPartitionRouter.java | 94 ++ .../load/LoadConsensusSubmitter.java | 202 +++ .../scheduler/load/LoadFallbackHandler.java | 212 ++++ .../load/LoadTsFileDispatcherImpl.java | 4 + .../scheduler/load/LoadTsFileScheduler.java | 1110 ++++------------- .../scheduler/load/LocalLoadStrategy.java | 203 +++ .../scheduler/load/MemoryBoundedBuffer.java | 78 ++ .../plan/scheduler/load/PieceDispatcher.java | 261 ++++ .../load/RegionConsensusContext.java | 93 ++ .../scheduler/load/TsFileLoadStrategy.java | 46 + .../scheduler/load/TsFileSplitConsumer.java | 137 ++ .../load/TwoPhaseConsensusLoadStrategy.java | 418 +++++++ .../load/splitter/AlignedChunkData.java | 5 + .../load/splitter/ChunkData.java | 2 + .../load/splitter/ChunkEncoder.java | 54 + .../load/splitter/EncodedChunk.java | 62 + .../load/splitter/EncodedChunkGroup.java | 68 + .../load/splitter/MemoryTsFileOutput.java | 72 ++ .../load/splitter/NonAlignedChunkData.java | 5 + .../load/splitter/TsFileDataType.java | 1 + .../planner/node/load/LoadTsFileNodeTest.java | 4 +- .../load/LoadTsFileConsensusNodeTest.java | 133 ++ .../load/LoadTsFileSchedulerTest.java | 90 +- .../plan/planner/plan/node/PlanNodeType.java | 1 + 35 files changed, 3706 insertions(+), 925 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusOp.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionBatchFetcher.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadConsensusSubmitter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LocalLoadStrategy.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/MemoryBoundedBuffer.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileLoadStrategy.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkEncoder.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunk.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunkGroup.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/MemoryTsFileOutput.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java diff --git a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/config/IoTConsensusConfig.java b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/config/IoTConsensusConfig.java index c6ca9a7e0dd5c..a84e67f87172b 100644 --- a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/config/IoTConsensusConfig.java +++ b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/config/IoTConsensusConfig.java @@ -311,6 +311,7 @@ public static class Replication { private final int maxLogEntriesNumPerBatch; private final int maxSizePerBatch; private final int maxPendingBatchesNum; + private final int shardMaxSize; private final int maxQueueLength; private final long maxWaitingTimeForWaitBatchInMs; @@ -344,10 +345,12 @@ private Replication( long regionMigrationSpeedLimitBytesPerSecond, long subscriptionWalRetentionSizeInBytes, long subscriptionWalRetentionTimeMs, - long snapshotTransmissionProgressLogIntervalMs) { + long snapshotTransmissionProgressLogIntervalMs, + int shardMaxSize) { this.maxLogEntriesNumPerBatch = maxLogEntriesNumPerBatch; this.maxSizePerBatch = maxSizePerBatch; this.maxPendingBatchesNum = maxPendingBatchesNum; + this.shardMaxSize = shardMaxSize; this.maxQueueLength = maxQueueLength; this.maxWaitingTimeForWaitBatchInMs = maxWaitingTimeForWaitBatchInMs; this.maxWaitingTimeForAccumulatingBatchInMs = maxWaitingTimeForAccumulatingBatchInMs; @@ -376,6 +379,10 @@ public int getMaxPendingBatchesNum() { return maxPendingBatchesNum; } + public int getShardMaxSize() { + return shardMaxSize; + } + public int getMaxQueueLength() { return maxQueueLength; } @@ -441,6 +448,7 @@ public static class Builder { private int maxLogEntriesNumPerBatch = 1024; private int maxSizePerBatch = 16 * 1024 * 1024; private int maxPendingBatchesNum = 5; + private int shardMaxSize = 4 * 1024 * 1024; private int maxQueueLength = 4096; private long maxWaitingTimeForWaitBatchInMs = 10 * 1000L; @@ -478,6 +486,11 @@ public Replication.Builder setMaxPendingBatchesNum(int maxPendingBatchesNum) { return this; } + public Replication.Builder setShardMaxSize(int shardMaxSize) { + this.shardMaxSize = shardMaxSize; + return this; + } + public Builder setMaxQueueLength(int maxQueueLength) { this.maxQueueLength = maxQueueLength; return this; @@ -571,7 +584,8 @@ public Replication build() { regionMigrationSpeedLimitBytesPerSecond, subscriptionWalRetentionSizeInBytes, subscriptionWalRetentionTimeMs, - snapshotTransmissionProgressLogIntervalMs); + snapshotTransmissionProgressLogIntervalMs, + shardMaxSize); } } } diff --git a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java index 304badbd8aef3..1ac46f1b3d2e3 100644 --- a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java +++ b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java @@ -24,6 +24,8 @@ import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory; import org.apache.iotdb.commons.concurrent.ThreadName; +import org.apache.iotdb.commons.request.ConsensusRequestShardSource; +import org.apache.iotdb.commons.request.ShardableConsensusRequest; import org.apache.iotdb.commons.service.metric.MetricService; import org.apache.iotdb.commons.subscription.config.SubscriptionConfig; import org.apache.iotdb.consensus.common.Peer; @@ -42,6 +44,8 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Collections; import java.util.Iterator; import java.util.LinkedList; import java.util.List; @@ -252,6 +256,10 @@ public class LogDispatcherThread implements Runnable { private final CountDownLatch runFinished = new CountDownLatch(1); private volatile long lastIdleWriterSafeTimeBarrierSentTimeMs = 0L; + // An in-flight shardable request that is still being streamed out shard by shard. While this + // is non-null, the dispatcher keeps appending its remaining shards to batches before any entry + // with a larger searchIndex, so the receiver can reassemble the group in order. + private PendingShard pendingShard; public LogDispatcherThread(Peer peer, IoTConsensusConfig config, long initialSyncIndex) { this.peer = peer; @@ -457,6 +465,15 @@ public Batch getBatch() { } Batch batches = new Batch(config); + // If a shardable request is still being streamed out, finish sending its remaining shards + // first so the group stays contiguous and is never interleaved with newer entries. + drainPendingShard(batches); + if (pendingShard != null || !batches.canAccumulate()) { + // The batch is full before the in-flight group completed (or there is nothing left to + // drain into it), so send only what has been accumulated so far. + batches.buildIndex(); + return batches; + } // This condition will be executed in several scenarios: // 1. restart // 2. The getBatch() is invoked immediately at the moment the PendingEntries are consumed @@ -654,16 +671,135 @@ private boolean constructBatchFromWAL(long currentIndex, long maxIndex, Batch lo private void constructBatchIndexedFromConsensusRequest( IndexedConsensusRequest request, Batch logBatches) { - TLogEntry logEntry = - new TLogEntry( - request.getSerializedRequests(), - request.getSearchIndex(), - false, - request.getMemorySize()); + appendSerializedRequests(request, false, logBatches); + } + + /** + * Sends the remaining shards of an in-flight shardable request into {@code batch}. Stops when + * the batch is full, leaving {@link #pendingShard} non-null so the next {@link #getBatch()} + * continues from where it stopped. + */ + private void drainPendingShard(Batch batch) { + while (pendingShard != null && pendingShard.hasNext() && batch.canAccumulate()) { + batch.addTLogEntry(pendingShard.nextEntry()); + } + if (pendingShard != null && !pendingShard.hasNext()) { + pendingShard = null; + } + } + + /** + * Appends a request's serialized bytes into {@code batch}. When the request is shardable and + * its serialized form is larger than the configured shard size, the payload is split into + * multiple {@link TLogEntry}s that share the same searchIndex, each carrying a single shard, so + * it is streamed out shard by shard instead of as one oversized message. Shards that do not fit + * into this batch are kept in {@link #pendingShard}. + * + * @return true if the whole request was consumed; false if the batch filled up before the + * request was fully appended (leftover shards are retained in {@link #pendingShard}). + */ + private boolean appendSerializedRequests( + IndexedConsensusRequest data, boolean fromWAL, Batch logBatches) { + final int shardMaxSize = config.getReplication().getShardMaxSize(); + if (isShardableAndOversized(data, shardMaxSize)) { + final long shardGroupId = + ((ShardableConsensusRequest) data.getRequests().get(0)).shardGroupId(); + final PendingShard group = + new PendingShard( + new ConsensusRequestShardSource( + toSingleBuffer(data.getSerializedRequests()), shardGroupId, shardMaxSize), + data.getSearchIndex(), + fromWAL, + data.getRoutingEpoch(), + data.getPhysicalTime()); + while (group.hasNext() && logBatches.canAccumulate()) { + logBatches.addTLogEntry(group.nextEntry()); + } + if (group.hasNext()) { + pendingShard = group; + return false; + } + return true; + } + addLogEntry(data.getSerializedRequests(), data, fromWAL, data.getMemorySize(), logBatches); + return true; + } + + private boolean isShardableAndOversized(IndexedConsensusRequest data, int shardMaxSize) { + if (data.getRequests().isEmpty() + || !(data.getRequests().get(0) instanceof ShardableConsensusRequest)) { + return false; + } + long total = 0; + for (ByteBuffer buffer : data.getSerializedRequests()) { + total += buffer.remaining(); + } + return total > shardMaxSize; + } + + private void addLogEntry( + List data, + IndexedConsensusRequest request, + boolean fromWAL, + long memorySize, + Batch logBatches) { + TLogEntry logEntry = new TLogEntry(data, request.getSearchIndex(), fromWAL, memorySize); logEntry.setRoutingEpoch(request.getRoutingEpoch()); logEntry.setPhysicalTime(request.getPhysicalTime()); logBatches.addTLogEntry(logEntry, request.containsUserData()); } + + private static ByteBuffer toSingleBuffer(List buffers) { + if (buffers.size() == 1) { + return buffers.get(0); + } + int total = 0; + for (ByteBuffer buffer : buffers) { + total += buffer.remaining(); + } + final ByteBuffer merged = ByteBuffer.allocate(total); + for (ByteBuffer buffer : buffers) { + merged.put(buffer.duplicate()); + } + merged.flip(); + return merged; + } + } + + /** In-flight metadata of a shardable request that is still being streamed out. */ + private static class PendingShard { + + private final ConsensusRequestShardSource source; + private final long searchIndex; + private final boolean fromWAL; + private final long routingEpoch; + private final long physicalTime; + + private PendingShard( + ConsensusRequestShardSource source, + long searchIndex, + boolean fromWAL, + long routingEpoch, + long physicalTime) { + this.source = source; + this.searchIndex = searchIndex; + this.fromWAL = fromWAL; + this.routingEpoch = routingEpoch; + this.physicalTime = physicalTime; + } + + private boolean hasNext() { + return source.hasNext(); + } + + private TLogEntry nextEntry() { + final ByteBuffer shard = source.next(); + final TLogEntry entry = + new TLogEntry(Collections.singletonList(shard), searchIndex, fromWAL, shard.remaining()); + entry.setRoutingEpoch(routingEpoch); + entry.setPhysicalTime(physicalTime); + return entry; + } } public static AtomicLong getReceiverMemSizeSum() { diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index bf7a36010030a..5615ee3eab443 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -1175,6 +1175,9 @@ public final class DataNodeQueryMessages { "Start load TsFile {} locally."; public static final String LOAD_ALL_FAILED_TSFILES_ARE_CONVERTED_TO_TABLETS = "Load: all failed TsFiles are converted to tablets and inserted."; + public static final String + LOG_LOAD_FAILED_TO_LOAD_SOME_TSFILES_BY_CONVERTING_THEM_INTO_TABLETS_FAILED_TSFILES_ARG_7D9DB9C3 = + "Load: failed to load some TsFiles by converting them into tablets. Failed TsFiles: %s"; // --- Plan / Statement --- @@ -2703,6 +2706,8 @@ public final class DataNodeQueryMessages { "LoadTsFileScheduler: Region migration was detected during loading TsFile {}, will convert to insertion to avoid data loss"; public static final String LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG = "Load TsFile {} Successfully, load process [{}/{}]"; + public static final String LOG_LOAD_BATCH_FINISHED_DELETING_ARG_SOURCE_TSFILES_AFTER_LOAD_D5EE56E9 = + "LOAD batch finished, deleting %d source TsFiles after load"; public static final String CAN_NOT_LOAD_TSFILE_ARG_LOAD_PROCESS_ARG_ARG = "Can not Load TsFile {}, load process [{}/{}]"; public static final String LOAD_TSFILE_S_FAILED_WILL_TRY_TO_CONVERT_TO_TABLETS_AND_INSERT_FAILED_TSFILES_ARG = @@ -2713,6 +2718,8 @@ public final class DataNodeQueryMessages { "Parse or send TsFile %s error."; public static final String DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG = "Dispatch one piece to ReplicaSet {} error. Result status code {}. "; + public static final String LOG_LOAD_CONSENSUS_SUBMIT_TRANSIENT_FAILURE_RETRY_D7E1D9A6 = + "Transient failure while submitting LOAD consensus {} (load {}) to {}, will retry ({}/{}): {}"; public static final String RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG = "Result status message {}. Dispatch piece node error:%n{}"; public static final String SUB_STATUS_CODE_ARG_SUB_STATUS_MESSAGE_ARG = @@ -3816,6 +3823,7 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_THE_SECOND_ARGUMENT_OF_PERCENTILE_FUNCTION_PERCENTAGE_MUST_BE_A_DOUBLE_LITERAL_D9464B46 = "The second argument of 'percentile' function percentage must be a double literal"; public static final String EXCEPTION_DATA_TYPE_MISMATCH_FOR_MEASUREMENT_ARGARGARG_TYPE_IN_TSFILE_ARG_TYPE_IN_IOTDB_ARG_C5BA7DBD = "Data type mismatch for measurement %s%s%s, type in TsFile: %s, type in IoTDB: %s"; public static final String MESSAGE_FAILED_TO_RELEASE_EXTERNAL_TSFILE_QUERY_RESOURCE_712EE978 = "Failed to release external TsFile query resource"; + public static final String EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 = "Unknown LoadTsFileConsensusOp ordinal: "; public static final String EXCEPTION_OUTER_QUERY_TIMEOUT_EXCEEDED_BEFORE_IOTDBLOCAL_QUERY_STARTS_800BFA63 = "Outer query timeout exceeded before IoTDBLocal query starts"; public static final String MESSAGE_FAILED_TO_CLOSE_UDF_RESULT_SET_AT_INDEX_ARG_A293B7EC = "Failed to close UDF result set at index {}"; public static final String EXCEPTION_INTERNAL_QUERY_EXECUTION_NOT_FOUND_62642542 = "Internal query execution not found"; @@ -3879,5 +3887,7 @@ private DataNodeQueryMessages() {} "No more DeviceEntry records are available"; public static final String EXCEPTION_ONLY_INMEMORYDEVICEENTRYDATASET_SUPPORTS_GET_INLINE_DEVICE_ENTRIES_07A52CAB = "Only InMemoryDeviceEntryDataSet supports get inline device entries"; + public static final String EXCEPTION_LOAD_CONSENSUS_INVALID_PIECE_REF_F3498507 = + "Invalid LOAD consensus piece ref: path %s, offset %d, size %d"; } diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java index ccde46c4881b9..fccc3a53dbd37 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -29,6 +29,54 @@ private StorageEngineMessages() {} // ======================== StorageEngine ======================== public static final String FAIL_TO_RECOVER_WAL = "Fail to recover wal."; + public static final String LOG_LOAD_CONSENSUS_WRITE_TO_REGION_ARG_VIA_PROTOCOL_ARG_EBB55042 = + "Write LOAD consensus node to region {} via protocol {}"; + public static final String LOG_LOAD_CONSENSUS_REFRESH_REPLICA_SET_FAILED_7C244C63 = + "Failed to refresh LOAD consensus replica set for region {}, using cached set: {}"; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_CHECKSUM_MISMATCH_CF261675 = + "LOAD consensus piece checksum mismatch, loadId: %s, pieceIndex: %d"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_EOF_8743387D = + "Unexpected end of file when reading staged piece %s at offset %s."; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_DATA_MISSING_AFTER_PULL_8269CB0B = + "LOAD piece %d data of load %s is still missing after pulling from the write node."; + public static final String MESSAGE_LOAD_CONSENSUS_PULL_PIECE_NOT_SUPPORTED_71EC4B46 = + "Fetching piece from leader is not supported."; + public static final String EXCEPTION_LOAD_CONSENSUS_PIECE_DATA_MISSING_OR_CHECKSUM_MISMATCH_AFTER_PULL_35F4972E = + "LOAD task %s piece %d data is missing or its checksum mismatches after pull."; + public static final String LOG_LOAD_CONSENSUS_RETAINED_PIECE_READ_FAILED_0659D19B = + "Failed to read retained LOAD piece {} of load {} from {}: {}"; + public static final String LOG_LOAD_CONSENSUS_RETAINED_PIECE_WRITE_FAILED_99697608 = + "Failed to write retained LOAD piece {} of load {} to {}: {}"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19 = + "Staged file %s of load %s is not continuous: expected offset %d but current file length is %d."; + public static final String MESSAGE_LOAD_CONSENSUS_PREPARE_WITHOUT_STAGED_DATA_FE8ADC37 = + "Cannot prepare load %s because no staged data exists on this node."; + public static final String LOG_LOAD_CONSENSUS_RECOVERED_TASK_02824CE6 = + "Recovered in-progress LOAD task {} from disk; staged data is kept until COMMIT or ABORT."; + public static final String LOG_LOAD_CONSENSUS_RECOVER_TASK_META_FAILED_C39E04BB = + "Failed to recover the task meta of LOAD task {}: {}"; + public static final String LOG_LOAD_CONSENSUS_TASK_META_WRITE_FAILED_5D2420BF = + "Failed to persist the task meta of LOAD task {} to {}: {}"; + public static final String LOG_LOAD_CONSENSUS_TERMINAL_MARKER_WRITE_FAILED_4D6D7433 = + "Failed to write the terminal marker of LOAD task {} to {}: {}"; + public static final String LOG_LOAD_CONSENSUS_RECOVER_TASK_UNRESUMABLE_A159436C = + "Cannot resume LOAD task {} from disk: its durable task meta is missing or corrupt; the next " + + "command will re-create it."; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_SHORT_WRITE_E7392FAD = + "Staged file %s of load %s was not fully written: expected %d bytes at offset %d, wrote %d"; + public static final String EXCEPTION_LOAD_CONSENSUS_PROGRESS_SERIALIZE_FAILED_28EFD091 = + "Failed to serialize LOAD progress index for time partition %d"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_INCOMPLETE_1CDE954B = + "Staged file %s of load %s is incomplete and cannot be committed."; + public static final String LOG_LOAD_CONSENSUS_SNAPSHOT_TAKEN_09A7DD4C = + "Snapshotted %d in-progress LOAD task(s) with %d staged file(s) for region %s into %s."; + public static final String LOG_LOAD_CONSENSUS_SNAPSHOT_RESTORED_90ABC1BF = + "Restored %d in-progress LOAD task(s) with %d staged file(s) from snapshot %s."; + public static final String EXCEPTION_LOAD_CONSENSUS_SNAPSHOT_RESTORE_FAILED_F8C29C64 = + "Failed to restore LOAD snapshot from %s: %s"; + public static final String EXCEPTION_LOAD_TSFILE_ALIGNED_VALUE_CHUNK_TIME_CHUNK_EEB00760 = + "Cannot attach value chunk of measurement %s in file %s: expected exactly one buffered " + + "aligned time chunk, found %d."; public static final String STORAGE_ENGINE_FAILED_TO_SET_UP = "Storage engine failed to set up."; public static final String SEQ_MEMTABLE_FLUSH_CHECK_THREAD_STARTED = "start sequence memtable timed flush check thread successfully."; public static final String UNSEQ_MEMTABLE_FLUSH_CHECK_THREAD_STARTED = "start unsequence memtable timed flush check thread successfully."; @@ -478,6 +526,8 @@ private StorageEngineMessages() {} public static final String CANNOT_CREATE_TSFILE_FOR_WRITING = "Can not create TsFile {} for writing."; public static final String CLOSE_TSFILE_IO_WRITER_ERROR = "Close TsFileIOWriter {} error."; public static final String CLOSE_MODIFICATION_FILE_ERROR = "Close ModificationFile {} error."; + public static final String EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93 = + "Table '%s.%s' does not exist when applying LOAD chunk data. It may have been dropped after the LOAD was analyzed."; public static final String TASK_DIR_NOT_EMPTY_SKIP_DELETE = "Task dir {} is not empty, skip deleting."; public static final String LOAD_CLEANUP_TASK_CANCELED = "Load cleanup task {} is canceled."; public static final String LOAD_CLEANUP_TASK_STARTS = "Load cleanup task {} starts."; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index f26f2306be2fa..f080afcc594bb 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -1156,6 +1156,9 @@ public final class DataNodeQueryMessages { "开始本地加载 TsFile {}。"; public static final String LOAD_ALL_FAILED_TSFILES_ARE_CONVERTED_TO_TABLETS = "加载：所有失败的 TsFile 已转换为 Tablet 并插入。"; + public static final String + LOG_LOAD_FAILED_TO_LOAD_SOME_TSFILES_BY_CONVERTING_THEM_INTO_TABLETS_FAILED_TSFILES_ARG_7D9DB9C3 = + "加载：部分 TsFile 通过转换为 Tablet 仍加载失败。失败的 TsFile：%s"; // --- Plan / Statement --- @@ -3180,6 +3183,8 @@ public final class DataNodeQueryMessages { "LoadTsFileScheduler：加载 TsFile {} 期间检测到 Region 迁移，将转换为 insertion 以避免数据丢失。"; public static final String LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG = "成功加载 TsFile {}，加载进度 [{}/{}]"; + public static final String LOG_LOAD_BATCH_FINISHED_DELETING_ARG_SOURCE_TSFILES_AFTER_LOAD_D5EE56E9 = + "LOAD 批次结束，删除 %d 个源 TsFile"; public static final String CAN_NOT_LOAD_TSFILE_ARG_LOAD_PROCESS_ARG_ARG = "无法加载 TsFile {}，加载进度 [{}/{}]。"; @@ -3195,6 +3200,8 @@ public final class DataNodeQueryMessages { public static final String DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG = "分发 TsFile 片段到 ReplicaSet {} 出错。结果状态码 {}。 "; + public static final String LOG_LOAD_CONSENSUS_SUBMIT_TRANSIENT_FAILURE_RETRY_D7E1D9A6 = + "提交 LOAD 共识 {}（load {}）到 {} 时遇到瞬时失败，将重试（{}/{}）：{}"; public static final String RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG = "结果状态消息 {}。分发片段节点出错：%n{}"; @@ -4572,6 +4579,7 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_THE_SECOND_ARGUMENT_OF_PERCENTILE_FUNCTION_PERCENTAGE_MUST_BE_A_DOUBLE_LITERAL_D9464B46 = "'percentile' 函数的第二个参数 percentage 必须是 double 字面量"; public static final String EXCEPTION_DATA_TYPE_MISMATCH_FOR_MEASUREMENT_ARGARGARG_TYPE_IN_TSFILE_ARG_TYPE_IN_IOTDB_ARG_C5BA7DBD = "测点 %s%s%s 的数据类型不匹配，TsFile 中类型：%s，IoTDB 中类型：%s"; public static final String MESSAGE_FAILED_TO_RELEASE_EXTERNAL_TSFILE_QUERY_RESOURCE_712EE978 = "释放外部 TsFile 查询资源失败"; + public static final String EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 = "未知的 LoadTsFileConsensusOp 序号："; public static final String EXCEPTION_OUTER_QUERY_TIMEOUT_EXCEEDED_BEFORE_IOTDBLOCAL_QUERY_STARTS_800BFA63 = "在 IoTDBLocal 查询开始前，外层查询已超时"; public static final String MESSAGE_FAILED_TO_CLOSE_UDF_RESULT_SET_AT_INDEX_ARG_A293B7EC = "关闭索引 {} 处的 UDF 结果集失败"; public static final String EXCEPTION_INTERNAL_QUERY_EXECUTION_NOT_FOUND_62642542 = "未找到内部查询执行"; @@ -4636,5 +4644,7 @@ private DataNodeQueryMessages() {} "没有更多可用的 DeviceEntry 记录"; public static final String EXCEPTION_ONLY_INMEMORYDEVICEENTRYDATASET_SUPPORTS_GET_INLINE_DEVICE_ENTRIES_07A52CAB = "只有 InMemoryDeviceEntryDataSet 支持获取内存中的设备条目"; + public static final String EXCEPTION_LOAD_CONSENSUS_INVALID_PIECE_REF_F3498507 = + "无效的 LOAD 共识分片引用：路径 %s，偏移 %d，大小 %d"; } diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 1244879ed425c..f0de191a06ad9 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -29,6 +29,52 @@ private StorageEngineMessages() {} // ======================== StorageEngine ======================== public static final String FAIL_TO_RECOVER_WAL = "WAL 恢复失败。"; + public static final String LOG_LOAD_CONSENSUS_WRITE_TO_REGION_ARG_VIA_PROTOCOL_ARG_EBB55042 = + "通过协议 {} 向 Region {} 写入 LOAD 共识节点"; + public static final String LOG_LOAD_CONSENSUS_REFRESH_REPLICA_SET_FAILED_7C244C63 = + "刷新 Region {} 的 LOAD 共识副本集失败，使用缓存的副本集：{}"; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_CHECKSUM_MISMATCH_CF261675 = + "LOAD 共识分片校验和不一致，loadId: %s，pieceIndex: %d"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_EOF_8743387D = + "读取已暂存的分片 %s 时意外到达文件末尾，offset: %s。"; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_DATA_MISSING_AFTER_PULL_8269CB0B = + "LOAD 分片 %d（load %s）的数据在向写节点回补后仍然缺失。"; + public static final String MESSAGE_LOAD_CONSENSUS_PULL_PIECE_NOT_SUPPORTED_71EC4B46 = + "已不支持向 Leader 回补 LOAD 分片。"; + public static final String EXCEPTION_LOAD_CONSENSUS_PIECE_DATA_MISSING_OR_CHECKSUM_MISMATCH_AFTER_PULL_35F4972E = + "回补后 LOAD 任务 %s 的分片 %d 数据仍缺失或校验和不一致。"; + public static final String LOG_LOAD_CONSENSUS_RETAINED_PIECE_READ_FAILED_0659D19B = + "读取保留的 LOAD 分片 {}（load {}）失败，文件：{}，原因：{}"; + public static final String LOG_LOAD_CONSENSUS_RETAINED_PIECE_WRITE_FAILED_99697608 = + "写入保留的 LOAD 分片 {}（load {}）失败，文件：{}，原因：{}"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19 = + "load %s 的暂存文件 %s 不连续：期望偏移 %d，但当前文件长度为 %d。"; + public static final String MESSAGE_LOAD_CONSENSUS_PREPARE_WITHOUT_STAGED_DATA_FE8ADC37 = + "无法准备（PREPARE）load %s，因为该节点上不存在暂存数据。"; + public static final String LOG_LOAD_CONSENSUS_RECOVERED_TASK_02824CE6 = + "已从磁盘恢复进行中的 LOAD 任务 {}；暂存数据将保留至 COMMIT 或 ABORT。"; + public static final String LOG_LOAD_CONSENSUS_RECOVER_TASK_META_FAILED_C39E04BB = + "恢复 LOAD 任务 {} 的任务元数据失败：{}"; + public static final String LOG_LOAD_CONSENSUS_TASK_META_WRITE_FAILED_5D2420BF = + "持久化 LOAD 任务 {} 的任务元数据到 {} 失败：{}"; + public static final String LOG_LOAD_CONSENSUS_TERMINAL_MARKER_WRITE_FAILED_4D6D7433 = + "写入 LOAD 任务 {} 的终止标记到 {} 失败：{}"; + public static final String LOG_LOAD_CONSENSUS_RECOVER_TASK_UNRESUMABLE_A159436C = + "无法从磁盘恢复 LOAD 任务 {}：其持久化任务元数据缺失或损坏，下一条命令将重新创建该任务。"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_SHORT_WRITE_E7392FAD = + "load %s 的暂存文件 %s 未完整写入：偏移 %d 处期望 %d 字节，实际写入 %d 字节"; + public static final String EXCEPTION_LOAD_CONSENSUS_PROGRESS_SERIALIZE_FAILED_28EFD091 = + "时间分区 %d 的 LOAD 进度索引序列化失败"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_INCOMPLETE_1CDE954B = + "load %s 的暂存文件 %s 不完整，无法提交（COMMIT）。"; + public static final String LOG_LOAD_CONSENSUS_SNAPSHOT_TAKEN_09A7DD4C = + "已将 region %s 的 %d 个进行中的 LOAD 任务（共 %d 个暂存文件）纳入快照 %s。"; + public static final String LOG_LOAD_CONSENSUS_SNAPSHOT_RESTORED_90ABC1BF = + "已从快照 %s 恢复 %d 个进行中的 LOAD 任务（共 %d 个暂存文件）。"; + public static final String EXCEPTION_LOAD_CONSENSUS_SNAPSHOT_RESTORE_FAILED_F8C29C64 = + "从 %s 恢复 LOAD 快照失败：%s"; + public static final String EXCEPTION_LOAD_TSFILE_ALIGNED_VALUE_CHUNK_TIME_CHUNK_EEB00760 = + "无法将测量 %s 的值 Chunk 挂载到文件 %s：预期恰好一个已缓冲的 Aligned 时间 Chunk，实际发现 %d 个。"; public static final String STORAGE_ENGINE_FAILED_TO_SET_UP = "存储引擎启动失败。"; public static final String SEQ_MEMTABLE_FLUSH_CHECK_THREAD_STARTED = "顺序 memtable 定时 flush 检查线程启动成功。"; public static final String UNSEQ_MEMTABLE_FLUSH_CHECK_THREAD_STARTED = "乱序 memtable 定时 flush 检查线程启动成功。"; @@ -478,6 +524,8 @@ private StorageEngineMessages() {} public static final String CANNOT_CREATE_TSFILE_FOR_WRITING = "无法创建 TsFile {} 用于写入。"; public static final String CLOSE_TSFILE_IO_WRITER_ERROR = "关闭 TsFileIOWriter {} 出错。"; public static final String CLOSE_MODIFICATION_FILE_ERROR = "关闭修改文件 {} 出错。"; + public static final String EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93 = + "应用 LOAD chunk 数据时表 '%s.%s' 不存在，可能在 LOAD 分析之后被删除了。"; public static final String TASK_DIR_NOT_EMPTY_SKIP_DELETE = "任务目录 {} 非空，跳过删除。"; public static final String LOAD_CLEANUP_TASK_CANCELED = "加载清理任务 {} 已取消。"; public static final String LOAD_CLEANUP_TASK_STARTS = "加载清理任务 {} 开始。"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java index 80f9f1f1b30f8..ef6f9b08ca688 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java @@ -4467,7 +4467,7 @@ public int getLoadTsFileSpiltPartitionMaxSize() { } public void setLoadTsFileSpiltPartitionMaxSize(int loadTsFileSpiltPartitionMaxSize) { - if (loadTsFileSpiltPartitionMaxSize <= 0) { + if (loadTsFileSpiltPartitionMaxSize < 0) { throw new IllegalArgumentException( DataNodeMiscMessages .MISC_EXCEPTION_LOADTSFILESPILTPARTITIONMAXSIZE_SHOULD_BE_GREATER_THAN_OR_95B4DB23); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java new file mode 100644 index 0000000000000..da8ffe144f500 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java @@ -0,0 +1,807 @@ +/* + * 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.queryengine.plan.planner.plan.node.load; + +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.IPlanVisitor; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; +import org.apache.iotdb.commons.request.ShardableConsensusRequest; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.WritePlanNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.SearchNode; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.IWALByteBufferView; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryType; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryValue; +import org.apache.iotdb.db.storageengine.load.LoadTsFileManager; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; + +import org.apache.tsfile.exception.NotImplementedException; +import org.apache.tsfile.exception.write.PageException; +import org.apache.tsfile.utils.PublicBAOS; +import org.apache.tsfile.utils.ReadWriteIOUtils; + +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.RandomAccessFile; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** + * Consensus-backed LOAD request carrying Begin / Piece / Seal / Commit phases. Submitted once to + * the DataRegion write peer (or Ratis leader); replicas apply via the consensus state machine. + */ +public class LoadTsFileConsensusNode extends SearchNode + implements WALEntryValue, ShardableConsensusRequest { + + private LoadTsFileConsensusOp op; + private String loadId; + private String tsFileId; + private boolean isTableModel; + private String database; + private int expectedPieceCount = -1; + private long pieceIndex = -1; + private long pieceOffset = -1; + private List pieceRefs = new ArrayList<>(); + private long dataSize; + private long checksum; + private int pieceCount = -1; + private long totalBytes; + private boolean isGeneratedByPipe; + private boolean deleteAfterLoad; + private List tsFileDataList = new ArrayList<>(); + private Map timePartition2ProgressIndex = new HashMap<>(); + private TRegionReplicaSet regionReplicaSet; + private ProgressIndex progressIndex; + private boolean isGeneratedByRemoteConsensusLeader; + + /** Endpoint ("ip:port") of the follower that issued a PULL, used to push the piece back. */ + private String pullSourceEndPoint; + + public LoadTsFileConsensusNode(PlanNodeId id) { + super(id); + } + + public static LoadTsFileConsensusNode begin( + PlanNodeId id, + String loadId, + String tsFileId, + boolean isTableModel, + String database, + int expectedPieceCount) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.BEGIN; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.isTableModel = isTableModel; + node.database = database == null ? "" : database; + node.expectedPieceCount = expectedPieceCount; + return node; + } + + public static LoadTsFileConsensusNode piece( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + long pieceOffset, + List tsFileDataList, + long checksum) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PIECE; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceIndex = pieceIndex; + node.pieceOffset = pieceOffset; + node.tsFileDataList = + tsFileDataList == null ? new ArrayList<>() : new ArrayList<>(tsFileDataList); + node.dataSize = node.tsFileDataList.stream().mapToLong(TsFileData::getDataSize).sum(); + // The checksum is part of the consensus contract and must be preserved exactly as supplied by + // the caller. Recomputing it here silently changes the value for callers that use a staged + // piece checksum (and made marker/piece validation disagree). + node.checksum = checksum; + return node; + } + + public static LoadTsFileConsensusNode pieceRef( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + String relativePath, + long offset, + long size, + long checksum) { + return pieceRefs( + id, + loadId, + tsFileId, + pieceIndex, + Collections.singletonList(new PieceRef(relativePath, offset, size)), + checksum, + size); + } + + public static LoadTsFileConsensusNode pieceRefs( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + List refs, + long checksum, + long dataSize) { + return pieceRefs(id, loadId, tsFileId, pieceIndex, refs, checksum, dataSize, null); + } + + public static LoadTsFileConsensusNode pieceRefs( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + List refs, + long checksum, + long dataSize, + List deletions) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PIECE; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceIndex = pieceIndex; + node.pieceRefs = refs == null ? new ArrayList<>() : new ArrayList<>(refs); + node.dataSize = dataSize; + node.checksum = checksum; + if (deletions != null) { + node.tsFileDataList = new ArrayList<>(deletions); + } + return node; + } + + /** + * Marker-only PIECE: carries the piece metadata (index, checksum, byte size) but neither chunk + * data nor piece refs. This is what the write node logs to the WAL after applying a chunk-data + * piece. The write node keeps the actual chunk bytes in its retained-piece store; a follower + * applies them only when this marker arrives, and a missing piece fails directly, so the WAL + * stays at dozens of bytes per piece instead of the full LOAD bytes. + */ + public static LoadTsFileConsensusNode pieceMarker( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + long checksum, + long dataSize) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PIECE; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceIndex = pieceIndex; + node.checksum = checksum; + node.dataSize = dataSize; + return node; + } + + /** Legacy PULL request; no longer supported and kept only for compatibility with older nodes. */ + public static LoadTsFileConsensusNode pull( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + long checksum, + String pullSourceEndPoint) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PULL; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceIndex = pieceIndex; + node.checksum = checksum; + node.pullSourceEndPoint = pullSourceEndPoint; + return node; + } + + public static LoadTsFileConsensusNode pieceRef( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + String relativePath, + long offset, + long size, + long checksum, + long dataSize) { + final LoadTsFileConsensusNode node = + pieceRef(id, loadId, tsFileId, pieceIndex, relativePath, offset, size, checksum); + node.dataSize = dataSize; + return node; + } + + public static LoadTsFileConsensusNode prepare( + PlanNodeId id, + String loadId, + String tsFileId, + int pieceCount, + long totalBytes, + long checksum, + Map timePartition2ProgressIndex) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PREPARE; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceCount = pieceCount; + node.totalBytes = totalBytes; + node.checksum = checksum; + node.timePartition2ProgressIndex = + timePartition2ProgressIndex == null + ? new HashMap<>() + : new HashMap<>(timePartition2ProgressIndex); + return node; + } + + public static LoadTsFileConsensusNode abort( + PlanNodeId id, String loadId, String tsFileId, boolean isGeneratedByPipe) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.ABORT; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.isGeneratedByPipe = isGeneratedByPipe; + return node; + } + + public static LoadTsFileConsensusNode commit( + PlanNodeId id, + String loadId, + String tsFileId, + boolean isGeneratedByPipe, + boolean deleteAfterLoad, + Map timePartition2ProgressIndex) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.COMMIT; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.isGeneratedByPipe = isGeneratedByPipe; + node.deleteAfterLoad = deleteAfterLoad; + node.timePartition2ProgressIndex = + timePartition2ProgressIndex == null + ? new HashMap<>() + : new HashMap<>(timePartition2ProgressIndex); + return node; + } + + public LoadTsFileConsensusOp getOp() { + return op; + } + + public String getLoadId() { + return loadId; + } + + public String getTsFileId() { + return tsFileId; + } + + public boolean isTableModel() { + return isTableModel; + } + + public String getDatabase() { + return database; + } + + public int getExpectedPieceCount() { + return expectedPieceCount; + } + + public long getPieceIndex() { + return pieceIndex; + } + + public long getPieceOffset() { + return pieceOffset; + } + + public String getRelativePath() { + return pieceRefs.isEmpty() ? null : pieceRefs.get(0).relativePath; + } + + public long getSize() { + return pieceRefs.stream().mapToLong(ref -> ref.size).sum(); + } + + public List getPieceRefs() { + return pieceRefs; + } + + public long getDataSize() { + return dataSize; + } + + public long getChecksum() { + return checksum; + } + + public boolean isGeneratedByRemoteConsensusLeader() { + return isGeneratedByRemoteConsensusLeader; + } + + public void markAsGeneratedByRemoteConsensusLeader() { + this.isGeneratedByRemoteConsensusLeader = true; + } + + public int getPieceCount() { + return pieceCount; + } + + public long getTotalBytes() { + return totalBytes; + } + + public boolean isGeneratedByPipe() { + return isGeneratedByPipe; + } + + public boolean isDeleteAfterLoad() { + return deleteAfterLoad; + } + + public List getTsFileDataList() { + return tsFileDataList; + } + + /** Whether this PIECE carries actual chunk/deletion data (write-node path) or is a marker. */ + public boolean hasChunkData() { + return !tsFileDataList.isEmpty(); + } + + public String getPullSourceEndPoint() { + return pullSourceEndPoint; + } + + public void setPullSourceEndPoint(String pullSourceEndPoint) { + this.pullSourceEndPoint = pullSourceEndPoint; + } + + public Map getTimePartition2ProgressIndex() { + return timePartition2ProgressIndex; + } + + public void setRegionReplicaSet(TRegionReplicaSet regionReplicaSet) { + this.regionReplicaSet = regionReplicaSet; + } + + @Override + public TRegionReplicaSet getRegionReplicaSet() { + return regionReplicaSet; + } + + @Override + public ProgressIndex getProgressIndex() { + return progressIndex; + } + + @Override + public void setProgressIndex(ProgressIndex progressIndex) { + this.progressIndex = progressIndex; + } + + @Override + public SearchNode merge(List searchNodes) { + if (searchNodes.size() == 1) { + return searchNodes.get(0); + } + throw new UnsupportedOperationException(DataNodeQueryMessages.MERGE_IS_NOT_SUPPORTED); + } + + @Override + public List splitByPartition(IAnalysis analysis) { + return Collections.singletonList(this); + } + + @Override + public List getChildren() { + return Collections.emptyList(); + } + + @Override + public void addChild(PlanNode child) { + // no children + } + + @Override + public PlanNodeType getType() { + return PlanNodeType.LOAD_TSFILE_CONSENSUS; + } + + @Override + public PlanNode clone() { + throw new NotImplementedException( + DataNodeQueryMessages.CLONE_OF_LOAD_PIECE_TSFILE_IS_NOT_IMPLEMENTED); + } + + @Override + public int allowedChildCount() { + return NO_CHILD_ALLOWED; + } + + @Override + public List getOutputColumnNames() { + return Collections.emptyList(); + } + + @Override + public R accept(IPlanVisitor visitor, C context) { + return ((PlanVisitor) visitor).visitLoadTsFileConsensus(this, context); + } + + @Override + protected void serializeAttributes(ByteBuffer byteBuffer) { + PlanNodeType.LOAD_TSFILE_CONSENSUS.serialize(byteBuffer); + try { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final DataOutputStream stream = new DataOutputStream(baos); + serializeBody(stream, true); + byteBuffer.put(baos.toByteArray()); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + @Override + protected void serializeAttributes(DataOutputStream stream) throws IOException { + PlanNodeType.LOAD_TSFILE_CONSENSUS.serialize(stream); + serializeBody(stream, true); + } + + private void serializeBody(DataOutputStream stream) throws IOException { + serializeBody(stream, false); + } + + private void serializeBody(DataOutputStream stream, boolean includeContent) throws IOException { + ReadWriteIOUtils.write(op.ordinal(), stream); + ReadWriteIOUtils.write(loadId, stream); + ReadWriteIOUtils.write(tsFileId, stream); + ReadWriteIOUtils.write(isTableModel, stream); + ReadWriteIOUtils.write(database == null ? "" : database, stream); + ReadWriteIOUtils.write(expectedPieceCount, stream); + ReadWriteIOUtils.write(pieceIndex, stream); + ReadWriteIOUtils.write(pieceOffset, stream); + ReadWriteIOUtils.write(dataSize, stream); + ReadWriteIOUtils.write(checksum, stream); + ReadWriteIOUtils.write(pieceCount, stream); + ReadWriteIOUtils.write(totalBytes, stream); + ReadWriteIOUtils.write(isGeneratedByPipe, stream); + ReadWriteIOUtils.write(deleteAfterLoad, stream); + ReadWriteIOUtils.write(pieceRefs.size(), stream); + for (PieceRef ref : pieceRefs) { + ReadWriteIOUtils.write(ref.relativePath, stream); + ReadWriteIOUtils.write(ref.offset, stream); + ReadWriteIOUtils.write(ref.size, stream); + if (includeContent) { + final byte[] content = readPieceContent(ref); + ReadWriteIOUtils.write(content.length, stream); + stream.write(content); + } + } + ReadWriteIOUtils.write(tsFileDataList.size(), stream); + for (TsFileData data : tsFileDataList) { + data.serialize(stream); + } + ReadWriteIOUtils.write(timePartition2ProgressIndex.size(), stream); + for (Map.Entry entry : timePartition2ProgressIndex.entrySet()) { + ReadWriteIOUtils.write(entry.getKey().getStartTime(), stream); + ReadWriteIOUtils.write(entry.getValue().length, stream); + stream.write(entry.getValue()); + } + ReadWriteIOUtils.write(pullSourceEndPoint == null ? "" : pullSourceEndPoint, stream); + } + + public static LoadTsFileConsensusNode deserialize(ByteBuffer buffer) { + try { + final LoadTsFilePieceNode.ByteBufferInputStream stream = + new LoadTsFilePieceNode.ByteBufferInputStream(buffer); + final LoadTsFileConsensusNode node = deserializeBody(stream, true); + final PlanNodeId planNodeId = PlanNodeId.deserialize(buffer); + node.setPlanNodeId(planNodeId); + ReadWriteIOUtils.readInt(buffer); + return node; + } catch (IOException | PageException | IllegalPathException e) { + throw new IllegalStateException(e); + } + } + + private static LoadTsFileConsensusNode deserializeBody(InputStream stream) + throws IOException, PageException, IllegalPathException { + return deserializeBody(stream, false); + } + + private static LoadTsFileConsensusNode deserializeBody(InputStream stream, boolean readContent) + throws IOException, PageException, IllegalPathException { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(new PlanNodeId("")); + node.op = LoadTsFileConsensusOp.fromOrdinal(ReadWriteIOUtils.readInt(stream)); + node.loadId = ReadWriteIOUtils.readString(stream); + node.tsFileId = ReadWriteIOUtils.readString(stream); + node.isTableModel = ReadWriteIOUtils.readBool(stream); + node.database = ReadWriteIOUtils.readString(stream); + node.expectedPieceCount = ReadWriteIOUtils.readInt(stream); + node.pieceIndex = ReadWriteIOUtils.readLong(stream); + node.pieceOffset = ReadWriteIOUtils.readLong(stream); + node.dataSize = ReadWriteIOUtils.readLong(stream); + node.checksum = ReadWriteIOUtils.readLong(stream); + node.pieceCount = ReadWriteIOUtils.readInt(stream); + node.totalBytes = ReadWriteIOUtils.readLong(stream); + node.isGeneratedByPipe = ReadWriteIOUtils.readBool(stream); + node.deleteAfterLoad = ReadWriteIOUtils.readBool(stream); + final int refCount = ReadWriteIOUtils.readInt(stream); + node.pieceRefs = new ArrayList<>(refCount); + for (int i = 0; i < refCount; i++) { + final String refPath = ReadWriteIOUtils.readString(stream); + final long refOffset = ReadWriteIOUtils.readLong(stream); + final long refSize = ReadWriteIOUtils.readLong(stream); + byte[] content = null; + if (readContent) { + final int contentLength = ReadWriteIOUtils.readInt(stream); + if (contentLength > 0) { + content = new byte[contentLength]; + int offset = 0; + while (offset < contentLength) { + final int read = stream.read(content, offset, contentLength - offset); + if (read < 0) { + throw new IOException( + DataNodeQueryMessages.EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + "pieceContent"); + } + offset += read; + } + } + } + node.pieceRefs.add(new PieceRef(refPath, refOffset, refSize, content)); + } + final int dataCount = ReadWriteIOUtils.readInt(stream); + node.tsFileDataList = new ArrayList<>(dataCount); + for (int i = 0; i < dataCount; i++) { + node.tsFileDataList.add(TsFileData.deserialize(stream)); + } + final int progressCount = ReadWriteIOUtils.readInt(stream); + node.timePartition2ProgressIndex = new HashMap<>(progressCount); + for (int i = 0; i < progressCount; i++) { + final long startTime = ReadWriteIOUtils.readLong(stream); + final int len = ReadWriteIOUtils.readInt(stream); + final byte[] bytes = new byte[len]; + int offset = 0; + while (offset < len) { + final int read = stream.read(bytes, offset, len - offset); + if (read < 0) { + throw new IOException( + DataNodeQueryMessages.EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + "progressIndex"); + } + offset += read; + } + node.timePartition2ProgressIndex.put(new TTimePartitionSlot(startTime), bytes); + } + node.pullSourceEndPoint = ReadWriteIOUtils.readString(stream); + return node; + } + + @Override + public void serializeToWAL(IWALByteBufferView buffer) { + serializeToWAL(buffer, getEncodedSearchIndex()); + } + + public void serializeToWAL(IWALByteBufferView buffer, long encodedSearchIndex) { + try { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final DataOutputStream stream = new DataOutputStream(baos); + stream.writeShort(getType().getNodeType()); + stream.writeLong(encodedSearchIndex); + serializeBody(stream); + buffer.put(baos.toByteArray()); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + /** Serialize this node into an IoTConsensusRequest-compatible buffer with content expanded. */ + public ByteBuffer serialize() { + try (PublicBAOS byteArrayOutputStream = new PublicBAOS(); + DataOutputStream stream = new DataOutputStream(byteArrayOutputStream)) { + ReadWriteIOUtils.write(WALEntryType.LOAD_TSFILE_CONSENSUS_NODE.getCode(), stream); + ReadWriteIOUtils.write(-1L, stream); + ReadWriteIOUtils.write(getType().getNodeType(), stream); + serializeBody(stream, true); + getPlanNodeId().serialize(stream); + ReadWriteIOUtils.write(0, stream); + return ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + /** Stable, collision-resistant identifier of the logical request this node belongs to. */ + @Override + public long shardGroupId() { + return Objects.hash(loadId, tsFileId, pieceIndex); + } + + private byte[] readPieceContent(PieceRef ref) { + if (ref.content != null) { + return ref.content; + } + // PieceRef construction already validated offset/size bounds (non-negative, no overflow, + // within int range), so the narrowing cast below is safe. + final byte[] content = new byte[(int) ref.size]; + if (content.length == 0) { + return content; + } + final File file = + LoadTsFileManager.findLoadTsFile(ref.relativePath) + .orElseThrow( + () -> + new IllegalStateException( + DataNodeQueryMessages + .EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + "piece file not found: " + + ref.relativePath)); + try (RandomAccessFile raf = new RandomAccessFile(file, "r")) { + raf.seek(ref.offset); + raf.readFully(content); + return content; + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + @Override + public int serializedSize() { + try { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final DataOutputStream stream = new DataOutputStream(baos); + serializeBody(stream); + return Short.BYTES + Long.BYTES + baos.size(); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + public static LoadTsFileConsensusNode deserializeFromWAL(DataInputStream stream) + throws IOException { + final long searchIndex = stream.readLong(); + try { + final LoadTsFileConsensusNode node = deserializeBody(stream, false); + node.setSearchIndexFromWAL(searchIndex); + return node; + } catch (PageException | IllegalPathException e) { + throw new IOException(e); + } + } + + public static LoadTsFileConsensusNode deserializeFromWAL(ByteBuffer buffer) { + final long searchIndex = buffer.getLong(); + try { + final LoadTsFileConsensusNode node = + deserializeBody(new LoadTsFilePieceNode.ByteBufferInputStream(buffer), false); + node.setSearchIndexFromWAL(searchIndex); + return node; + } catch (IOException | PageException | IllegalPathException e) { + throw new IllegalStateException(e); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof LoadTsFileConsensusNode)) { + return false; + } + final LoadTsFileConsensusNode that = (LoadTsFileConsensusNode) o; + return Objects.equals(loadId, that.loadId) + && Objects.equals(tsFileId, that.tsFileId) + && op == that.op + && pieceIndex == that.pieceIndex + && checksum == that.checksum; + } + + @Override + public int hashCode() { + return Objects.hash(op, loadId, tsFileId, pieceIndex, checksum); + } + + @Override + public String toString() { + return "LoadTsFileConsensusNode{op=" + + op + + ", loadId=" + + loadId + + ", tsFileId=" + + tsFileId + + ", pieceIndex=" + + pieceIndex + + '}'; + } + + /** A staging-file reference recorded by the consensus/WAL entry instead of the raw content. */ + public static class PieceRef { + private final String relativePath; + private final long offset; + private final long size; + private final byte[] content; + + public PieceRef(String relativePath, long size) { + this(relativePath, 0L, size, null); + } + + public PieceRef(String relativePath, long offset, long size) { + this(relativePath, offset, size, null); + } + + public PieceRef(String relativePath, long offset, long size, byte[] content) { + // Bounds/overflow guard before the legacy raw-ref path casts size to int and seeks the file: + // a malformed or hostile ref must fail here instead of truncating, overflowing or escaping + // into an arbitrary byte array. + if (offset < 0 + || size < 0 + || size > Integer.MAX_VALUE + || offset + size < 0 + || offset + size > Integer.MAX_VALUE) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages.EXCEPTION_LOAD_CONSENSUS_INVALID_PIECE_REF_F3498507, + relativePath, + offset, + size)); + } + this.relativePath = relativePath; + this.offset = offset; + this.size = size; + this.content = content; + } + + public String getRelativePath() { + return relativePath; + } + + public long getOffset() { + return offset; + } + + public long getSize() { + return size; + } + + public byte[] getContent() { + return content; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusOp.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusOp.java new file mode 100644 index 0000000000000..9543a0c7d664f --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusOp.java @@ -0,0 +1,51 @@ +/* + * 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.queryengine.plan.planner.plan.node.load; + +/** + * Consensus LOAD request phases, sequenced by the scheduler/client. + * + *

{@link #BEGIN} opens a load, {@link #PIECE} stages one data piece, {@link #PREPARE} seals the + * data of a load, {@link #COMMIT} imports the staged data, {@link #ABORT} discards it, and {@link + * #PULL} asks the current write node to re-deliver the serialized bytes of one already-applied + * piece; this legacy request is no longer supported and is kept only for compatibility with older + * nodes. Phase enforcement and idempotency are the client's responsibility; the server applies each + * request against the staged data directory identified by the load id without keeping per-load + * transaction state in memory. + */ +public enum LoadTsFileConsensusOp { + BEGIN, + PIECE, + PREPARE, + COMMIT, + ABORT, + PULL; + + public static LoadTsFileConsensusOp fromOrdinal(int ordinal) { + final LoadTsFileConsensusOp[] values = values(); + if (ordinal < 0 || ordinal >= values.length) { + throw new IllegalArgumentException( + org.apache.iotdb.db.i18n.DataNodeQueryMessages + .EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + ordinal); + } + return values[ordinal]; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFilePieceNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFilePieceNode.java index df49b0e52e5da..cb89122439d22 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFilePieceNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFilePieceNode.java @@ -50,6 +50,7 @@ public class LoadTsFilePieceNode extends WritePlanNode { private static final Logger LOGGER = LoggerFactory.getLogger(LoadTsFilePieceNode.class); private File tsFile; + private long pieceIndex; private long dataSize; private List tsFileDataList; @@ -59,12 +60,21 @@ public LoadTsFilePieceNode(PlanNodeId id) { } public LoadTsFilePieceNode(PlanNodeId id, File tsFile) { + this(id, tsFile, 0L); + } + + public LoadTsFilePieceNode(PlanNodeId id, File tsFile, long pieceIndex) { super(id); this.tsFile = tsFile; + this.pieceIndex = pieceIndex; this.dataSize = 0; this.tsFileDataList = new ArrayList<>(); } + public long getPieceIndex() { + return pieceIndex; + } + public long getDataSize() { return dataSize; } @@ -133,6 +143,7 @@ protected void serializeAttributes(ByteBuffer byteBuffer) { @Override protected void serializeAttributes(DataOutputStream stream) throws IOException { PlanNodeType.LOAD_TSFILE.serialize(stream); + ReadWriteIOUtils.write(pieceIndex, stream); ReadWriteIOUtils.write(tsFile.getPath(), stream); // TODO: can save this space ReadWriteIOUtils.write(tsFileDataList.size(), stream); for (TsFileData tsFileData : tsFileDataList) { @@ -160,8 +171,10 @@ public static PlanNode deserialize(ByteBuffer buffer) { ByteBufferInputStream stream = new ByteBufferInputStream(buffer); try { ReadWriteIOUtils.readShort(stream); // read PlanNodeType + final long pieceIndex = ReadWriteIOUtils.readLong(stream); final File tsFile = new File(ReadWriteIOUtils.readString(stream)); - final LoadTsFilePieceNode pieceNode = new LoadTsFilePieceNode(new PlanNodeId(""), tsFile); + final LoadTsFilePieceNode pieceNode = + new LoadTsFilePieceNode(new PlanNodeId(""), tsFile, pieceIndex); final int tsFileDataSize = ReadWriteIOUtils.readInt(stream); for (int i = 0; i < tsFileDataSize; i++) { TsFileData tsFileData = TsFileData.deserialize(stream); @@ -185,13 +198,14 @@ public boolean equals(Object o) { } LoadTsFilePieceNode loadTsFilePieceNode = (LoadTsFilePieceNode) o; return Objects.equals(tsFile, loadTsFilePieceNode.tsFile) + && pieceIndex == loadTsFilePieceNode.pieceIndex && Objects.equals(dataSize, loadTsFilePieceNode.dataSize) && Objects.equals(tsFileDataList, loadTsFilePieceNode.tsFileDataList); } @Override public int hashCode() { - return Objects.hash(tsFile, dataSize, tsFileDataList); + return Objects.hash(tsFile, pieceIndex, dataSize, tsFileDataList); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionBatchFetcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionBatchFetcher.java new file mode 100644 index 0000000000000..d6e3da78998dc --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionBatchFetcher.java @@ -0,0 +1,114 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.partition.DataPartition; +import org.apache.iotdb.commons.partition.DataPartitionQueryParam; +import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.Pair; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Batch partition fetcher for LOAD. It wraps the query-engine {@link IPartitionFetcher} and hides + * two details from the rest of the pipeline: + * + *

    + *
  • Transmit limit. The requested (device, time-partition) pairs are split into batches + * of at most {@code TTimePartitionSlotTransmitLimit} entries; each batch is resolved with one + * {@code getOrCreateDataPartition} call. + *
  • Database hint. {@link #setDatabase(String)} enables the explicit database lookup + * used by table-model loads and pipe-generated tree-model loads. + *
+ * + * {@link #queryDataPartition(List, String)} returns one {@link TRegionReplicaSet} per input pair, + * in the same order, which {@link DataPartitionRouter} maps back onto chunks. + */ +class DataPartitionBatchFetcher { + + private static final int TRANSMIT_LIMIT = + CommonDescriptor.getInstance().getConfig().getTTimePartitionSlotTransmitLimit(); + + private final IPartitionFetcher fetcher; + private String database; + + DataPartitionBatchFetcher(IPartitionFetcher fetcher) { + this.fetcher = fetcher; + } + + void setDatabase(String database) { + this.database = database; + } + + List queryDataPartition( + List> slotList, String userName) { + List replicaSets = new ArrayList<>(slotList.size()); + int size = slotList.size(); + + for (int i = 0; i < size; i += TRANSMIT_LIMIT) { + List> subSlotList = + slotList.subList(i, Math.min(size, i + TRANSMIT_LIMIT)); + DataPartition dataPartition = + fetcher.getOrCreateDataPartition(toQueryParam(subSlotList), userName); + for (final Pair pair : subSlotList) { + // database is an explicit database hint for table-model loads and + // pipe-generated tree-model loads. + replicaSets.add( + database != null + ? dataPartition.getDataRegionReplicaSetForWriting(pair.left, pair.right, database) + : dataPartition.getDataRegionReplicaSetForWriting(pair.left, pair.right)); + } + } + return replicaSets; + } + + private List toQueryParam( + List> slots) { + final Map> device2TimePartitionSlots = new HashMap<>(); + for (final Pair slot : slots) { + device2TimePartitionSlots.computeIfAbsent(slot.left, key -> new HashSet<>()).add(slot.right); + } + + final List queryParams = + new ArrayList<>(device2TimePartitionSlots.size()); + for (final Map.Entry> entry : + device2TimePartitionSlots.entrySet()) { + final DataPartitionQueryParam queryParam = + new DataPartitionQueryParam(entry.getKey(), new ArrayList<>(entry.getValue())); + // database is an explicit database hint for table-model loads and + // pipe-generated tree-model loads. + if (database != null) { + queryParam.setDatabaseName(database); + } + queryParams.add(queryParam); + } + return queryParams; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java new file mode 100644 index 0000000000000..16c53c7816e03 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java @@ -0,0 +1,94 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.Pair; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * LOAD chunk router: maps chunks to their target regions. {@link #route(List)} takes the + * directionless chunk buffer and: + * + *
    + *
  1. deduplicates the (device, time-partition) pairs, so many chunks over the same slot issue + * only one partition query per distinct slot; + *
  2. resolves every distinct pair through {@link DataPartitionBatchFetcher}; + *
  3. returns the replica set of every input chunk, preserving the input order. + *
+ * + * The caller ({@code TsFileSplitConsumer.routeChunkData()}) feeds the result to {@link + * PieceDispatcher}, which performs the replica-set-change detection and appends the chunk to the + * per-region piece. + */ +class DataPartitionRouter { + + private final DataPartitionBatchFetcher partitionFetcher; + private final String userName; + + DataPartitionRouter(DataPartitionBatchFetcher partitionFetcher, String userName) { + this.partitionFetcher = partitionFetcher; + this.userName = userName; + } + + /** + * Returns, for every chunk in the input list (same order), the region replica set it must be + * written to. + */ + List route(List chunkDataList) { + if (chunkDataList.isEmpty()) { + return new ArrayList<>(); + } + + final List> partitionSlotList = new ArrayList<>(); + final int[] chunkPartitionIndexes = new int[chunkDataList.size()]; + final Map> partitionSlotIndexes = new HashMap<>(); + for (int i = 0, size = chunkDataList.size(); i < size; i++) { + final ChunkData chunkData = chunkDataList.get(i); + final IDeviceID device = chunkData.getDevice(); + final TTimePartitionSlot timePartitionSlot = chunkData.getTimePartitionSlot(); + final Map slotIndexes = + partitionSlotIndexes.computeIfAbsent(device, key -> new HashMap<>()); + Integer partitionSlotIndex = slotIndexes.get(timePartitionSlot); + if (partitionSlotIndex == null) { + partitionSlotIndex = partitionSlotList.size(); + slotIndexes.put(timePartitionSlot, partitionSlotIndex); + partitionSlotList.add(new Pair<>(device, timePartitionSlot)); + } + chunkPartitionIndexes[i] = partitionSlotIndex; + } + + final List replicaSets = + partitionFetcher.queryDataPartition(partitionSlotList, userName); + final List routedReplicaSets = new ArrayList<>(chunkDataList.size()); + for (int i = 0, size = chunkDataList.size(); i < size; i++) { + routedReplicaSets.add(replicaSets.get(chunkPartitionIndexes[i])); + } + return routedReplicaSets; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadConsensusSubmitter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadConsensusSubmitter.java new file mode 100644 index 0000000000000..f428236700eff --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadConsensusSubmitter.java @@ -0,0 +1,202 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; +import org.apache.iotdb.commons.consensus.ConsensusGroupId; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.consensus.common.Peer; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.consensus.DataRegionConsensusImpl; +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult; +import org.apache.iotdb.db.queryengine.execution.executor.RegionWriteExecutor; +import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; +import org.apache.iotdb.mpp.rpc.thrift.TPlanNode; +import org.apache.iotdb.mpp.rpc.thrift.TSendBatchPlanNodeReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendSinglePlanNodeReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendSinglePlanNodeResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Collections; +import java.util.List; + +/** + * Transport for LOAD consensus commands (BEGIN / PIECE / PREPARE / COMMIT / ABORT). One instance is + * shared by all files of a scheduler; it is stateless besides the local endpoint and the client + * manager. + * + *

{@link #submit(TRegionReplicaSet, LoadTsFileConsensusNode)}: + * + *

    + *
  1. stamps the target {@code regionReplicaSet} onto the node for correlation only - no follower + * endpoints are carried, replicas receive the command through consensus log replication + * exactly like ordinary writes; + *
  2. resolves the single write peer of the partition: the current Ratis leader when the protocol + * is Ratis, otherwise the first replica-set location (the IoTConsensus write node, the same + * target the normal write path dispatches to); + *
  3. writes the command through {@link RegionWriteExecutor} (local) or the internal RPC ({@code + * sendBatchPlanNode}) on that peer, which applies it via {@code + * DataRegionConsensusImpl.write} like any other write plan. + *
+ * + *

IoTConsensus replicates the WAL entries (marker-only for LOAD pieces) to the followers, whose + * own {@code TsFileWriterManager} rebuilds the staged files; the chunk bytes are pulled back from + * the write node on demand. Ratis replicates the full command through its own log, so every replica + * applies the chunk data directly and keeps its own writer. + */ +public class LoadConsensusSubmitter { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadConsensusSubmitter.class); + + private final String localhostIp; + private final int localhostPort; + private final IClientManager clientManager; + + public LoadConsensusSubmitter( + IClientManager clientManager) { + this.clientManager = clientManager; + this.localhostIp = IoTDBDescriptor.getInstance().getConfig().getInternalAddress(); + this.localhostPort = IoTDBDescriptor.getInstance().getConfig().getInternalPort(); + } + + public TSStatus submit(TRegionReplicaSet replicaSet, LoadTsFileConsensusNode node) { + final ConsensusGroupId regionId = + ConsensusGroupId.Factory.createFromTConsensusGroupId(replicaSet.getRegionId()); + // Follow the freshest partition route: after a write-node switch (leader change / region + // migration) the replica set captured at split time may be stale, so re-resolve it from the + // local partition table (a cache miss fetches the latest route map from the ConfigNode) and + // submit to the current write node, exactly like normal writes. + final TRegionReplicaSet currentReplicaSet = refreshReplicaSet(replicaSet, regionId); + node.setRegionReplicaSet(currentReplicaSet); + + final String protocol = + IoTDBDescriptor.getInstance().getConfig().getDataRegionConsensusProtocolClass(); + LOGGER.info( + StorageEngineMessages.LOG_LOAD_CONSENSUS_WRITE_TO_REGION_ARG_VIA_PROTOCOL_ARG_EBB55042, + regionId, + protocol); + + final TDataNodeLocation writePeer = resolveWritePeer(currentReplicaSet, regionId, protocol); + if (writePeer == null) { + return new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode()) + .setMessage(String.valueOf(replicaSet)); + } + return isLocal(writePeer.getInternalEndPoint()) + ? writeLocal(regionId, node) + : writeRemote(writePeer.getInternalEndPoint(), regionId, node); + } + + /** + * Re-resolves the partition replica set from the local partition table (falling back to the + * passed set when the lookup fails). IoTConsensus V1 has no leader election, so after a write + * node switch the coordinator must route by the refreshed route map instead of the stale set. + */ + private TRegionReplicaSet refreshReplicaSet( + TRegionReplicaSet replicaSet, ConsensusGroupId regionId) { + try { + final List replicaSets = + ClusterPartitionFetcher.getInstance() + .getRegionReplicaSet( + Collections.singletonList(regionId.convertToTConsensusGroupId())); + if (!replicaSets.isEmpty()) { + return replicaSets.get(0); + } + } catch (Exception e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_REFRESH_REPLICA_SET_FAILED_7C244C63, + regionId, + e.getMessage()); + } + return replicaSet; + } + + /** + * Resolves the single write peer of the partition. Ratis writes must land on the current leader, + * so the leader endpoint is matched against the replica-set locations first (falling back to the + * first location while the leader is not known yet); IoTConsensus routes every write to the first + * replica-set location, which is the partition's write node - the same target normal writes use. + */ + private TDataNodeLocation resolveWritePeer( + TRegionReplicaSet replicaSet, ConsensusGroupId regionId, String protocol) { + final List locations = replicaSet.getDataNodeLocations(); + if (locations == null || locations.isEmpty()) { + return null; + } + if (ConsensusFactory.RATIS_CONSENSUS.equals(protocol)) { + final Peer leader = DataRegionConsensusImpl.getInstance().getLeader(regionId); + if (leader != null) { + for (TDataNodeLocation location : locations) { + final TEndPoint endPoint = location.getInternalEndPoint(); + if (endPoint != null && endPoint.getIp().equals(leader.getEndpoint().getIp())) { + return location; + } + } + } + } + return locations.get(0); + } + + private TSStatus writeLocal(ConsensusGroupId regionId, LoadTsFileConsensusNode node) { + final RegionWriteExecutor executor = new RegionWriteExecutor(); + final RegionExecutionResult result = executor.execute(regionId, node); + return result.getStatus(); + } + + private TSStatus writeRemote( + TEndPoint endPoint, ConsensusGroupId regionId, LoadTsFileConsensusNode node) { + try (SyncDataNodeInternalServiceClient client = clientManager.borrowClient(endPoint)) { + final TSendSinglePlanNodeReq singleReq = + new TSendSinglePlanNodeReq( + new TPlanNode(node.serializeToByteBuffer()), regionId.convertToTConsensusGroupId()); + final TSendBatchPlanNodeReq batchReq = + new TSendBatchPlanNodeReq(Collections.singletonList(singleReq)); + final List responses = + client.sendBatchPlanNode(batchReq).getResponses(); + if (responses == null || responses.isEmpty()) { + return new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } + final TSendSinglePlanNodeResp resp = responses.get(0); + if (resp.isAccepted()) { + return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()); + } + return resp.getStatus() != null + ? resp.getStatus() + : new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } catch (Exception e) { + return new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode()).setMessage(e.getMessage()); + } + } + + private boolean isLocal(TEndPoint endPoint) { + return endPoint != null + && localhostIp.equals(endPoint.getIp()) + && localhostPort == endPoint.getPort(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java new file mode 100644 index 0000000000000..c73fb444ff69b --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java @@ -0,0 +1,212 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadTsFile; +import org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement; +import org.apache.iotdb.db.storageengine.load.converter.LoadTsFileDataTypeConverter; +import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.FileNotFoundException; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.ListIterator; +import java.util.stream.Collectors; + +/** + * Failure fallback of the LOAD scheduler: when at least one TsFile fails, this handler converts the + * failed files into tablets and retries the insertion, then resolves the final result. Extracted + * from the scheduler so the failure path is a single, testable component. + * + *

Flow ({@link #convertFailedTsFilesToTablets()}): + * + *

    + *
  1. builds the comma-separated list of failed file paths for logging; + *
  2. for every failed file, converts it with {@link LoadTsFileDataTypeConverter} - table-model + * files through {@code convertForTableModel}, tree-model files through {@code + * convertForTreeModel} with a retry statement built by {@code buildRetryTreeLoadStatement}; + *
  3. removes successfully converted files from the failure list; + *
  4. if no failure remains the load is considered successful (FINISHED); otherwise the state + * machine transitions to FAILED with the remaining file list. + *
+ * + * The whole conversion is measured as the {@code SCHEDULER_CAST_TABLETS} phase metric. + */ +public class LoadFallbackHandler { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadFallbackHandler.class); + + private static final LoadTsFileCostMetricsSet LOAD_TSFILE_COST_METRICS_SET = + LoadTsFileCostMetricsSet.getInstance(); + + private final MPPQueryContext queryContext; + private final boolean isGeneratedByPipe; + private final List tsFileNodeList; + private final List failedTsFileNodeIndexes; + private final QueryStateMachine stateMachine; + + public LoadFallbackHandler( + MPPQueryContext queryContext, + boolean isGeneratedByPipe, + List tsFileNodeList, + List failedTsFileNodeIndexes, + QueryStateMachine stateMachine) { + this.queryContext = queryContext; + this.isGeneratedByPipe = isGeneratedByPipe; + this.tsFileNodeList = tsFileNodeList; + this.failedTsFileNodeIndexes = failedTsFileNodeIndexes; + this.stateMachine = stateMachine; + } + + public void convertFailedTsFilesToTablets() { + final StringBuilder failedTsFiles = + new StringBuilder( + !tsFileNodeList.isEmpty() + ? tsFileNodeList + .get(failedTsFileNodeIndexes.get(0)) + .getTsFileResource() + .getTsFilePath() + : ""); + final ListIterator iterator = failedTsFileNodeIndexes.listIterator(1); + while (iterator.hasNext()) { + failedTsFiles + .append(", ") + .append(tsFileNodeList.get(iterator.next()).getTsFileResource().getTsFilePath()); + } + + final long startTime = System.nanoTime(); + try { + // if failed to load some TsFiles, then try to convert the TsFiles to Tablets + LOGGER.info( + DataNodeQueryMessages + .LOAD_TSFILE_S_FAILED_WILL_TRY_TO_CONVERT_TO_TABLETS_AND_INSERT_FAILED_TSFILES_ARG, + failedTsFiles); + convertFailedTsFilesToTabletsAndRetry(); + } finally { + LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.SCHEDULER_CAST_TABLETS, System.nanoTime() - startTime); + } + } + + private void convertFailedTsFilesToTabletsAndRetry() { + final LoadTsFileDataTypeConverter loadTsFileDataTypeConverter = + new LoadTsFileDataTypeConverter(queryContext, isGeneratedByPipe); + + final Iterator iterator = failedTsFileNodeIndexes.listIterator(); + while (iterator.hasNext()) { + final int failedLoadTsFileIndex = iterator.next(); + final LoadSingleTsFileNode failedNode = tsFileNodeList.get(failedLoadTsFileIndex); + final String filePath = failedNode.getTsFileResource().getTsFilePath(); + + try { + final TSStatus status = + failedNode.isTableModel() + ? loadTsFileDataTypeConverter + .convertForTableModel( + (isGeneratedByPipe + ? LoadTsFile.createForPipe(null, filePath, Collections.emptyMap()) + : LoadTsFile.createUnchecked( + null, filePath, Collections.emptyMap())) + .setDatabase(failedNode.getDatabase()) + .setDeleteAfterLoad(failedNode.isDeleteAfterLoad()) + .setConvertOnTypeMismatch(true)) + .orElse(null) + : loadTsFileDataTypeConverter + .convertForTreeModel( + buildRetryTreeLoadStatement( + filePath, + failedNode.isDeleteAfterLoad(), + LoadTsFileScheduler.getPartitionQueryDatabase( + failedNode, isGeneratedByPipe))) + .orElse(null); + + if (loadTsFileDataTypeConverter.isSuccessful(status)) { + iterator.remove(); + LOGGER.info( + DataNodeQueryMessages + .LOAD_SUCCESSFULLY_CONVERTED_TSFILE_ARG_INTO_TABLETS_AND_INSERTED, + failedNode.getTsFileResource().getTsFilePath()); + } else { + LOGGER.warn( + DataNodeQueryMessages.LOAD_FAILED_TO_CONVERT_TO_TABLETS_FROM_TSFILE_ARG_STATUS_ARG, + failedNode.getTsFileResource().getTsFilePath(), + status); + } + } catch (final Exception e) { + LOGGER.warn( + DataNodeQueryMessages.LOAD_FAILED_TO_CONVERT_TO_TABLETS_FROM_TSFILE_ARG_EXCEPTION_ARG, + failedNode.getTsFileResource().getTsFilePath(), + e.getMessage(), + e); + } + } + + // If all failed TsFiles are converted into tablets and inserted, + // we can consider the load process as successful. + if (failedTsFileNodeIndexes.isEmpty()) { + LOGGER.info(DataNodeQueryMessages.LOAD_ALL_FAILED_TSFILES_ARE_CONVERTED_TO_TABLETS); + stateMachine.transitionToFinished(); + } else { + final String failedFiles = + failedTsFileNodeIndexes.stream() + .map(i -> tsFileNodeList.get(i).getTsFileResource().getTsFilePath()) + .collect(Collectors.joining(", ")); + LOGGER.warn( + DataNodeQueryMessages + .LOG_LOAD_FAILED_TO_LOAD_SOME_TSFILES_BY_CONVERTING_THEM_INTO_TABLETS_FAILED_TSFILES_ARG_7D9DB9C3, + failedFiles); + stateMachine.transitionToFailed( + new LoadFileException( + String.format( + DataNodeQueryMessages + .LOG_LOAD_FAILED_TO_LOAD_SOME_TSFILES_BY_CONVERTING_THEM_INTO_TABLETS_FAILED_TSFILES_ARG_7D9DB9C3, + failedFiles))); + } + } + + private LoadTsFileStatement buildRetryTreeLoadStatement( + final String filePath, final boolean deleteAfterLoad, final String database) + throws FileNotFoundException { + final LoadTsFileStatement statement = + (isGeneratedByPipe + ? LoadTsFileStatement.createForPipe(filePath) + : LoadTsFileStatement.createUnchecked(filePath)) + .setDeleteAfterLoad(deleteAfterLoad) + .setConvertOnTypeMismatch(true); + if (database != null) { + statement.setDatabase(database); + statement.updateDatabaseLevelByTreeDatabase(); + } + if (isGeneratedByPipe) { + statement.markIsGeneratedByPipe(); + } + return statement; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java index 1ba39ffd0bd3b..5d2054c4f12a5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java @@ -76,6 +76,10 @@ import static com.google.common.util.concurrent.Futures.immediateFuture; +/** + * LOAD dispatcher: legacy local dispatcher of the LOAD local-load path (no decode needed) and the + * per-file uuid holder used for executor naming and log correlation. + */ public class LoadTsFileDispatcherImpl implements IFragInstanceDispatcher, AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(LoadTsFileDispatcherImpl.class); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java index a116f3ac1d84b..d46936d1bc93b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java @@ -19,33 +19,11 @@ package org.apache.iotdb.db.queryengine.plan.scheduler.load; -import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; -import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; -import org.apache.iotdb.common.rpc.thrift.TSStatus; -import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; -import org.apache.iotdb.commons.conf.CommonDescriptor; -import org.apache.iotdb.commons.consensus.ConsensusGroupId; -import org.apache.iotdb.commons.consensus.DataRegionId; -import org.apache.iotdb.commons.consensus.index.ProgressIndex; -import org.apache.iotdb.commons.exception.IoTDBException; -import org.apache.iotdb.commons.partition.DataPartition; -import org.apache.iotdb.commons.partition.DataPartitionQueryParam; -import org.apache.iotdb.commons.partition.StorageExecutor; -import org.apache.iotdb.commons.service.metric.MetricService; -import org.apache.iotdb.commons.service.metric.enums.Metric; -import org.apache.iotdb.commons.service.metric.enums.Tag; -import org.apache.iotdb.db.conf.IoTDBConfig; -import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.load.LoadFileException; -import org.apache.iotdb.db.exception.load.LoadReadOnlyException; -import org.apache.iotdb.db.exception.load.RegionReplicaSetChangedException; -import org.apache.iotdb.db.exception.mpp.FragmentInstanceDispatchException; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; -import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.PlanFragmentId; import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; @@ -53,86 +31,179 @@ import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher; import org.apache.iotdb.db.queryengine.plan.planner.plan.DistributedQueryPlan; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; -import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; -import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; -import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadTsFile; -import org.apache.iotdb.db.queryengine.plan.scheduler.FragInstanceDispatchResult; import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler; -import org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement; import org.apache.iotdb.db.service.RegionMigrateService; -import org.apache.iotdb.db.storageengine.StorageEngine; -import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; -import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; -import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.ArrayDeviceTimeIndex; -import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.PlainDeviceTimeIndex; -import org.apache.iotdb.db.storageengine.load.converter.LoadTsFileDataTypeConverter; import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileMemoryManager; -import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; -import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; -import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; -import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; -import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; -import org.apache.iotdb.metrics.utils.MetricLevel; -import org.apache.iotdb.mpp.rpc.thrift.TLoadCommandReq; -import org.apache.iotdb.rpc.TSStatusCode; import io.airlift.units.Duration; -import org.apache.tsfile.file.metadata.IDeviceID; -import org.apache.tsfile.file.metadata.StringArrayDeviceID; -import org.apache.tsfile.utils.Pair; -import org.apache.tsfile.utils.PublicBAOS; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.DataOutputStream; -import java.io.File; -import java.io.FileNotFoundException; -import java.io.IOException; -import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Collections; -import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; -import java.util.Iterator; import java.util.List; -import java.util.ListIterator; import java.util.Map; -import java.util.Objects; -import java.util.Optional; import java.util.Set; -import java.util.UUID; -import java.util.concurrent.CancellationException; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.stream.Collectors; /** - * {@link LoadTsFileScheduler} is used for scheduling {@link LoadSingleTsFileNode} and {@link - * LoadTsFilePieceNode}. because these two nodes need two phases to finish transfer. + * LOAD scheduler: {@link LoadTsFileScheduler} is the coordinator of a batch of {@link + * LoadSingleTsFileNode} loads (one node per source TsFile). It owns only the lifecycle - concurrent + * file lock, per-node guard clauses, state-machine transitions and the failure fallback - and + * routes every file to a {@link TsFileLoadStrategy}. All per-region bookkeeping, memory management + * and consensus submission live in the strategy components. * - *

for more details please check: ...; + *

Overall structure

+ * + *

All components below are part of the LOAD pipeline (LOAD TSFILE): the scheduler, the load + * strategies and every routing/buffering/dispatching helper are LOAD-only classes under {@code + * org.apache.iotdb.db.queryengine.plan.scheduler.load}. + * + * 

{@code
+ * LoadTsFileScheduler.start()
+ *     |
+ *     +--> per LoadSingleTsFileNode
+ *     |       lock file -> empty? -> strategy -> migration check
+ *     |
+ *     +--> needDecodeTsFile?
+ *     |       |-- false -> LocalLoadStrategy
+ *     |       |            `-- FragmentInstance -> local region (no decode)
+ *     |       `-- true  -> TwoPhaseConsensusLoadStrategy
+ *     |                     phase1: TsFileSplitConsumer
+ *     |                       DataPartitionRouter -> MemoryBoundedBuffer
+ *     |                       -> PieceDispatcher
+ *     |                     phase2: BEGIN -> PIECE* -> PREPARE -> COMMIT
+ *     |                       or ABORT, via RegionConsensusContext +
+ *     |                       LoadConsensusSubmitter
+ *     |
+ *     +--> success -> register pending deletion (the source file is kept
+ *     |               until the whole LOAD batch ends) + log
+ *     |       failure -> record failed index
+ *     |
+ *     `--> all success -> FINISHED
+ *             else -> LoadFallbackHandler (convert to tablets, retry)
+ *                      -> FINISHED / FAILED
+ *
+ *     finally -> deletePendingDeletionFiles()
+ * }
+ * + *

Consensus pipeline (phase 1)

+ * + * 
{@code
+ * TsFileSplitter
+ *      |
+ *      | TsFileData (CHUNK / DELETION)
+ *      v
+ * TsFileSplitConsumer
+ *      |
+ *      +-- CHUNK:  buffer -> DataPartitionRouter -> per-region piece
+ *      |           over budget? -> PieceDispatcher: dispatch largest first
+ *      +-- end of file: flush remaining pieces
+ *      |
+ *      v
+ * PieceDispatcher
+ *      | dispatch callback
+ *      v
+ * TwoPhaseConsensusLoadStrategy.dispatchConsensusPiece
+ *      |
+ *      +-- first piece of a region: BEGIN(loadId) then PIECE(0)
+ *      +-- later pieces:            PIECE(1), PIECE(2), ...
+ *      |
+ *      v
+ * RegionConsensusContext.accumulate(bytes, checksum)
+ *      |
+ *      v
+ * LoadConsensusSubmitter (submit to the partition write node; bounded retry)
+ * }
+ * + *

Two-phase protocol timeline

+ * + * 
{@code
+ * coordinator                          region write peer
+ *      |                                      |
+ *      |---- BEGIN(loadId) ------------------>| create staged writer
+ *      |---- PIECE(0, chunks) --------------->| append chunks
+ *      |---- PIECE(1, chunks) --------------->| append chunks
+ *      |---- ...                              |
+ *      |---- PREPARE(count, bytes, checksum)->| seal staged TsFile
+ *      |---- COMMIT ------------------------->| load staged TsFile
+ *      |                                      |
+ *   on failure:
+ *      |---- ABORT -------------------------->| drop staged data
+ * }
+ * + *

Result handling

+ * + * Successful files are only registered for deferred deletion and logged (debug for pipe-generated + * loads, info otherwise). Failed indexes are collected; when all files are done the scheduler + * either transitions to FINISHED or hands the failures to {@link LoadFallbackHandler}, which + * converts the failed TsFiles into tablets, retries the insertion and finally transitions to + * FINISHED or FAILED. Only then - in the finally block - are the registered source TsFiles + * physically deleted, so a file with deleteAfterLoad is never removed before the whole LOAD batch + * has ended. + * + *

Component responsibilities

+ * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + *
Components of the LOAD pipeline
ComponentResponsibility
{@link DataPartitionBatchFetcher}LOAD partition fetcher: batches (device, time-partition) queries with the transmit limit + * and applies the table-model/pipe database hint
{@link DataPartitionRouter}LOAD chunk router: deduplicates (device, slot) pairs and maps every chunk to its target + * {@code TRegionReplicaSet}
{@link MemoryBoundedBuffer}LOAD memory budget: pure memory-pool accounting; emits the "over budget" signal that + * triggers eviction
{@link PieceDispatcher}LOAD piece dispatcher: buffered per-region pieces, largest-first eviction heap and + * flushing
{@link TsFileSplitConsumer}LOAD split consumer: the {@code TsFileDataConsumer} composing router, buffer and + * dispatcher into the route -> buffer -> dispatch pipeline
{@link RegionConsensusContext}LOAD per-region two-phase state: one context per region with load id, piece count, total + * bytes, XOR checksum and BEGIN state
{@link LoadConsensusSubmitter}LOAD consensus transport for BEGIN/PIECE/PREPARE/COMMIT/ABORT: resolves the partition's + * write node and submits via local consensus write or internal RPC, like the normal write + * path
{@link LoadTsFileDispatcherImpl}legacy LOAD local dispatcher (local-load path) and the per-file uuid holder for + * correlation
{@link LoadFallbackHandler}LOAD failure fallback: converts failed TsFiles into tablets and retries, then resolves + * the final state-machine result
*/ public class LoadTsFileScheduler implements IScheduler { private static final Logger LOGGER = LoggerFactory.getLogger(LoadTsFileScheduler.class); - private static final IoTDBConfig CONFIG = IoTDBDescriptor.getInstance().getConfig(); - - private static final LoadTsFileCostMetricsSet LOAD_TSFILE_COST_METRICS_SET = - LoadTsFileCostMetricsSet.getInstance(); - - private static final long SINGLE_SCHEDULER_MAX_MEMORY_SIZE = - IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize() >> 2; - private static final int TRANSMIT_LIMIT = - CommonDescriptor.getInstance().getConfig().getTTimePartitionSlotTransmitLimit(); - private static final Set LOADING_FILE_SET = new HashSet<>(); private final MPPQueryContext queryContext; @@ -142,10 +213,16 @@ public class LoadTsFileScheduler implements IScheduler { private final List tsFileNodeList; private final List failedTsFileNodeIndexes; private final PlanFragmentId fragmentId; - private final Set allReplicaSets; private final boolean isGeneratedByPipe; - private final Map timePartitionSlotToProgressIndex; private final LoadTsFileDataCacheMemoryBlock block; + private final LoadConsensusSubmitter consensusSubmitter; + + /** + * Source TsFiles whose COMMIT has already been sent and applied, but whose physical deletion must + * wait until the whole LOAD batch has finished (all nodes resolved or fallback done). Keyed by + * the absolute TsFile path so the same file is never registered twice. + */ + private final Map pendingDeletionFiles = new HashMap<>(); public LoadTsFileScheduler( DistributedQueryPlan distributedQueryPlan, @@ -161,10 +238,9 @@ public LoadTsFileScheduler( this.fragmentId = distributedQueryPlan.getRootSubPlan().getPlanFragment().getId(); this.dispatcher = new LoadTsFileDispatcherImpl(internalServiceClientManager, isGeneratedByPipe); this.partitionFetcher = new DataPartitionBatchFetcher(partitionFetcher); - this.allReplicaSets = new HashSet<>(); this.isGeneratedByPipe = isGeneratedByPipe; - this.timePartitionSlotToProgressIndex = new HashMap<>(); this.block = LoadTsFileMemoryManager.getInstance().allocateDataCacheMemoryBlock(); + this.consensusSubmitter = new LoadConsensusSubmitter(internalServiceClientManager); for (FragmentInstance fragmentInstance : distributedQueryPlan.getInstances()) { tsFileNodeList.add((LoadSingleTsFileNode) fragmentInstance.getFragment().getPlanNodeTree()); @@ -175,531 +251,166 @@ public LoadTsFileScheduler( public void start() { try { stateMachine.transitionToRunning(); - int tsFileNodeListSize = tsFileNodeList.size(); boolean isLoadSuccess = true; - for (int i = 0; i < tsFileNodeListSize; ++i) { + for (int i = 0; i < tsFileNodeList.size(); ++i) { final LoadSingleTsFileNode node = tsFileNodeList.get(i); final String filePath = node.getTsFileResource().getTsFilePath(); + final String userName = queryContext.getSession().getUserName(); partitionFetcher.setDatabase(getPartitionQueryDatabase(node, isGeneratedByPipe)); - boolean isLoadSingleTsFileSuccess = true; - boolean shouldRemoveFileFromLoadingSet = false; - try { - synchronized (LOADING_FILE_SET) { - if (LOADING_FILE_SET.contains(filePath)) { - throw new LoadFileException( - String.format( - DataNodeQueryMessages - .QUERY_EXCEPTION_TSFILE_S_IS_LOADING_BY_ANOTHER_SCHEDULER_55077B82, - filePath)); - } - LOADING_FILE_SET.add(filePath); - } - shouldRemoveFileFromLoadingSet = true; - - final long startTimeMs = System.currentTimeMillis(); - - if (node.isTsFileEmpty()) { - LOGGER.info(DataNodeQueryMessages.LOAD_SKIP_TSFILE_BECAUSE_IT_HAS_NO_DATA, filePath); - } else if (!node.needDecodeTsFile( - slotList -> - partitionFetcher.queryDataPartition( - slotList, queryContext.getSession().getUserName()))) { - // do not decode, load locally - final long startTime = System.nanoTime(); - try { - isLoadSingleTsFileSuccess = loadLocally(node); - } finally { - LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( - LoadTsFileCostMetricsSet.LOAD_LOCALLY, System.nanoTime() - startTime); - } - } else { - // need decode, load locally or remotely, use two phases method - String uuid = UUID.randomUUID().toString(); - dispatcher.setUuid(uuid); - allReplicaSets.clear(); - - long startTime = System.nanoTime(); - final boolean isFirstPhaseSuccess; - try { - isFirstPhaseSuccess = firstPhase(node); - } finally { - LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( - LoadTsFileCostMetricsSet.FIRST_PHASE, System.nanoTime() - startTime); - } - - startTime = System.nanoTime(); - final boolean isSecondPhaseSuccess; - try { - isSecondPhaseSuccess = - secondPhase(isFirstPhaseSuccess, uuid, node.getTsFileResource()); - } finally { - LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( - LoadTsFileCostMetricsSet.SECOND_PHASE, System.nanoTime() - startTime); - } - - if (!isFirstPhaseSuccess || !isSecondPhaseSuccess) { - isLoadSingleTsFileSuccess = false; - } - } - - if (RegionMigrateService.getInstance().getLastNotifyMigratingTime() > startTimeMs - || RegionMigrateService.getInstance().mayHaveMigratingRegions()) { - LOGGER.warn( - DataNodeQueryMessages - .LOADTSFILESCHEDULER_REGION_MIGRATION_WAS_DETECTED_DURING_LOADING_TSFILE_ARG_WILL_CONVERT, - filePath); - isLoadSingleTsFileSuccess = false; - } - - if (isLoadSingleTsFileSuccess) { - node.clean(); - if (isGeneratedByPipe) { - LOGGER.debug( - DataNodeQueryMessages.LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG, - filePath, - i + 1, - tsFileNodeListSize); - } else { - LOGGER.info( - DataNodeQueryMessages.LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG, - filePath, - i + 1, - tsFileNodeListSize); - } - } else { - isLoadSuccess = false; - failedTsFileNodeIndexes.add(i); - LOGGER.warn( - DataNodeQueryMessages.CAN_NOT_LOAD_TSFILE_ARG_LOAD_PROCESS_ARG_ARG, - filePath, - i + 1, - tsFileNodeListSize); - } - } catch (Exception e) { + if (!processSingleNode(node, i, tsFileNodeList.size(), userName)) { isLoadSuccess = false; failedTsFileNodeIndexes.add(i); - LOGGER.warn(DataNodeQueryMessages.LOADTSFILESCHEDULER_LOADS_TSFILE_ERROR, filePath, e); - } finally { - if (shouldRemoveFileFromLoadingSet) { - synchronized (LOADING_FILE_SET) { - LOADING_FILE_SET.remove(filePath); - } - } + continue; + } + + // The COMMIT marker has been sent (or the file was loaded locally); the source file is no + // longer needed by the write path, but it must survive until the whole LOAD batch ends - + // record it in memory and delete it together in the finally block. + registerPendingDeletion(node); + if (isGeneratedByPipe) { + LOGGER.debug( + DataNodeQueryMessages.LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG, + filePath, + i + 1, + tsFileNodeList.size()); + } else { + LOGGER.info( + DataNodeQueryMessages.LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG, + filePath, + i + 1, + tsFileNodeList.size()); } } if (isLoadSuccess) { stateMachine.transitionToFinished(); } else { - final StringBuilder failedTsFiles = - new StringBuilder( - !tsFileNodeList.isEmpty() - ? tsFileNodeList - .get(failedTsFileNodeIndexes.get(0)) - .getTsFileResource() - .getTsFilePath() - : ""); - final ListIterator iterator = failedTsFileNodeIndexes.listIterator(1); - while (iterator.hasNext()) { - failedTsFiles - .append(", ") - .append(tsFileNodeList.get(iterator.next()).getTsFileResource().getTsFilePath()); - } - final long startTime = System.nanoTime(); - try { - // if failed to load some TsFiles, then try to convert the TsFiles to Tablets - LOGGER.info( - DataNodeQueryMessages - .LOAD_TSFILE_S_FAILED_WILL_TRY_TO_CONVERT_TO_TABLETS_AND_INSERT_FAILED_TSFILES_ARG, - failedTsFiles); - convertFailedTsFilesToTabletsAndRetry(); - } finally { - LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( - LoadTsFileCostMetricsSet.SCHEDULER_CAST_TABLETS, System.nanoTime() - startTime); - } + new LoadFallbackHandler( + queryContext, + isGeneratedByPipe, + tsFileNodeList, + failedTsFileNodeIndexes, + stateMachine) + .convertFailedTsFilesToTablets(); } } finally { + deletePendingDeletionFiles(); dispatcher.close(); LoadTsFileMemoryManager.getInstance().releaseDataCacheMemoryBlock(); } } - private boolean firstPhase(LoadSingleTsFileNode node) { - final TsFileDataManager tsFileDataManager = new TsFileDataManager(this, node, block); - try { - new TsFileSplitter( - node.getTsFileResource().getTsFile(), tsFileDataManager::addOrSendTsFileData) - .splitTsFileByDataPartition(); - if (!tsFileDataManager.sendAllTsFileData()) { - return false; - } - } catch (IllegalStateException e) { - LOGGER.warn( - String.format( - DataNodeQueryMessages.DISPATCH_TSFILEDATA_ERROR_WHEN_PARSING_TSFILE_S, - node.getTsFileResource().getTsFile()), - e); - return false; - } catch (Exception e) { - LOGGER.warn( - String.format( - DataNodeQueryMessages.PARSE_OR_SEND_TSFILE_S_ERROR, - node.getTsFileResource().getTsFile()), - e); - return false; - } finally { - tsFileDataManager.clear(); + /** + * Records a successfully loaded TsFile for deferred deletion. Only files with {@code + * deleteAfterLoad} are tracked; they are physically removed by {@link + * #deletePendingDeletionFiles()} after the whole LOAD batch finishes. + */ + private void registerPendingDeletion(final LoadSingleTsFileNode node) { + if (node.isDeleteAfterLoad()) { + pendingDeletionFiles.put(node.getTsFileResource().getTsFilePath(), node); } - return true; } - static boolean isSameRegionReplicaSet(TRegionReplicaSet original, TRegionReplicaSet current) { - if (!Objects.equals(original.getRegionId(), current.getRegionId())) { - return false; - } - - final Map locationCounts = new HashMap<>(); - for (TDataNodeLocation location : original.getDataNodeLocations()) { - locationCounts.merge(location, 1, Integer::sum); + /** + * Deletes every source TsFile whose load has finished (COMMIT sent and applied). This runs when + * the whole LOAD batch has ended - after the state machine resolved to FINISHED/FAILED - instead + * of right after each individual COMMIT, so a source file is never removed while the batch may + * still need it. + */ + private void deletePendingDeletionFiles() { + if (pendingDeletionFiles.isEmpty()) { + return; } - for (TDataNodeLocation location : current.getDataNodeLocations()) { - final Integer count = locationCounts.get(location); - if (count == null) { - return false; - } - if (count == 1) { - locationCounts.remove(location); - } else { - locationCounts.put(location, count - 1); - } + LOGGER.info( + DataNodeQueryMessages + .LOG_LOAD_BATCH_FINISHED_DELETING_ARG_SOURCE_TSFILES_AFTER_LOAD_D5EE56E9, + pendingDeletionFiles.size()); + for (final LoadSingleTsFileNode node : pendingDeletionFiles.values()) { + node.clean(); } - return locationCounts.isEmpty(); + pendingDeletionFiles.clear(); } - private boolean dispatchOnePieceNode( - LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { - allReplicaSets.add(replicaSet); - FragmentInstance instance = - new FragmentInstance( - new PlanFragment(fragmentId, pieceNode), - fragmentId.genFragmentInstanceId(), - null, - queryContext.getQueryType(), - queryContext.getTimeOut() - (System.currentTimeMillis() - queryContext.getStartTime()), - queryContext.getSession(), - queryContext.isDebug(), - queryContext.isVerbose()); - instance.setExecutorAndHost(new StorageExecutor(replicaSet)); - Future dispatchResultFuture = - dispatcher.dispatch(null, Collections.singletonList(instance)); - + /** + * Loads one TsFile: guard clauses first (concurrent-load lock, empty file), then route to the + * strategy and finally detect whether a region migration raced with the load. + */ + private boolean processSingleNode( + LoadSingleTsFileNode node, int index, int listSize, String userName) { + final String filePath = node.getTsFileResource().getTsFilePath(); + final long startTimeMs = System.currentTimeMillis(); + boolean shouldRemoveFileFromLoadingSet = false; try { - FragInstanceDispatchResult result = - dispatchResultFuture.get( - CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds(), TimeUnit.SECONDS); - if (!result.isSuccessful()) { - LOGGER.warn( - DataNodeQueryMessages.DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG - + DataNodeQueryMessages - .RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG, - replicaSet, - TSStatusCode.representOf(result.getFailureStatus().getCode()).name(), - result.getFailureStatus().getMessage(), - pieceNode); - if (result.getFailureStatus().getSubStatus() != null) { - for (TSStatus status : result.getFailureStatus().getSubStatus()) { - LOGGER.warn( - DataNodeQueryMessages.SUB_STATUS_CODE_ARG_SUB_STATUS_MESSAGE_ARG, - TSStatusCode.representOf(status.getCode()).toString(), - status.getMessage()); - } + synchronized (LOADING_FILE_SET) { + if (LOADING_FILE_SET.contains(filePath)) { + throw new LoadFileException( + String.format( + DataNodeQueryMessages + .QUERY_EXCEPTION_TSFILE_S_IS_LOADING_BY_ANOTHER_SCHEDULER_55077B82, + filePath)); } - TSStatus status = result.getFailureStatus(); - status.setMessage( - String.format( - DataNodeQueryMessages.MESSAGE_LOAD_ARG_PIECE_ERROR_1ST_PHASE_BECAUSE_F3D9672C, - pieceNode.getTsFile()) - + status.getMessage()); - return false; - } - } catch (InterruptedException | ExecutionException | CancellationException e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); + LOADING_FILE_SET.add(filePath); } - LOGGER.warn(DataNodeQueryMessages.INTERRUPT_OR_EXECUTION_ERROR, e); - return false; - } catch (TimeoutException e) { - dispatchResultFuture.cancel(true); - LOGGER.warn( - String.format( - DataNodeQueryMessages.WAIT_FOR_LOADING_S_TIME_OUT, - LoadTsFilePieceNode.class.getName()), - e); - return false; - } - return true; - } + shouldRemoveFileFromLoadingSet = true; - private boolean secondPhase( - boolean isFirstPhaseSuccess, String uuid, TsFileResource tsFileResource) { - if (isGeneratedByPipe) { - LOGGER.debug(DataNodeQueryMessages.START_DISPATCHING_LOAD_COMMAND_FOR_UUID, uuid); - } else { - LOGGER.info(DataNodeQueryMessages.START_DISPATCHING_LOAD_COMMAND_FOR_UUID, uuid); - } - final File tsFile = tsFileResource.getTsFile(); - final TLoadCommandReq loadCommandReq = - new TLoadCommandReq( - (isFirstPhaseSuccess ? LoadCommand.EXECUTE : LoadCommand.ROLLBACK).ordinal(), uuid); + if (node.isTsFileEmpty()) { + LOGGER.info(DataNodeQueryMessages.LOAD_SKIP_TSFILE_BECAUSE_IT_HAS_NO_DATA, filePath); + return true; + } - try { - loadCommandReq.setIsGeneratedByPipe(isGeneratedByPipe); - loadCommandReq.setTimePartition2ProgressIndex( - timePartitionSlotToProgressIndex.entrySet().stream() - .collect( - Collectors.toMap( - Map.Entry::getKey, - entry -> { - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); - final DataOutputStream dataOutputStream = - new DataOutputStream(byteArrayOutputStream)) { - entry.getValue().serialize(dataOutputStream); - return ByteBuffer.wrap( - byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); - } catch (final IOException e) { - throw new RuntimeException( - String.format( - DataNodeQueryMessages - .QUERY_EXCEPTION_SERIALIZE_PROGRESS_INDEX_ERROR_ISFIRSTPHASESUCCESS_S_UUID_690F0419, - isFirstPhaseSuccess, - uuid, - tsFile.getAbsolutePath()), - e); - } - }))); - Future dispatchResultFuture = - dispatcher.dispatchCommand(loadCommandReq, allReplicaSets); + final TsFileLoadStrategy strategy; + if (!node.needDecodeTsFile( + slotList -> partitionFetcher.queryDataPartition(slotList, userName))) { + // do not decode, load locally + strategy = new LocalLoadStrategy(queryContext, fragmentId, dispatcher); + } else { + // need decode, use the consensus two-phase pipeline + strategy = + new TwoPhaseConsensusLoadStrategy( + dispatcher, + partitionFetcher, + block, + consensusSubmitter, + userName, + isGeneratedByPipe); + } + final boolean isLoadSingleTsFileSuccess = strategy.execute(node); - FragInstanceDispatchResult result = dispatchResultFuture.get(); - if (!result.isSuccessful()) { + if (RegionMigrateService.getInstance().getLastNotifyMigratingTime() > startTimeMs + || RegionMigrateService.getInstance().mayHaveMigratingRegions()) { LOGGER.warn( DataNodeQueryMessages - .DISPATCH_LOAD_COMMAND_ARG_OF_TSFILE_ARG_ERROR_TO_REPLICASETS_ARG_ERROR - + DataNodeQueryMessages.RESULT_STATUS_CODE_ARG_RESULT_STATUS_MESSAGE_ARG, - loadCommandReq, - tsFile, - allReplicaSets, - TSStatusCode.representOf(result.getFailureStatus().getCode()).name(), - result.getFailureStatus().getMessage()); - TSStatus status = result.getFailureStatus(); - status.setMessage( - String.format( - DataNodeQueryMessages - .MESSAGE_LOAD_ARG_ERROR_SECOND_PHASE_BECAUSE_ARG_FIRST_PHASE_ARG_CBA980FC, - tsFile, - status.getMessage(), - isFirstPhaseSuccess - ? DataNodeQueryMessages.MESSAGE_SUCCESS_260CA9DD - : DataNodeQueryMessages.MESSAGE_FAILED_26934EB3)); + .LOADTSFILESCHEDULER_REGION_MIGRATION_WAS_DETECTED_DURING_LOADING_TSFILE_ARG_WILL_CONVERT, + filePath); + logCannotLoad(node, index, listSize); return false; } - } catch (InterruptedException | ExecutionException e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); + if (!isLoadSingleTsFileSuccess) { + logCannotLoad(node, index, listSize); + return false; } - LOGGER.warn(DataNodeQueryMessages.INTERRUPT_OR_EXECUTION_ERROR, e); - return false; + return true; } catch (Exception e) { - LOGGER.warn( - DataNodeQueryMessages.EXCEPTION_OCCURRED_DURING_SECOND_PHASE_OF_LOADING_TSFILE, - tsFile, - e); + LOGGER.warn(DataNodeQueryMessages.LOADTSFILESCHEDULER_LOADS_TSFILE_ERROR, filePath, e); return false; - } - return true; - } - - private ByteBuffer assignProgressIndex(TsFileResource tsFileResource) throws IOException { - PipeDataNodeAgent.runtime().assignProgressIndexForTsFileLoad(tsFileResource); - - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); - final DataOutputStream dataOutputStream = new DataOutputStream(byteArrayOutputStream)) { - tsFileResource.getMaxProgressIndex().serialize(dataOutputStream); - return ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); - } - } - - private boolean loadLocally(LoadSingleTsFileNode node) throws IoTDBException { - LOGGER.info( - DataNodeQueryMessages.START_LOAD_TSFILE_LOCALLY, - node.getTsFileResource().getTsFile().getPath()); - - if (CommonDescriptor.getInstance().getConfig().isReadOnly()) { - throw new LoadReadOnlyException(); - } - - // if the time index is PlainDeviceTimeIndex, convert it to ArrayDeviceTimeIndex - if (node.getTsFileResource().getTimeIndex() instanceof PlainDeviceTimeIndex) { - final PlainDeviceTimeIndex timeIndex = - (PlainDeviceTimeIndex) node.getTsFileResource().getTimeIndex(); - final Map convertedDeviceToIndex = new ConcurrentHashMap<>(); - for (final Map.Entry entry : timeIndex.getDeviceToIndex().entrySet()) { - convertedDeviceToIndex.put( - entry.getKey() instanceof StringArrayDeviceID - ? entry.getKey() - : new StringArrayDeviceID(entry.getKey().toString()), - entry.getValue()); - } - node.getTsFileResource() - .setTimeIndex( - new ArrayDeviceTimeIndex( - convertedDeviceToIndex, timeIndex.getStartTimes(), timeIndex.getEndTimes())); - } - - try { - FragmentInstance instance = - new FragmentInstance( - new PlanFragment(fragmentId, node), - fragmentId.genFragmentInstanceId(), - null, - queryContext.getQueryType(), - queryContext.getTimeOut() - - (System.currentTimeMillis() - queryContext.getStartTime()), - queryContext.getSession(), - queryContext.isDebug(), - queryContext.isVerbose()); - instance.setExecutorAndHost(new StorageExecutor(node.getLocalRegionReplicaSet())); - dispatcher.dispatchLocally(instance); - } catch (FragmentInstanceDispatchException e) { - LOGGER.warn( - String.format( - DataNodeQueryMessages.DISPATCH_TSFILE_S_ERROR_TO_LOCAL_ERROR_RESULT_STATUS_CODE_S - + DataNodeQueryMessages.RESULT_STATUS_MESSAGE_S, - node.getTsFileResource().getTsFile(), - TSStatusCode.representOf(e.getFailureStatus().getCode()).name(), - e.getFailureStatus().getMessage())); - return false; - } - - // add metrics - Optional.ofNullable( - StorageEngine.getInstance() - .getDataRegion( - (DataRegionId) - ConsensusGroupId.Factory.createFromTConsensusGroupId( - node.getLocalRegionReplicaSet().getRegionId()))) - .ifPresent( - dataRegion -> - dataRegion - .getNonSystemDatabaseName() - .ifPresent( - databaseName -> { - // Report load tsFile points to IoTDB flush metrics - MemTableFlushTask.recordFlushPointsMetricInternal( - node.getWritePointCount(), - databaseName, - dataRegion.getDataRegionIdString()); - - MetricService.getInstance() - .count( - node.getWritePointCount(), - Metric.QUANTITY.toString(), - MetricLevel.CORE, - Tag.NAME.toString(), - Metric.POINTS_IN.toString(), - Tag.DATABASE.toString(), - databaseName, - Tag.REGION.toString(), - dataRegion.getDataRegionIdString(), - Tag.TYPE.toString(), - Metric.LOAD_POINT_COUNT.toString()); - MetricService.getInstance() - .count( - node.getWritePointCount(), - Metric.LEADER_QUANTITY.toString(), - MetricLevel.CORE, - Tag.NAME.toString(), - Metric.POINTS_IN.toString(), - Tag.DATABASE.toString(), - databaseName, - Tag.REGION.toString(), - dataRegion.getDataRegionIdString(), - Tag.TYPE.toString(), - Metric.LOAD_POINT_COUNT.toString()); - })); - - return true; - } - - private void convertFailedTsFilesToTabletsAndRetry() { - final LoadTsFileDataTypeConverter loadTsFileDataTypeConverter = - new LoadTsFileDataTypeConverter(queryContext, isGeneratedByPipe); - - final Iterator iterator = failedTsFileNodeIndexes.listIterator(); - while (iterator.hasNext()) { - final int failedLoadTsFileIndex = iterator.next(); - final LoadSingleTsFileNode failedNode = tsFileNodeList.get(failedLoadTsFileIndex); - final String filePath = failedNode.getTsFileResource().getTsFilePath(); - - try { - final TSStatus status = - failedNode.isTableModel() - ? loadTsFileDataTypeConverter - .convertForTableModel( - (isGeneratedByPipe - ? LoadTsFile.createForPipe(null, filePath, Collections.emptyMap()) - : LoadTsFile.createUnchecked( - null, filePath, Collections.emptyMap())) - .setDatabase(failedNode.getDatabase()) - .setDeleteAfterLoad(failedNode.isDeleteAfterLoad()) - .setConvertOnTypeMismatch(true)) - .orElse(null) - : loadTsFileDataTypeConverter - .convertForTreeModel( - buildRetryTreeLoadStatement( - filePath, - failedNode.isDeleteAfterLoad(), - getPartitionQueryDatabase(failedNode, isGeneratedByPipe))) - .orElse(null); - - if (loadTsFileDataTypeConverter.isSuccessful(status)) { - iterator.remove(); - LOGGER.info( - DataNodeQueryMessages - .LOAD_SUCCESSFULLY_CONVERTED_TSFILE_ARG_INTO_TABLETS_AND_INSERTED, - failedNode.getTsFileResource().getTsFilePath()); - } else { - LOGGER.warn( - DataNodeQueryMessages.LOAD_FAILED_TO_CONVERT_TO_TABLETS_FROM_TSFILE_ARG_STATUS_ARG, - failedNode.getTsFileResource().getTsFilePath(), - status); + } finally { + if (shouldRemoveFileFromLoadingSet) { + synchronized (LOADING_FILE_SET) { + LOADING_FILE_SET.remove(filePath); } - } catch (final Exception e) { - LOGGER.warn( - DataNodeQueryMessages.LOAD_FAILED_TO_CONVERT_TO_TABLETS_FROM_TSFILE_ARG_EXCEPTION_ARG, - failedNode.getTsFileResource().getTsFilePath(), - e.getMessage(), - e); } } + } - // If all failed TsFiles are converted into tablets and inserted, - // we can consider the load process as successful. - if (failedTsFileNodeIndexes.isEmpty()) { - LOGGER.info(DataNodeQueryMessages.LOAD_ALL_FAILED_TSFILES_ARE_CONVERTED_TO_TABLETS); - stateMachine.transitionToFinished(); - } else { - final String errorMsg = - "Load: failed to load some TsFiles by converting them into tablets. Failed TsFiles: " - + failedTsFileNodeIndexes.stream() - .map(i -> tsFileNodeList.get(i).getTsFileResource().getTsFilePath()) - .collect(Collectors.joining(", ")); - LOGGER.warn(errorMsg); - stateMachine.transitionToFailed(new LoadFileException(errorMsg)); - } + private void logCannotLoad(LoadSingleTsFileNode node, int index, int listSize) { + LOGGER.warn( + DataNodeQueryMessages.CAN_NOT_LOAD_TSFILE_ARG_LOAD_PROCESS_ARG_ARG, + node.getTsFileResource().getTsFilePath(), + index + 1, + listSize); } static String getPartitionQueryDatabase( @@ -707,25 +418,6 @@ static String getPartitionQueryDatabase( return node.isTableModel() || isGeneratedByPipe ? node.getDatabase() : null; } - private LoadTsFileStatement buildRetryTreeLoadStatement( - final String filePath, final boolean deleteAfterLoad, final String database) - throws FileNotFoundException { - final LoadTsFileStatement statement = - (isGeneratedByPipe - ? LoadTsFileStatement.createForPipe(filePath) - : LoadTsFileStatement.createUnchecked(filePath)) - .setDeleteAfterLoad(deleteAfterLoad) - .setConvertOnTypeMismatch(true); - if (database != null) { - statement.setDatabase(database); - statement.updateDatabaseLevelByTreeDatabase(); - } - if (isGeneratedByPipe) { - statement.markIsGeneratedByPipe(); - } - return statement; - } - @Override public void stop(Throwable t) { dispatcher.abort(); @@ -741,266 +433,8 @@ public FragmentInfo getFragmentInfo() { return null; } - private void computeTimePartitionSlotToProgressIndexIfAbsent( - final TTimePartitionSlot timePartitionSlot) { - timePartitionSlotToProgressIndex.putIfAbsent( - timePartitionSlot, PipeDataNodeAgent.runtime().getNextProgressIndexForTsFileLoad()); - } - public enum LoadCommand { EXECUTE, ROLLBACK } - - private static class TsFileDataManager { - private final LoadTsFileScheduler scheduler; - private final LoadSingleTsFileNode singleTsFileNode; - - private long dataSize; - private final Map> - regionId2ReplicaSetAndNode; - private final List nonDirectionalChunkData; - private final LoadTsFileDataCacheMemoryBlock block; - - public TsFileDataManager( - LoadTsFileScheduler scheduler, - LoadSingleTsFileNode singleTsFileNode, - LoadTsFileDataCacheMemoryBlock block) { - this.scheduler = scheduler; - this.singleTsFileNode = singleTsFileNode; - this.dataSize = 0; - this.regionId2ReplicaSetAndNode = new HashMap<>(); - this.nonDirectionalChunkData = new ArrayList<>(); - this.block = block; - } - - private boolean addOrSendTsFileData(TsFileData tsFileData) throws LoadFileException { - switch (tsFileData.getType()) { - case CHUNK: - return addOrSendChunkData((ChunkData) tsFileData); - case DELETION: - return addOrSendDeletionData((DeletionData) tsFileData); - default: - throw new UnsupportedOperationException( - String.format( - DataNodeQueryMessages.QUERY_EXCEPTION_UNSUPPORTED_TSFILEDATATYPE_S_374475FA, - tsFileData.getType())); - } - } - - private boolean isMemoryEnough() { - return dataSize <= SINGLE_SCHEDULER_MAX_MEMORY_SIZE && block.hasEnoughMemory(); - } - - private boolean addOrSendChunkData(ChunkData chunkData) throws LoadFileException { - nonDirectionalChunkData.add(chunkData); - dataSize += chunkData.getDataSize(); - block.addMemoryUsage(chunkData.getDataSize()); - scheduler.computeTimePartitionSlotToProgressIndexIfAbsent(chunkData.getTimePartitionSlot()); - - if (!isMemoryEnough()) { - routeChunkData(); - - // start to dispatch from the biggest TsFilePieceNode - List sortedRegionIds = - regionId2ReplicaSetAndNode.keySet().stream() - .sorted( - Comparator.comparingLong( - o -> regionId2ReplicaSetAndNode.get(o).getRight().getDataSize()) - .reversed()) - .collect(Collectors.toList()); - - for (TConsensusGroupId sortedRegionId : sortedRegionIds) { - final TRegionReplicaSet replicaSet = - regionId2ReplicaSetAndNode.get(sortedRegionId).getLeft(); - final LoadTsFilePieceNode pieceNode = - regionId2ReplicaSetAndNode.get(sortedRegionId).getRight(); - if (pieceNode.getDataSize() == 0) { // total data size has been reduced to 0 - break; - } - final boolean isDispatchSuccess = scheduler.dispatchOnePieceNode(pieceNode, replicaSet); - - regionId2ReplicaSetAndNode.replace( - sortedRegionId, - new Pair<>( - replicaSet, - new LoadTsFilePieceNode( - singleTsFileNode.getPlanNodeId(), - singleTsFileNode - .getTsFileResource() - .getTsFile()))); // can not just remove, because of deletion - releaseMemoryUsage(pieceNode.getDataSize()); - - if (!isDispatchSuccess) { - // Currently there is no retry, so return directly - return false; - } - - if (isMemoryEnough()) { - break; - } - } - } - - return true; - } - - private void routeChunkData() throws LoadFileException { - if (nonDirectionalChunkData.isEmpty()) { - return; - } - - final List> partitionSlotList = new ArrayList<>(); - final int[] chunkPartitionIndexes = new int[nonDirectionalChunkData.size()]; - final Map> partitionSlotIndexes = new HashMap<>(); - for (int i = 0, size = nonDirectionalChunkData.size(); i < size; i++) { - final ChunkData chunkData = nonDirectionalChunkData.get(i); - final IDeviceID device = chunkData.getDevice(); - final TTimePartitionSlot timePartitionSlot = chunkData.getTimePartitionSlot(); - final Map slotIndexes = - partitionSlotIndexes.computeIfAbsent(device, key -> new HashMap<>()); - Integer partitionSlotIndex = slotIndexes.get(timePartitionSlot); - if (partitionSlotIndex == null) { - partitionSlotIndex = partitionSlotList.size(); - slotIndexes.put(timePartitionSlot, partitionSlotIndex); - partitionSlotList.add(new Pair<>(device, timePartitionSlot)); - } - chunkPartitionIndexes[i] = partitionSlotIndex; - } - - List replicaSets = - scheduler.partitionFetcher.queryDataPartition( - partitionSlotList, scheduler.queryContext.getSession().getUserName()); - for (int i = 0, size = nonDirectionalChunkData.size(); i < size; i++) { - final TRegionReplicaSet replicaSet = replicaSets.get(chunkPartitionIndexes[i]); - final TConsensusGroupId regionId = replicaSet.getRegionId(); - if (regionId2ReplicaSetAndNode.containsKey(regionId) - && !isSameRegionReplicaSet( - regionId2ReplicaSetAndNode.get(regionId).getLeft(), replicaSet)) { - // Detected region replica set changed (maybe due to region migration), throw an exception - throw new RegionReplicaSetChangedException( - regionId2ReplicaSetAndNode.get(regionId).getLeft(), replicaSet); - } - - regionId2ReplicaSetAndNode - .computeIfAbsent( - replicaSet.getRegionId(), - o -> - new Pair<>( - replicaSet, - new LoadTsFilePieceNode( - singleTsFileNode.getPlanNodeId(), - singleTsFileNode.getTsFileResource().getTsFile()))) - .getRight() - .addTsFileData(nonDirectionalChunkData.get(i)); - } - nonDirectionalChunkData.clear(); - } - - private boolean addOrSendDeletionData(DeletionData deletionData) throws LoadFileException { - routeChunkData(); // ensure chunk data will be added before deletion - - for (Map.Entry> entry : - regionId2ReplicaSetAndNode.entrySet()) { - dataSize += deletionData.getDataSize(); - block.addMemoryUsage(deletionData.getDataSize()); - entry.getValue().getRight().addTsFileData(deletionData); - } - return true; - } - - private boolean sendAllTsFileData() throws LoadFileException { - routeChunkData(); - - boolean isAllSuccess = true; - for (Map.Entry> entry : - regionId2ReplicaSetAndNode.entrySet()) { - releaseMemoryUsage(entry.getValue().getRight().getDataSize()); - if (isAllSuccess - && !scheduler.dispatchOnePieceNode( - entry.getValue().getRight(), entry.getValue().getLeft())) { - LOGGER.warn( - DataNodeQueryMessages.DISPATCH_PIECE_NODE_ARG_OF_TSFILE_ARG_ERROR, - entry.getValue(), - singleTsFileNode.getTsFileResource().getTsFile()); - isAllSuccess = false; - } - } - return isAllSuccess; - } - - private void releaseMemoryUsage(final long memorySize) { - dataSize -= memorySize; - block.reduceMemoryUsage(memorySize); - } - - private void clear() { - if (dataSize > 0) { - block.reduceMemoryUsage(dataSize); - dataSize = 0; - } - nonDirectionalChunkData.clear(); - regionId2ReplicaSetAndNode.clear(); - } - } - - private static class DataPartitionBatchFetcher { - private final IPartitionFetcher fetcher; - private String database; - - public DataPartitionBatchFetcher(IPartitionFetcher fetcher) { - this.fetcher = fetcher; - } - - public void setDatabase(String database) { - this.database = database; - } - - public List queryDataPartition( - List> slotList, String userName) { - List replicaSets = new ArrayList<>(slotList.size()); - int size = slotList.size(); - - for (int i = 0; i < size; i += TRANSMIT_LIMIT) { - List> subSlotList = - slotList.subList(i, Math.min(size, i + TRANSMIT_LIMIT)); - DataPartition dataPartition = - fetcher.getOrCreateDataPartition(toQueryParam(subSlotList), userName); - for (final Pair pair : subSlotList) { - // database is an explicit database hint for table-model loads and - // pipe-generated tree-model loads. - replicaSets.add( - database != null - ? dataPartition.getDataRegionReplicaSetForWriting(pair.left, pair.right, database) - : dataPartition.getDataRegionReplicaSetForWriting(pair.left, pair.right)); - } - } - return replicaSets; - } - - private List toQueryParam( - List> slots) { - final Map> device2TimePartitionSlots = new HashMap<>(); - for (final Pair slot : slots) { - device2TimePartitionSlots - .computeIfAbsent(slot.left, key -> new HashSet<>()) - .add(slot.right); - } - - final List queryParams = - new ArrayList<>(device2TimePartitionSlots.size()); - for (final Map.Entry> entry : - device2TimePartitionSlots.entrySet()) { - final DataPartitionQueryParam queryParam = - new DataPartitionQueryParam(entry.getKey(), new ArrayList<>(entry.getValue())); - // database is an explicit database hint for table-model loads and - // pipe-generated tree-model loads. - if (database != null) { - queryParam.setDatabaseName(database); - } - queryParams.add(queryParam); - } - return queryParams; - } - } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LocalLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LocalLoadStrategy.java new file mode 100644 index 0000000000000..0be512ff696a3 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LocalLoadStrategy.java @@ -0,0 +1,203 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.consensus.ConsensusGroupId; +import org.apache.iotdb.commons.consensus.DataRegionId; +import org.apache.iotdb.commons.exception.IoTDBException; +import org.apache.iotdb.commons.partition.StorageExecutor; +import org.apache.iotdb.commons.service.metric.MetricService; +import org.apache.iotdb.commons.service.metric.enums.Metric; +import org.apache.iotdb.commons.service.metric.enums.Tag; +import org.apache.iotdb.db.exception.load.LoadReadOnlyException; +import org.apache.iotdb.db.exception.mpp.FragmentInstanceDispatchException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.storageengine.StorageEngine; +import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.ArrayDeviceTimeIndex; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.PlainDeviceTimeIndex; +import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; +import org.apache.iotdb.metrics.utils.MetricLevel; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; + +/** + * LOAD local-load strategy: loads a TsFile without decoding it, used when the file's device/time + * ranges map to a single local region. The whole file is wrapped into one {@code FragmentInstance} + * (reusing the scheduler's fragment id) and dispatched to the local data region through {@link + * LoadTsFileDispatcherImpl#dispatchLocally}; nothing crosses the network. + * + *

The strategy also: + * + *

    + *
  • rejects loads while the node is read-only ({@link LoadReadOnlyException}); + *
  • converts a {@code PlainDeviceTimeIndex} into an {@code ArrayDeviceTimeIndex} so the local + * writer can use the time index directly; + *
  • records flush/points metrics on the target data region; + *
  • measures the whole execution as the {@code LOAD_LOCALLY} phase metric. + *
+ */ +public class LocalLoadStrategy implements TsFileLoadStrategy { + + private static final Logger LOGGER = LoggerFactory.getLogger(LocalLoadStrategy.class); + + private static final LoadTsFileCostMetricsSet LOAD_TSFILE_COST_METRICS_SET = + LoadTsFileCostMetricsSet.getInstance(); + + private final MPPQueryContext queryContext; + private final PlanFragmentId fragmentId; + private final LoadTsFileDispatcherImpl dispatcher; + + public LocalLoadStrategy( + MPPQueryContext queryContext, + PlanFragmentId fragmentId, + LoadTsFileDispatcherImpl dispatcher) { + this.queryContext = queryContext; + this.fragmentId = fragmentId; + this.dispatcher = dispatcher; + } + + @Override + public boolean execute(LoadSingleTsFileNode node) throws IoTDBException { + final long startTime = System.nanoTime(); + try { + return loadLocally(node); + } finally { + LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.LOAD_LOCALLY, System.nanoTime() - startTime); + } + } + + private boolean loadLocally(LoadSingleTsFileNode node) throws IoTDBException { + LOGGER.info( + DataNodeQueryMessages.START_LOAD_TSFILE_LOCALLY, + node.getTsFileResource().getTsFile().getPath()); + + if (CommonDescriptor.getInstance().getConfig().isReadOnly()) { + throw new LoadReadOnlyException(); + } + + // if the time index is PlainDeviceTimeIndex, convert it to ArrayDeviceTimeIndex + if (node.getTsFileResource().getTimeIndex() instanceof PlainDeviceTimeIndex) { + final PlainDeviceTimeIndex timeIndex = + (PlainDeviceTimeIndex) node.getTsFileResource().getTimeIndex(); + final Map convertedDeviceToIndex = new ConcurrentHashMap<>(); + for (final Map.Entry entry : timeIndex.getDeviceToIndex().entrySet()) { + convertedDeviceToIndex.put( + entry.getKey() instanceof StringArrayDeviceID + ? entry.getKey() + : new StringArrayDeviceID(entry.getKey().toString()), + entry.getValue()); + } + node.getTsFileResource() + .setTimeIndex( + new ArrayDeviceTimeIndex( + convertedDeviceToIndex, timeIndex.getStartTimes(), timeIndex.getEndTimes())); + } + + try { + FragmentInstance instance = + new FragmentInstance( + new PlanFragment(fragmentId, node), + fragmentId.genFragmentInstanceId(), + null, + queryContext.getQueryType(), + queryContext.getTimeOut() + - (System.currentTimeMillis() - queryContext.getStartTime()), + queryContext.getSession(), + queryContext.isDebug(), + queryContext.isVerbose()); + instance.setExecutorAndHost(new StorageExecutor(node.getLocalRegionReplicaSet())); + dispatcher.dispatchLocally(instance); + } catch (FragmentInstanceDispatchException e) { + LOGGER.warn( + String.format( + DataNodeQueryMessages.DISPATCH_TSFILE_S_ERROR_TO_LOCAL_ERROR_RESULT_STATUS_CODE_S + + DataNodeQueryMessages.RESULT_STATUS_MESSAGE_S, + node.getTsFileResource().getTsFile(), + TSStatusCode.representOf(e.getFailureStatus().getCode()).name(), + e.getFailureStatus().getMessage())); + return false; + } + + // add metrics + Optional.ofNullable( + StorageEngine.getInstance() + .getDataRegion( + (DataRegionId) + ConsensusGroupId.Factory.createFromTConsensusGroupId( + node.getLocalRegionReplicaSet().getRegionId()))) + .ifPresent( + dataRegion -> + dataRegion + .getNonSystemDatabaseName() + .ifPresent( + databaseName -> { + // Report load tsFile points to IoTDB flush metrics + MemTableFlushTask.recordFlushPointsMetricInternal( + node.getWritePointCount(), + databaseName, + dataRegion.getDataRegionIdString()); + + MetricService.getInstance() + .count( + node.getWritePointCount(), + Metric.QUANTITY.toString(), + MetricLevel.CORE, + Tag.NAME.toString(), + Metric.POINTS_IN.toString(), + Tag.DATABASE.toString(), + databaseName, + Tag.REGION.toString(), + dataRegion.getDataRegionIdString(), + Tag.TYPE.toString(), + Metric.LOAD_POINT_COUNT.toString()); + MetricService.getInstance() + .count( + node.getWritePointCount(), + Metric.LEADER_QUANTITY.toString(), + MetricLevel.CORE, + Tag.NAME.toString(), + Metric.POINTS_IN.toString(), + Tag.DATABASE.toString(), + databaseName, + Tag.REGION.toString(), + dataRegion.getDataRegionIdString(), + Tag.TYPE.toString(), + Metric.LOAD_POINT_COUNT.toString()); + })); + + return true; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/MemoryBoundedBuffer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/MemoryBoundedBuffer.java new file mode 100644 index 0000000000000..05b0204e58753 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/MemoryBoundedBuffer.java @@ -0,0 +1,78 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; + +/** + * LOAD memory budget: memory-pool accounting for the pieces of one source TsFile. It keeps two + * views in sync: + * + *
    + *
  • its own {@code dataSize} - the bytes currently buffered into pieces, and + *
  • the shared {@link LoadTsFileDataCacheMemoryBlock} - so the cluster-wide LOAD data cache is + * aware of this file's footprint. + *
+ * + * The budget is {@code thriftMaxFrameSize >> 2}; {@link #isMemoryEnough()} is the signal the + * pipeline polls after every chunk. When it turns false, {@link PieceDispatcher} evicts the largest + * buffered piece first. {@link #add(long)} / {@link #release(long)} keep both views consistent on + * every buffered/dispatched piece, and {@link #clear()} is the last-chance cleanup that returns any + * leftover accounting to the shared block. + */ +class MemoryBoundedBuffer { + + private static final long MAX_MEMORY_SIZE = + IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize() >> 2; + + private final LoadTsFileDataCacheMemoryBlock block; + private long dataSize = 0; + + MemoryBoundedBuffer(LoadTsFileDataCacheMemoryBlock block) { + this.block = block; + } + + boolean isMemoryEnough() { + return dataSize <= MAX_MEMORY_SIZE && block.hasEnoughMemory(); + } + + void add(long memorySize) { + dataSize += memorySize; + block.addMemoryUsage(memorySize); + } + + void release(long memorySize) { + dataSize -= memorySize; + block.reduceMemoryUsage(memorySize); + } + + long getDataSize() { + return dataSize; + } + + /** Returns all buffered accounting to the shared memory block; safe to call multiple times. */ + void clear() { + if (dataSize > 0) { + block.reduceMemoryUsage(dataSize); + dataSize = 0; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java new file mode 100644 index 0000000000000..2c5674f2dd111 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java @@ -0,0 +1,261 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.exception.load.RegionReplicaSetChangedException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkEncoder; +import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; +import org.apache.iotdb.db.storageengine.load.splitter.EncodedChunk; +import org.apache.iotdb.db.storageengine.load.splitter.EncodedChunkGroup; + +import org.apache.tsfile.utils.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.PriorityQueue; + +/** + * LOAD piece dispatcher: holds the buffered pieces of one source TsFile - one {@link + * LoadTsFilePieceNode} per target region - and decides when they are sent. + * + *

State: + * + *

    + *
  • {@code regionId2ReplicaSetAndNode} - the current piece of every touched region; + *
  • {@code largestPieceRegions} - max-heap of (region, piece size at offer time) used for + * largest-first eviction; entries become stale once a piece grows or is dispatched and are + * skipped lazily on poll. + *
+ * + *

Behaviors: + * + *

    + *
  • {@link #offerChunk(ChunkData, TRegionReplicaSet)} appends a chunk and throws {@link + * RegionReplicaSetChangedException} if the same region suddenly maps to a different replica + * set (region migration); + *
  • {@link #addDeletionToAll(DeletionData)} replicates a deletion into every buffered piece + * (memory is accounted once per region); + *
  • {@link #dispatchLargestUntilMemoryEnough()} evicts largest pieces while over budget; + *
  • {@link #flushAll()} flushes the remainder at end of file. + *
+ * + * Dispatch is delegated back through {@link DispatchCallback}, so this class never talks to + * consensus itself; {@link TwoPhaseConsensusLoadStrategy} implements the callback with its + * BEGIN/PIECE submission logic. + */ +class PieceDispatcher { + + private record PartitionKey(TConsensusGroupId regionId, long timePartitionStart) {} + + private static final Logger LOGGER = LoggerFactory.getLogger(PieceDispatcher.class); + + @FunctionalInterface + interface DispatchCallback { + boolean dispatch(LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet); + } + + private final LoadSingleTsFileNode singleTsFileNode; + private final MemoryBoundedBuffer memoryBuffer; + private final DispatchCallback dispatchCallback; + + private final Map> + regionId2ReplicaSetAndNode = new HashMap<>(); + private final Map partition2NextPieceIndex = new HashMap<>(); + private final Map partition2NextOffset = new HashMap<>(); + + /** + * Max-heap of (regionId, buffered piece size at offer time) used to dispatch the largest pieces + * first when the data cache is over budget. Entries become stale once the piece size changes or + * the piece is dispatched, and are skipped lazily on poll, which avoids re-sorting all buffered + * pieces on every over-budget event. + */ + private final PriorityQueue> largestPieceRegions = + new PriorityQueue<>((a, b) -> Long.compare(b.getValue(), a.getValue())); + + PieceDispatcher( + LoadSingleTsFileNode singleTsFileNode, + MemoryBoundedBuffer memoryBuffer, + DispatchCallback dispatchCallback) { + this.singleTsFileNode = singleTsFileNode; + this.memoryBuffer = memoryBuffer; + this.dispatchCallback = dispatchCallback; + } + + void offerChunk(ChunkData chunkData, TRegionReplicaSet replicaSet) + throws LoadFileException, IOException { + final TConsensusGroupId regionId = replicaSet.getRegionId(); + final EncodedChunkGroup encodedChunkGroup = ChunkEncoder.encode(chunkData); + final PartitionKey key = + new PartitionKey(regionId, encodedChunkGroup.getTimePartitionSlot().getStartTime()); + if (regionId2ReplicaSetAndNode.containsKey(key) + && !Objects.equals(regionId2ReplicaSetAndNode.get(key).getLeft(), replicaSet)) { + // Detected region replica set changed (maybe due to region migration), throw an exception + throw new RegionReplicaSetChangedException( + regionId2ReplicaSetAndNode.get(key).getLeft(), replicaSet); + } + + long nextOffset = partition2NextOffset.getOrDefault(key, 0L); + for (EncodedChunk encodedChunk : encodedChunkGroup.getChunks()) { + encodedChunk.setOffset(nextOffset); + nextOffset += encodedChunk.getBytes().length; + } + partition2NextOffset.put(key, nextOffset); + final long memoryDelta = encodedChunkGroup.getDataSize() - chunkData.getDataSize(); + if (memoryDelta > 0) { + memoryBuffer.add(memoryDelta); + } else if (memoryDelta < 0) { + memoryBuffer.release(-memoryDelta); + } + regionId2ReplicaSetAndNode + .computeIfAbsent(key, o -> new Pair<>(replicaSet, newPieceNode(key))) + .getRight() + .addTsFileData(encodedChunkGroup); + offerPieceRegion(key); + } + + /** Replicates the deletion into every buffered piece; memory is accounted once per region. */ + void addDeletionToAll(DeletionData deletionData) { + for (Map.Entry> entry : + regionId2ReplicaSetAndNode.entrySet()) { + memoryBuffer.add(deletionData.getDataSize()); + entry.getValue().getRight().addTsFileData(deletionData); + offerPieceRegion(entry.getKey()); + } + } + + /** Dispatches from the biggest buffered piece until the data cache is back under budget. */ + boolean dispatchLargestUntilMemoryEnough() throws LoadFileException { + while (!memoryBuffer.isMemoryEnough()) { + final PartitionKey key = pollLargestPieceRegion(); + if (key == null) { + // No dispatchable piece remains; the remaining buffered data stays buffered until the + // next flush (end of file, a later deletion, or another over-budget event). + break; + } + final Pair pair = regionId2ReplicaSetAndNode.get(key); + final LoadTsFilePieceNode pieceNode = pair.getRight(); + memoryBuffer.release(pieceNode.getDataSize()); + if (!dispatchOne(pieceNode, pair.getLeft())) { + return false; + } + replacePieceNode(key, pair.getLeft()); + } + return true; + } + + /** Dispatches every non-empty buffered piece, e.g. at the end of the source TsFile. */ + boolean flushAll() throws LoadFileException { + for (Map.Entry> entry : + regionId2ReplicaSetAndNode.entrySet()) { + final LoadTsFilePieceNode pieceNode = entry.getValue().getRight(); + if (pieceNode.getDataSize() == 0) { + continue; + } + if (!dispatchPieces(Collections.singleton(entry.getKey()))) { + return false; + } + } + return true; + } + + private boolean dispatchPieces(Collection keys) throws LoadFileException { + for (PartitionKey key : keys) { + final Pair pair = regionId2ReplicaSetAndNode.get(key); + if (pair == null) { + continue; + } + final LoadTsFilePieceNode pieceNode = pair.getRight(); + if (pieceNode.getDataSize() == 0) { + continue; + } + memoryBuffer.release(pieceNode.getDataSize()); + if (!dispatchOne(pieceNode, pair.getLeft())) { + LOGGER.warn( + DataNodeQueryMessages.DISPATCH_PIECE_NODE_ARG_OF_TSFILE_ARG_ERROR, + pieceNode, + singleTsFileNode.getTsFileResource().getTsFile()); + return false; + } + replacePieceNode(key, pair.getLeft()); + } + return true; + } + + private boolean dispatchOne(LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { + return dispatchCallback.dispatch(pieceNode, replicaSet); + } + + private void replacePieceNode(PartitionKey key, TRegionReplicaSet replicaSet) { + regionId2ReplicaSetAndNode.replace(key, new Pair<>(replicaSet, newPieceNode(key))); + } + + private LoadTsFilePieceNode newPieceNode(PartitionKey key) { + final long pieceIndex = partition2NextPieceIndex.getOrDefault(key, 0L); + partition2NextPieceIndex.put(key, pieceIndex + 1); + return new LoadTsFilePieceNode( + singleTsFileNode.getPlanNodeId(), + singleTsFileNode.getTsFileResource().getTsFile(), + pieceIndex); + } + + private void offerPieceRegion(final PartitionKey key) { + final Pair pair = regionId2ReplicaSetAndNode.get(key); + if (pair != null) { + largestPieceRegions.offer(Map.entry(key, pair.getRight().getDataSize())); + } + } + + /** Pops the region with the largest non-empty buffered piece, skipping stale heap entries. */ + private PartitionKey pollLargestPieceRegion() { + while (!largestPieceRegions.isEmpty()) { + final Map.Entry entry = largestPieceRegions.poll(); + final Pair pair = + regionId2ReplicaSetAndNode.get(entry.getKey()); + if (pair == null) { + continue; + } + final long currentSize = pair.getRight().getDataSize(); + if (entry.getValue() == currentSize && currentSize > 0) { + return entry.getKey(); + } + } + return null; + } + + void clear() { + regionId2ReplicaSetAndNode.clear(); + partition2NextPieceIndex.clear(); + partition2NextOffset.clear(); + largestPieceRegions.clear(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java new file mode 100644 index 0000000000000..eaba46488aab8 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java @@ -0,0 +1,93 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.db.storageengine.load.LoadTsFileChecksumUtils; + +import java.util.UUID; + +/** + * One high-cohesion state object per target region of the LOAD consensus two-phase protocol. It + * replaces the five parallel maps ({@code regionPieceCount}, {@code regionPieceTotalBytes}, {@code + * regionPieceChecksum}, {@code begunConsensusRegions}, {@code regionLoadId}) that used to live in + * {@link LoadTsFileScheduler}, so the load id, the BEGIN state and the three accumulated counters + * can never diverge. + * + *

Lifecycle, driven by {@link TwoPhaseConsensusLoadStrategy}: + * + *

    + *
  • Created lazily when the first piece of a region is dispatched ({@code + * consensusContexts.computeIfAbsent}). + *
  • {@link #markBegun()} runs before the BEGIN command is submitted, so a region is begun at + * most once per load. + *
  • {@link #accumulate(long, long)} runs after every successful PIECE: it increments the piece + * count, adds the piece bytes and XORs the piece checksum. + *
  • Phase two reads {@link #getPieceCount()}, {@link #getTotalBytes()} and {@link + * #getChecksum()} to build the PREPARE command. + *
+ * + * The load id is generated per instance, so every region of the same source file is isolated from + * the others on the write nodes. + */ +public class RegionConsensusContext { + + /** Each region gets its own load id so its staged data is isolated from other regions. */ + private final String loadId = UUID.randomUUID().toString(); + + private long pieceCount = 0; + private long totalBytes = 0; + private long checksum = 0; + + /** Whether the BEGIN command has already been sent to this region. */ + private boolean begun = false; + + public String getLoadId() { + return loadId; + } + + public long getPieceCount() { + return pieceCount; + } + + public long getTotalBytes() { + return totalBytes; + } + + public long getChecksum() { + return checksum; + } + + public boolean isBegun() { + return begun; + } + + public void markBegun() { + begun = true; + } + + /** Records one successfully applied piece: increments the count, adds bytes and XORs checksum. */ + public void accumulate(long bytes, long pieceChecksum) { + pieceCount++; + totalBytes += bytes; + // Order-sensitive aggregation (rotate-then-XOR): the aggregate must depend on the exact piece + // order, because the plain XOR used before could not detect reordered or swapped pieces. + checksum = LoadTsFileChecksumUtils.combine(checksum, pieceChecksum); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileLoadStrategy.java new file mode 100644 index 0000000000000..97b89f8208af1 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileLoadStrategy.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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.commons.exception.IoTDBException; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; + +/** + * Strategy of the LOAD pipeline for loading one source TsFile. The scheduler ({@link + * LoadTsFileScheduler}) decides which strategy applies via {@code needDecodeTsFile}; each strategy + * owns its complete pipeline, phase bookkeeping and phase time metrics. + * + *

Implementations: + * + *

    + *
  • {@link LocalLoadStrategy} - no decode needed, dispatch the file to the local region; + *
  • {@link TwoPhaseConsensusLoadStrategy} - decode, stream PIECE batches through consensus, + * then PREPARE+COMMIT or ABORT. + *
+ */ +public interface TsFileLoadStrategy { + + /** + * Loads the given TsFile. + * + * @return true if the file was loaded successfully + */ + boolean execute(LoadSingleTsFileNode node) throws IoTDBException; +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java new file mode 100644 index 0000000000000..2fe9e25d0a91e --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java @@ -0,0 +1,137 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Consumer; + +/** + * LOAD split consumer: receives every {@link TsFileData} of one source TsFile's split output from + * {@link TsFileSplitter} and routes -> buffers -> dispatches it as per-region consensus + * pieces. + * + *
    + *
  • {@code CHUNK} data is buffered directionless first; before dispatch {@link + * DataPartitionRouter} resolves the target regions and {@link PieceDispatcher} appends every + * chunk to its per-region piece. + *
  • {@code DELETION} data is replicated into every buffered piece (chunks are routed first so a + * deletion never overtakes chunk data). + *
  • {@link MemoryBoundedBuffer} guards the shared data cache; over budget, the largest piece is + * dispatched immediately. + *
  • At end of file {@link #sendAllTsFileData()} flushes the remainder. + *
+ * + * The pipeline also notifies the progress-index callback for every chunk's time partition, so the + * strategy can track pipe progress while splitting. {@link #clear()} releases all buffered + * accounting and piece references. + */ +public class TsFileSplitConsumer implements TsFileSplitter.TsFileDataConsumer { + + private final LoadSingleTsFileNode singleTsFileNode; + private final DataPartitionRouter router; + private final MemoryBoundedBuffer memoryBuffer; + private final PieceDispatcher dispatcher; + private final Consumer progressIndexCallback; + + private final List nonDirectionalChunkData = new ArrayList<>(); + + public TsFileSplitConsumer( + LoadSingleTsFileNode singleTsFileNode, + LoadTsFileDataCacheMemoryBlock block, + DataPartitionBatchFetcher partitionFetcher, + String userName, + Consumer progressIndexCallback, + PieceDispatcher.DispatchCallback dispatchCallback) { + this.singleTsFileNode = singleTsFileNode; + this.router = new DataPartitionRouter(partitionFetcher, userName); + this.memoryBuffer = new MemoryBoundedBuffer(block); + this.dispatcher = new PieceDispatcher(singleTsFileNode, memoryBuffer, dispatchCallback); + this.progressIndexCallback = progressIndexCallback; + } + + @Override + public boolean apply(TsFileData tsFileData) throws LoadFileException { + return switch (tsFileData.getType()) { + case CHUNK -> addOrSendChunkData((ChunkData) tsFileData); + case DELETION -> addOrSendDeletionData((DeletionData) tsFileData); + default -> + throw new UnsupportedOperationException( + String.format( + DataNodeQueryMessages.QUERY_EXCEPTION_UNSUPPORTED_TSFILEDATATYPE_S_374475FA, + tsFileData.getType())); + }; + } + + private boolean addOrSendChunkData(ChunkData chunkData) throws LoadFileException { + nonDirectionalChunkData.add(chunkData); + memoryBuffer.add(chunkData.getDataSize()); + progressIndexCallback.accept(chunkData.getTimePartitionSlot()); + + if (!memoryBuffer.isMemoryEnough()) { + routeChunkData(); + if (!dispatcher.dispatchLargestUntilMemoryEnough()) { + return false; + } + } + return true; + } + + private boolean addOrSendDeletionData(DeletionData deletionData) throws LoadFileException { + routeChunkData(); // ensure chunk data will be added before deletion + dispatcher.addDeletionToAll(deletionData); + return true; + } + + private void routeChunkData() throws LoadFileException { + if (nonDirectionalChunkData.isEmpty()) { + return; + } + + final List replicaSets = router.route(nonDirectionalChunkData); + for (int i = 0, size = nonDirectionalChunkData.size(); i < size; i++) { + dispatcher.offerChunk(nonDirectionalChunkData.get(i), replicaSets.get(i)); + } + nonDirectionalChunkData.clear(); + } + + boolean sendAllTsFileData() throws LoadFileException { + routeChunkData(); + return dispatcher.flushAll(); + } + + /** Last-chance cleanup: returns all buffered accounting and drops every piece reference. */ + void clear() { + memoryBuffer.clear(); + nonDirectionalChunkData.clear(); + dispatcher.clear(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java new file mode 100644 index 0000000000000..5a40933c4c4d7 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java @@ -0,0 +1,418 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; +import org.apache.iotdb.db.storageengine.load.LoadTsFileChecksumUtils; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; +import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Two-phase consensus LOAD strategy for files that need decoding. + * + *

Phase 1 (split & stream). {@link #execute(LoadSingleTsFileNode)} resets the + * per-file state, assigns a fresh uuid to {@link LoadTsFileDispatcherImpl} (executor naming / log + * correlation) and feeds the source TsFile through {@link TsFileSplitter} into {@link + * TsFileSplitConsumer}. Every dispatched piece goes through {@code dispatchConsensusPiece}: the + * first piece of a region sends BEGIN with a fresh per-region load id, then PIECE commands with a + * monotonically increasing {@code pieceIndex}; {@link RegionConsensusContext#accumulate(long, + * long)} records piece count, total bytes and the XOR checksum. Submission goes through {@link + * LoadConsensusSubmitter} with bounded retries for transient failures only. + * + *

Phase 2 (commit or abort). If every region received all its pieces, each touched region + * gets PREPARE (with the accumulated count/bytes/checksum) followed by COMMIT; otherwise every + * touched region gets ABORT so the staged data is dropped. + * + *

Per-file state: {@code allReplicaSets} (the regions to prepare/commit/abort), {@code + * consensusContexts} (per-region two-phase state) and {@code timePartitionSlotToProgressIndex} + * (pipe progress index per time partition, collected while splitting for the upcoming + * progress-index sync). + */ +public class TwoPhaseConsensusLoadStrategy implements TsFileLoadStrategy { + + private static final Logger LOGGER = LoggerFactory.getLogger(TwoPhaseConsensusLoadStrategy.class); + + private static final LoadTsFileCostMetricsSet LOAD_TSFILE_COST_METRICS_SET = + LoadTsFileCostMetricsSet.getInstance(); + + /** + * Bounded retry for transient LOAD consensus submission failures (network errors, region + * migration, transient server errors). The write node deduplicates pieces by (loadId, pieceIndex, + * checksum), so a retried request whose first attempt actually applied is acknowledged as success + * instead of being applied twice. + */ + private static final int LOAD_CONSENSUS_SUBMIT_MAX_RETRIES = 3; + + private static final long LOAD_CONSENSUS_SUBMIT_RETRY_BACKOFF_MS = 100L; + + private final LoadTsFileDispatcherImpl dispatcher; + private final DataPartitionBatchFetcher partitionFetcher; + private final LoadTsFileDataCacheMemoryBlock block; + private final LoadConsensusSubmitter consensusSubmitter; + private final String userName; + private final boolean isGeneratedByPipe; + + /** The source file being loaded, kept for the phase-two commands and BEGIN metadata. */ + private LoadSingleTsFileNode currentNode; + + /** Regions touched by the current file; used to send ABORT/PREPARE+COMMIT in phase two. */ + private final Set allReplicaSets = new HashSet<>(); + + /** Per-region two-phase state of the current file; replaces the old five parallel maps. */ + private final Map consensusContexts = + new ConcurrentHashMap<>(); + + /** + * Progress index per time partition, assigned while the file is being split. Kept for the + * upcoming progress-index sync with the consensus prepare phase. + */ + private final Map timePartitionSlotToProgressIndex = + new HashMap<>(); + + public TwoPhaseConsensusLoadStrategy( + LoadTsFileDispatcherImpl dispatcher, + DataPartitionBatchFetcher partitionFetcher, + LoadTsFileDataCacheMemoryBlock block, + LoadConsensusSubmitter consensusSubmitter, + String userName, + boolean isGeneratedByPipe) { + this.dispatcher = dispatcher; + this.partitionFetcher = partitionFetcher; + this.block = block; + this.consensusSubmitter = consensusSubmitter; + this.userName = userName; + this.isGeneratedByPipe = isGeneratedByPipe; + } + + @Override + public boolean execute(LoadSingleTsFileNode node) { + this.currentNode = node; + dispatcher.setUuid(UUID.randomUUID().toString()); + allReplicaSets.clear(); + consensusContexts.clear(); + timePartitionSlotToProgressIndex.clear(); + + long startTime = System.nanoTime(); + final boolean isFirstPhaseSuccess; + try { + isFirstPhaseSuccess = firstPhase(node); + } finally { + LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.FIRST_PHASE, System.nanoTime() - startTime); + } + + startTime = System.nanoTime(); + final boolean isSecondPhaseSuccess; + try { + isSecondPhaseSuccess = secondPhase(node, isFirstPhaseSuccess); + } finally { + LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.SECOND_PHASE, System.nanoTime() - startTime); + } + + return isFirstPhaseSuccess && isSecondPhaseSuccess; + } + + private boolean firstPhase(LoadSingleTsFileNode node) { + final TsFileSplitConsumer pipeline = + new TsFileSplitConsumer( + node, + block, + partitionFetcher, + userName, + this::computeTimePartitionSlotToProgressIndexIfAbsent, + this::dispatchOnePieceNode); + try { + new TsFileSplitter(node.getTsFileResource().getTsFile(), pipeline) + .splitTsFileByDataPartition(); + return pipeline.sendAllTsFileData(); + } catch (IllegalStateException e) { + LOGGER.warn( + String.format( + DataNodeQueryMessages.DISPATCH_TSFILEDATA_ERROR_WHEN_PARSING_TSFILE_S, + node.getTsFileResource().getTsFile()), + e); + return false; + } catch (Exception e) { + LOGGER.warn( + String.format( + DataNodeQueryMessages.PARSE_OR_SEND_TSFILE_S_ERROR, + node.getTsFileResource().getTsFile()), + e); + return false; + } finally { + pipeline.clear(); + } + } + + private boolean dispatchOnePieceNode( + LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { + allReplicaSets.add(replicaSet); + return dispatchConsensusPiece(pieceNode, replicaSet); + } + + /** + * Submits a LOAD consensus request with a bounded number of attempts. Only transient failures are + * retried; permanent rejections (checksum mismatch, missing staged writer) are returned to the + * caller immediately so the scheduler can abort. + */ + private TSStatus submitConsensusWithRetry( + TRegionReplicaSet replicaSet, LoadTsFileConsensusNode node) { + TSStatus status = null; + for (int attempt = 1; attempt <= LOAD_CONSENSUS_SUBMIT_MAX_RETRIES; attempt++) { + status = consensusSubmitter.submit(replicaSet, node); + if (status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() + || !isTransientConsensusFailure(status) + || attempt == LOAD_CONSENSUS_SUBMIT_MAX_RETRIES) { + break; + } + LOGGER.warn( + DataNodeQueryMessages.LOG_LOAD_CONSENSUS_SUBMIT_TRANSIENT_FAILURE_RETRY_D7E1D9A6, + node.getOp(), + node.getLoadId(), + replicaSet, + attempt, + LOAD_CONSENSUS_SUBMIT_MAX_RETRIES, + status.getMessage()); + try { + Thread.sleep(LOAD_CONSENSUS_SUBMIT_RETRY_BACKOFF_MS * attempt); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + return status; + } + + private boolean isTransientConsensusFailure(TSStatus status) { + switch (TSStatusCode.representOf(status.getCode())) { + case DISPATCH_ERROR: + case INTERNAL_SERVER_ERROR: + case NO_AVAILABLE_REGION_GROUP: + case EXECUTE_STATEMENT_ERROR: + return true; + default: + return false; + } + } + + private boolean dispatchConsensusPiece( + LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { + final TConsensusGroupId regionId = replicaSet.getRegionId(); + final RegionConsensusContext context = + consensusContexts.computeIfAbsent(regionId, o -> new RegionConsensusContext()); + final String loadId = context.getLoadId(); + + if (!context.isBegun()) { + context.markBegun(); + final LoadTsFileConsensusNode begin = + LoadTsFileConsensusNode.begin( + new PlanNodeId("load-begin-" + loadId), + loadId, + pieceNode.getTsFile() == null ? null : pieceNode.getTsFile().getName(), + currentNode.isTableModel(), + currentNode.getDatabase(), + // The total piece count is only known after phase one; PREPARE carries the real + // count, BEGIN keeps the "unknown" sentinel so the two never disagree. + -1); + final TSStatus beginStatus = submitConsensusWithRetry(replicaSet, begin); + if (beginStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages.DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG + + DataNodeQueryMessages + .RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG, + replicaSet, + TSStatusCode.representOf(beginStatus.getCode()).name(), + beginStatus.getMessage(), + pieceNode); + return false; + } + } + + final long pieceIndex = pieceNode.getPieceIndex(); + // The checksum is part of the consensus contract: both the write node and every follower + // verify it, and the per-piece digest includes the piece index so reordered payloads are + // detected. Sending a constant here would silently disable checksum validation. + final long checksum = + LoadTsFileChecksumUtils.checksum(pieceIndex, pieceNode.getAllTsFileData()); + final LoadTsFileConsensusNode piece = + LoadTsFileConsensusNode.piece( + new PlanNodeId("load-piece-" + loadId + "-" + pieceIndex), + loadId, + pieceNode.getTsFile() == null ? null : pieceNode.getTsFile().getName(), + pieceIndex, + 0L, + pieceNode.getAllTsFileData(), + checksum); + final TSStatus pieceStatus = submitConsensusWithRetry(replicaSet, piece); + if (pieceStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages.DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG + + DataNodeQueryMessages + .RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG, + replicaSet, + TSStatusCode.representOf(pieceStatus.getCode()).name(), + pieceStatus.getMessage(), + pieceNode); + return false; + } + context.accumulate(piece.getDataSize(), piece.getChecksum()); + return true; + } + + private boolean secondPhase(LoadSingleTsFileNode node, boolean isFirstPhaseSuccess) { + if (!isFirstPhaseSuccess) { + return abortAllRegions(); + } + return prepareAndCommitAllRegions(node); + } + + private boolean abortAllRegions() { + for (TRegionReplicaSet replicaSet : allReplicaSets) { + final String loadId = consensusContexts.get(replicaSet.getRegionId()).getLoadId(); + final LoadTsFileConsensusNode abort = + LoadTsFileConsensusNode.abort( + new PlanNodeId("load-abort-" + loadId), loadId, null, isGeneratedByPipe); + final TSStatus status = consensusSubmitter.submit(replicaSet, abort); + if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages + .DISPATCH_LOAD_COMMAND_ARG_OF_TSFILE_ARG_ERROR_TO_REPLICASETS_ARG_ERROR + + DataNodeQueryMessages.RESULT_STATUS_CODE_ARG_RESULT_STATUS_MESSAGE_ARG, + abort, + loadId, + allReplicaSets, + TSStatusCode.representOf(status.getCode()).name(), + status.getMessage()); + return false; + } + } + return true; + } + + private boolean prepareAndCommitAllRegions(LoadSingleTsFileNode node) { + final Map timePartition2ProgressIndex = + serializeTimePartitionProgressIndexes(); + for (TRegionReplicaSet replicaSet : allReplicaSets) { + final RegionConsensusContext context = consensusContexts.get(replicaSet.getRegionId()); + final String loadId = context.getLoadId(); + final LoadTsFileConsensusNode prepare = + LoadTsFileConsensusNode.prepare( + new PlanNodeId("load-prepare-" + loadId), + loadId, + null, + (int) context.getPieceCount(), + context.getTotalBytes(), + context.getChecksum(), + timePartition2ProgressIndex); + final TSStatus prepareStatus = consensusSubmitter.submit(replicaSet, prepare); + if (prepareStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages + .DISPATCH_LOAD_COMMAND_ARG_OF_TSFILE_ARG_ERROR_TO_REPLICASETS_ARG_ERROR + + DataNodeQueryMessages.RESULT_STATUS_CODE_ARG_RESULT_STATUS_MESSAGE_ARG, + prepare, + loadId, + allReplicaSets, + TSStatusCode.representOf(prepareStatus.getCode()).name(), + prepareStatus.getMessage()); + return false; + } + + final LoadTsFileConsensusNode commit = + LoadTsFileConsensusNode.commit( + new PlanNodeId("load-commit-" + loadId), + loadId, + null, + isGeneratedByPipe, + node.isDeleteAfterLoad(), + timePartition2ProgressIndex); + final TSStatus commitStatus = consensusSubmitter.submit(replicaSet, commit); + if (commitStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages + .DISPATCH_LOAD_COMMAND_ARG_OF_TSFILE_ARG_ERROR_TO_REPLICASETS_ARG_ERROR + + DataNodeQueryMessages.RESULT_STATUS_CODE_ARG_RESULT_STATUS_MESSAGE_ARG, + commit, + loadId, + allReplicaSets, + TSStatusCode.representOf(commitStatus.getCode()).name(), + commitStatus.getMessage()); + return false; + } + } + return true; + } + + /** + * Serializes the per-time-partition {@link ProgressIndex} collected during splitting into the + * byte form carried by PREPARE/COMMIT, so the receiving side can restore the real progress + * instead of degrading it to {@code MinimumProgressIndex}. + */ + private Map serializeTimePartitionProgressIndexes() { + final Map result = new HashMap<>(); + for (Map.Entry entry : + timePartitionSlotToProgressIndex.entrySet()) { + try { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + entry.getValue().serialize(baos); + result.put(entry.getKey(), baos.toByteArray()); + } catch (IOException e) { + throw new IllegalStateException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_PROGRESS_SERIALIZE_FAILED_28EFD091, + entry.getKey().getStartTime()), + e); + } + } + return result; + } + + private void computeTimePartitionSlotToProgressIndexIfAbsent( + final TTimePartitionSlot timePartitionSlot) { + timePartitionSlotToProgressIndex.putIfAbsent( + timePartitionSlot, PipeDataNodeAgent.runtime().getNextProgressIndexForTsFileLoad()); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java index f7350c87323a3..5ca04b4d7d3af 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java @@ -146,6 +146,11 @@ public boolean isAligned() { return true; } + @Override + public boolean isEntireChunk() { + return !needDecodeChunk; + } + @Override public void writeToFileWriter(final TsFileIOWriter writer) throws IOException, PageException { writeTsFileData(writer); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java index 7cc5db029950e..cddd35deafa57 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java @@ -42,6 +42,8 @@ public interface ChunkData extends TsFileData { boolean isAligned(); + boolean isEntireChunk(); + void writeEntireChunk(ByteBuffer chunkData, IChunkMetadata chunkMetadata) throws IOException; void writeEntirePage(PageHeader pageHeader, ByteBuffer pageData) throws IOException; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkEncoder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkEncoder.java new file mode 100644 index 0000000000000..6299c77a8a078 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkEncoder.java @@ -0,0 +1,54 @@ +/* + * 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 + */ +package org.apache.iotdb.db.storageengine.load.splitter; + +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.write.writer.TsFileIOWriter; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +/** Encodes an undecoded source chunk into its physical TsFile chunk bytes. */ +public final class ChunkEncoder { + + private ChunkEncoder() {} + + public static EncodedChunkGroup encode(final ChunkData chunkData) throws IOException { + final MemoryTsFileOutput output = new MemoryTsFileOutput(); + final TsFileIOWriter writer = new TsFileIOWriter(output); + writer.startChunkGroup(chunkData.getDevice()); + final int chunkStart = (int) writer.getPos(); + try { + chunkData.writeToFileWriter(writer); + } catch (Exception e) { + throw new IOException(e); + } + final int chunkEnd = (int) writer.getPos(); + final List metadata = writer.getChunkMetadataListOfCurrentDeviceInMemory(); + final List chunks = new ArrayList<>(metadata.size()); + for (int i = 0; i < metadata.size(); i++) { + final int start = (int) metadata.get(i).getOffsetOfChunkHeader(); + final int end = + i + 1 < metadata.size() ? (int) metadata.get(i + 1).getOffsetOfChunkHeader() : chunkEnd; + if (end > start && start >= chunkStart) { + chunks.add( + new EncodedChunk( + chunkData.getDevice(), + chunkData.getTimePartitionSlot(), + output.copyRange(start, end - start), + new ChunkMetadata(metadata.get(i)))); + } + } + return new EncodedChunkGroup(chunkData.getDevice(), chunkData.getTimePartitionSlot(), chunks); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunk.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunk.java new file mode 100644 index 0000000000000..512ea5cef55f1 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunk.java @@ -0,0 +1,62 @@ +/* + * 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 + */ +package org.apache.iotdb.db.storageengine.load.splitter; + +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; + +import org.apache.tsfile.file.metadata.IChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; + +/** Physical Chunk bytes and the metadata required to seal the TsFile footer. */ +public final class EncodedChunk { + private final IDeviceID device; + private final TTimePartitionSlot timePartitionSlot; + private final byte[] bytes; + private final IChunkMetadata metadata; + private long offset; + + public EncodedChunk( + IDeviceID device, + TTimePartitionSlot timePartitionSlot, + byte[] bytes, + IChunkMetadata metadata) { + this.device = device; + this.timePartitionSlot = timePartitionSlot; + this.bytes = bytes; + this.metadata = metadata; + this.offset = 0; + } + + public IDeviceID getDevice() { + return device; + } + + public TTimePartitionSlot getTimePartitionSlot() { + return timePartitionSlot; + } + + public byte[] getBytes() { + return bytes; + } + + public IChunkMetadata getMetadata() { + return metadata; + } + + public long getOffset() { + return offset; + } + + public void setOffset(long offset) { + this.offset = offset; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunkGroup.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunkGroup.java new file mode 100644 index 0000000000000..6761b312a049a --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunkGroup.java @@ -0,0 +1,68 @@ +/* + * 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. + * The ASF licenses this file to you under the Apache License, Version 2.0. + */ +package org.apache.iotdb.db.storageengine.load.splitter; + +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.apache.tsfile.utils.ReadWriteIOUtils; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.util.List; + +public final class EncodedChunkGroup implements TsFileData { + private final IDeviceID device; + private final TTimePartitionSlot timePartitionSlot; + private final List chunks; + + public EncodedChunkGroup( + IDeviceID device, TTimePartitionSlot timePartitionSlot, List chunks) { + this.device = device; + this.timePartitionSlot = timePartitionSlot; + this.chunks = List.copyOf(chunks); + } + + public IDeviceID getDevice() { + return device; + } + + public TTimePartitionSlot getTimePartitionSlot() { + return timePartitionSlot; + } + + public List getChunks() { + return chunks; + } + + @Override + public long getDataSize() { + return chunks.stream().mapToLong(chunk -> chunk.getBytes().length).sum(); + } + + @Override + public TsFileDataType getType() { + return TsFileDataType.ENCODED_CHUNK; + } + + @Override + public void serialize(DataOutputStream stream) throws IOException { + ReadWriteIOUtils.write(getType().ordinal(), stream); + ReadWriteIOUtils.write(timePartitionSlot.getStartTime(), stream); + ReadWriteIOUtils.write(device instanceof StringArrayDeviceID, stream); + device.serialize(stream); + ReadWriteIOUtils.write(chunks.size(), stream); + for (EncodedChunk chunk : chunks) { + ReadWriteIOUtils.write(chunk.getMetadata().getMeasurementUid(), stream); + ReadWriteIOUtils.write(chunk.getMetadata().getDataType().ordinal(), stream); + ReadWriteIOUtils.write(chunk.getOffset(), stream); + ReadWriteIOUtils.write(chunk.getBytes().length, stream); + stream.write(chunk.getBytes()); + chunk.getMetadata().serializeTo(stream, true); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/MemoryTsFileOutput.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/MemoryTsFileOutput.java new file mode 100644 index 0000000000000..cdcf2365a4f9b --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/MemoryTsFileOutput.java @@ -0,0 +1,72 @@ +/* + * 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. + * The ASF licenses this file to you under the Apache License, Version 2.0. + */ +package org.apache.iotdb.db.storageengine.load.splitter; + +import org.apache.tsfile.utils.PublicBAOS; +import org.apache.tsfile.write.writer.TsFileOutput; + +import java.io.IOException; +import java.io.OutputStream; +import java.nio.ByteBuffer; + +final class MemoryTsFileOutput implements TsFileOutput { + private final PublicBAOS output = new PublicBAOS(); + + @Override + public void write(byte[] bytes) throws IOException { + output.write(bytes); + } + + @Override + public void write(byte value) { + output.write(value); + } + + @Override + public void write(ByteBuffer buffer) { + final int length = buffer.remaining(); + if (buffer.hasArray()) { + output.write(buffer.array(), buffer.arrayOffset() + buffer.position(), length); + buffer.position(buffer.position() + length); + } else { + final byte[] bytes = new byte[length]; + buffer.get(bytes); + output.write(bytes, 0, bytes.length); + } + } + + @Override + public long getPosition() { + return output.size(); + } + + @Override + public void close() throws IOException { + output.close(); + } + + @Override + public OutputStream wrapAsStream() { + return output; + } + + @Override + public void flush() {} + + @Override + public void truncate(long size) { + output.reset(); + } + + @Override + public void force() {} + + byte[] copyRange(int offset, int length) { + final byte[] result = new byte[length]; + System.arraycopy(output.getBuf(), offset, result, 0, length); + return result; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java index 7c385c3107bf9..fdba8fa3d88ca 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java @@ -114,6 +114,11 @@ public boolean isAligned() { return false; } + @Override + public boolean isEntireChunk() { + return !needDecodeChunk; + } + @Override public void writeToFileWriter(final TsFileIOWriter writer) throws IOException { ensureDataReadyForWriting(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileDataType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileDataType.java index 9e3f73f724dd5..bf983ec923d6a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileDataType.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileDataType.java @@ -21,5 +21,6 @@ public enum TsFileDataType { CHUNK, + ENCODED_CHUNK, DELETION } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/load/LoadTsFileNodeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/load/LoadTsFileNodeTest.java index 37cedc50b3430..238eef0112ff9 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/load/LoadTsFileNodeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/load/LoadTsFileNodeTest.java @@ -71,6 +71,7 @@ public void testLoadSingleTsFileNode() { public void testLoadTsFilePieceNode() { LoadTsFilePieceNode node = new LoadTsFilePieceNode(new PlanNodeId(""), new File("1")); Assert.assertEquals(0, node.getDataSize()); + Assert.assertEquals(0L, node.getPieceIndex()); Assert.assertEquals(new ArrayList<>(), node.getAllTsFileData()); Assert.assertEquals(node.getTsFile(), new File("1")); Assert.assertNull(node.getRegionReplicaSet()); @@ -86,11 +87,12 @@ public void testLoadTsFilePieceNode() { } catch (NotImplementedException ignored) { } Assert.assertEquals(0, node.allowedChildCount()); - Assert.assertEquals("LoadTsFilePieceNode{tsFile=1, dataSize=0}", node.toString()); + Assert.assertEquals("LoadTsFilePieceNode{tsFile=1, pieceIndex=0, dataSize=0}", node.toString()); ByteBuffer buffer = ByteBuffer.allocate(1024); node.serialize(buffer); LoadTsFilePieceNode node1 = (LoadTsFilePieceNode) LoadTsFilePieceNode.deserialize(buffer); Assert.assertEquals(node.getTsFile(), node1.getTsFile()); + Assert.assertEquals(node.getPieceIndex(), node1.getPieceIndex()); } @Test diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java new file mode 100644 index 0000000000000..7173f5cdc5884 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java @@ -0,0 +1,133 @@ +/* + * 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.queryengine.plan.planner.plan.node.load; + +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; + +import org.junit.Assert; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Collections; + +public class LoadTsFileConsensusNodeTest { + + @Test + public void testBeginSerializeRoundTrip() { + final LoadTsFileConsensusNode begin = + LoadTsFileConsensusNode.begin(new PlanNodeId("n1"), "load-1", "file-1", true, "db", 3); + final ByteBuffer buffer = begin.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.BEGIN, node.getOp()); + Assert.assertEquals("load-1", node.getLoadId()); + Assert.assertEquals("file-1", node.getTsFileId()); + Assert.assertTrue(node.isTableModel()); + Assert.assertEquals("db", node.getDatabase()); + Assert.assertEquals(3, node.getExpectedPieceCount()); + } + + @Test + public void testPieceSerializeRoundTrip() { + final LoadTsFileConsensusNode commit = + LoadTsFileConsensusNode.commit( + new PlanNodeId("n2"), "load-2", "file-2", true, false, Collections.emptyMap()); + final ByteBuffer buffer = commit.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.COMMIT, node.getOp()); + Assert.assertTrue(node.isGeneratedByPipe()); + } + + @Test + public void testPiecePreservesCallerChecksum() { + final LoadTsFileConsensusNode piece = + LoadTsFileConsensusNode.piece( + new PlanNodeId("checksum"), + "load", + "file", + 0L, + 0L, + Collections.emptyList(), + 987654321L); + + Assert.assertEquals(987654321L, piece.getChecksum()); + } + + @Test + public void testPieceRefSerializeRoundTrip() { + final LoadTsFileConsensusNode.PieceRef ref = + new LoadTsFileConsensusNode.PieceRef( + "task-1/partition.tsfile", 4096L, 3L, new byte[] {1, 2, 3}); + final LoadTsFileConsensusNode piece = + LoadTsFileConsensusNode.pieceRefs( + new PlanNodeId("n3"), "load-3", "file-3", 0L, Collections.singletonList(ref), 0L, 3L); + final ByteBuffer buffer = piece.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.PIECE, node.getOp()); + Assert.assertEquals(1, node.getPieceRefs().size()); + final LoadTsFileConsensusNode.PieceRef actual = node.getPieceRefs().get(0); + Assert.assertEquals("task-1/partition.tsfile", actual.getRelativePath()); + Assert.assertEquals(4096L, actual.getOffset()); + Assert.assertEquals(3L, actual.getSize()); + Assert.assertTrue(Arrays.equals(new byte[] {1, 2, 3}, actual.getContent())); + } + + @Test + public void testPieceMarkerSerializeRoundTrip() { + final LoadTsFileConsensusNode marker = + LoadTsFileConsensusNode.pieceMarker( + new PlanNodeId("n4"), "load-4", "file-4", 7L, 123456L, 1024L); + final ByteBuffer buffer = marker.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.PIECE, node.getOp()); + Assert.assertEquals("load-4", node.getLoadId()); + Assert.assertEquals("file-4", node.getTsFileId()); + Assert.assertEquals(7L, node.getPieceIndex()); + Assert.assertEquals(123456L, node.getChecksum()); + Assert.assertEquals(1024L, node.getDataSize()); + Assert.assertFalse(node.hasChunkData()); + Assert.assertTrue(node.getPieceRefs().isEmpty()); + } + + @Test + public void testPullSerializeRoundTrip() { + final LoadTsFileConsensusNode pull = + LoadTsFileConsensusNode.pull( + new PlanNodeId("n5"), "load-5", "file-5", 3L, 999L, "192.168.1.10:6667"); + final ByteBuffer buffer = pull.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.PULL, node.getOp()); + Assert.assertEquals(3L, node.getPieceIndex()); + Assert.assertEquals(999L, node.getChecksum()); + Assert.assertEquals("192.168.1.10:6667", node.getPullSourceEndPoint()); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java index 4605e6cf814f3..1dd8a45c9b61b 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java @@ -19,11 +19,6 @@ package org.apache.iotdb.db.queryengine.plan.scheduler.load; -import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; -import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType; -import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.PlanFragmentId; @@ -44,9 +39,9 @@ import java.io.File; import java.lang.reflect.Constructor; -import java.lang.reflect.Field; import java.lang.reflect.Method; -import java.util.Arrays; +import java.util.ArrayList; +import java.util.Collections; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; @@ -103,16 +98,15 @@ public void testGetPartitionQueryDatabaseForTableModelLoad() { @Test public void testBuildRetryTreeLoadStatementUpdatesDatabaseLevel() throws Exception { - final LoadTsFileScheduler scheduler = - new LoadTsFileScheduler( - distributedQueryPlan, + final LoadFallbackHandler fallbackHandler = + new LoadFallbackHandler( mock(MPPQueryContext.class), - mock(QueryStateMachine.class), - mock(IClientManager.class), - mock(IPartitionFetcher.class), - true); + true, + Collections.emptyList(), + new ArrayList<>(), + mock(QueryStateMachine.class)); final Method method = - LoadTsFileScheduler.class.getDeclaredMethod( + LoadFallbackHandler.class.getDeclaredMethod( "buildRetryTreeLoadStatement", String.class, boolean.class, String.class); method.setAccessible(true); @@ -121,7 +115,7 @@ public void testBuildRetryTreeLoadStatementUpdatesDatabaseLevel() throws Excepti final LoadTsFileStatement statement = (LoadTsFileStatement) - method.invoke(scheduler, tsFile.getAbsolutePath(), true, "root.test.sg_0"); + method.invoke(fallbackHandler, tsFile.getAbsolutePath(), true, "root.test.sg_0"); Assert.assertEquals("root.test.sg_0", statement.getDatabase()); Assert.assertEquals(2, statement.getDatabaseLevel()); @@ -129,78 +123,24 @@ public void testBuildRetryTreeLoadStatementUpdatesDatabaseLevel() throws Excepti } @Test - public void testTsFileDataManagerClearReleasesCachedMemory() throws Exception { + public void testMemoryBoundedBufferClearReleasesCachedMemory() throws Exception { final Constructor memoryBlockConstructor = LoadTsFileDataCacheMemoryBlock.class.getDeclaredConstructor(long.class); memoryBlockConstructor.setAccessible(true); final LoadTsFileDataCacheMemoryBlock memoryBlock = memoryBlockConstructor.newInstance(1024 * 1024L); - final Class dataManagerClass = - Class.forName(LoadTsFileScheduler.class.getName() + "$TsFileDataManager"); - final Constructor dataManagerConstructor = - dataManagerClass.getDeclaredConstructor( - LoadTsFileScheduler.class, - LoadSingleTsFileNode.class, - LoadTsFileDataCacheMemoryBlock.class); - dataManagerConstructor.setAccessible(true); - final Object dataManager = - dataManagerConstructor.newInstance( - mock(LoadTsFileScheduler.class), mock(LoadSingleTsFileNode.class), memoryBlock); - // Simulate data buffered before split or routing aborts. clear() is the last chance to return // this accounting to the shared LOAD memory block. final long cachedMemorySize = 128L; - memoryBlock.addMemoryUsage(cachedMemorySize); - final Field dataSizeField = dataManagerClass.getDeclaredField("dataSize"); - dataSizeField.setAccessible(true); - dataSizeField.setLong(dataManager, cachedMemorySize); - - final Method clearMethod = dataManagerClass.getDeclaredMethod("clear"); - clearMethod.setAccessible(true); - clearMethod.invoke(dataManager); + final MemoryBoundedBuffer memoryBoundedBuffer = new MemoryBoundedBuffer(memoryBlock); + memoryBoundedBuffer.add(cachedMemorySize); + memoryBoundedBuffer.clear(); final Method getMemoryUsageMethod = LoadTsFileDataCacheMemoryBlock.class.getDeclaredMethod("getMemoryUsageInBytes"); getMemoryUsageMethod.setAccessible(true); Assert.assertEquals(0L, getMemoryUsageMethod.invoke(memoryBlock)); - Assert.assertEquals(0L, dataSizeField.getLong(dataManager)); - } - - @Test - public void testRegionReplicaSetComparison() { - final TDataNodeLocation dataNode1 = createDataNodeLocation(1, 10731); - final TDataNodeLocation dataNode3 = createDataNodeLocation(3, 10733); - final TDataNodeLocation dataNode5 = createDataNodeLocation(5, 10735); - final TConsensusGroupId regionId = new TConsensusGroupId(TConsensusGroupType.DataRegion, 1); - final TRegionReplicaSet original = - new TRegionReplicaSet(regionId, Arrays.asList(dataNode5, dataNode3, dataNode1)); - - Assert.assertTrue( - LoadTsFileScheduler.isSameRegionReplicaSet( - original, - new TRegionReplicaSet(regionId, Arrays.asList(dataNode3, dataNode5, dataNode1)))); - Assert.assertFalse( - LoadTsFileScheduler.isSameRegionReplicaSet( - original, - new TRegionReplicaSet( - regionId, Arrays.asList(dataNode3, dataNode5, createDataNodeLocation(7, 10737))))); - Assert.assertFalse( - LoadTsFileScheduler.isSameRegionReplicaSet( - original, - new TRegionReplicaSet( - regionId, Arrays.asList(dataNode3, dataNode5, createDataNodeLocation(1, 11731))))); - Assert.assertFalse( - LoadTsFileScheduler.isSameRegionReplicaSet( - original, - new TRegionReplicaSet( - new TConsensusGroupId(TConsensusGroupType.DataRegion, 2), - Arrays.asList(dataNode3, dataNode5, dataNode1)))); - } - - private static TDataNodeLocation createDataNodeLocation(int dataNodeId, int internalPort) { - return new TDataNodeLocation() - .setDataNodeId(dataNodeId) - .setInternalEndPoint(new TEndPoint("127.0.0.1", internalPort)); + Assert.assertEquals(0L, memoryBoundedBuffer.getDataSize()); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java index 20c3c1bf879a2..2cb87b94e39ec 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java @@ -210,6 +210,7 @@ public enum PlanNodeType { RELATIONAL_INSERT_ROWS((short) 2002), RELATIONAL_DELETE_DATA((short) 2003), OBJECT_FILE_NODE((short) 2004), + LOAD_TSFILE_CONSENSUS((short) 2010), ; private static final IPlanNodeDeserializer DESERIALIZER; From d8d5bc17fd11645c3d0f381c594c439d95d45c78 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Fri, 11 Sep 2026 09:35:39 +0800 Subject: [PATCH 2/8] feat(load): route decoded tsfile pieces through consensus --- .../consensus/config/IoTConsensusConfig.java | 18 +- .../iot/logdispatcher/LogDispatcher.java | 148 +---- .../dataregion/DataExecutionVisitor.java | 47 ++ .../node/DataNodePlanNodeDeserializer.java | 7 + .../plan/planner/plan/node/PlanVisitor.java | 5 + .../node/load/LoadTsFileConsensusNode.java | 74 +-- .../plan/scheduler/load/ChunkLayout.java | 27 + .../scheduler/load/ChunkOffsetCalculator.java | 144 ++++ .../scheduler/load/LoadTsFileScheduler.java | 10 +- .../plan/scheduler/load/PieceDispatcher.java | 40 +- .../load/RegionConsensusContext.java | 93 --- .../scheduler/load/TsFileSplitConsumer.java | 7 +- .../load/TwoPhaseConsensusLoadStrategy.java | 62 +- .../iotdb/db/storageengine/StorageEngine.java | 60 +- .../storageengine/dataregion/DataRegion.java | 51 ++ .../dataregion/wal/buffer/WALEntryType.java | 4 +- .../storageengine/load/LoadTsFileManager.java | 197 +++--- .../load/splitter/AlignedChunkData.java | 620 ++++++++++-------- .../BatchedAlignedValueChunkData.java | 248 +++---- .../load/splitter/ChunkData.java | 21 + .../load/splitter/ChunkEncoder.java | 54 -- .../load/splitter/EncodedChunk.java | 62 -- .../load/splitter/EncodedChunkGroup.java | 68 -- .../load/splitter/MemoryTsFileOutput.java | 72 -- .../load/splitter/NonAlignedChunkData.java | 400 ++++++----- .../load/splitter/TsFileData.java | 6 +- .../load/splitter/TsFileDataType.java | 1 - .../load/splitter/TsFileSplitter.java | 46 +- .../write/TsFilePrecalculatedChunkWriter.java | 294 +++++++++ .../load/LoadTsFileConsensusNodeTest.java | 15 - .../splitter/ChunkDataDirectWriteTest.java | 68 ++ 31 files changed, 1609 insertions(+), 1360 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkLayout.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkOffsetCalculator.java delete mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java delete mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkEncoder.java delete mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunk.java delete mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunkGroup.java delete mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/MemoryTsFileOutput.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java diff --git a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/config/IoTConsensusConfig.java b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/config/IoTConsensusConfig.java index a84e67f87172b..c6ca9a7e0dd5c 100644 --- a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/config/IoTConsensusConfig.java +++ b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/config/IoTConsensusConfig.java @@ -311,7 +311,6 @@ public static class Replication { private final int maxLogEntriesNumPerBatch; private final int maxSizePerBatch; private final int maxPendingBatchesNum; - private final int shardMaxSize; private final int maxQueueLength; private final long maxWaitingTimeForWaitBatchInMs; @@ -345,12 +344,10 @@ private Replication( long regionMigrationSpeedLimitBytesPerSecond, long subscriptionWalRetentionSizeInBytes, long subscriptionWalRetentionTimeMs, - long snapshotTransmissionProgressLogIntervalMs, - int shardMaxSize) { + long snapshotTransmissionProgressLogIntervalMs) { this.maxLogEntriesNumPerBatch = maxLogEntriesNumPerBatch; this.maxSizePerBatch = maxSizePerBatch; this.maxPendingBatchesNum = maxPendingBatchesNum; - this.shardMaxSize = shardMaxSize; this.maxQueueLength = maxQueueLength; this.maxWaitingTimeForWaitBatchInMs = maxWaitingTimeForWaitBatchInMs; this.maxWaitingTimeForAccumulatingBatchInMs = maxWaitingTimeForAccumulatingBatchInMs; @@ -379,10 +376,6 @@ public int getMaxPendingBatchesNum() { return maxPendingBatchesNum; } - public int getShardMaxSize() { - return shardMaxSize; - } - public int getMaxQueueLength() { return maxQueueLength; } @@ -448,7 +441,6 @@ public static class Builder { private int maxLogEntriesNumPerBatch = 1024; private int maxSizePerBatch = 16 * 1024 * 1024; private int maxPendingBatchesNum = 5; - private int shardMaxSize = 4 * 1024 * 1024; private int maxQueueLength = 4096; private long maxWaitingTimeForWaitBatchInMs = 10 * 1000L; @@ -486,11 +478,6 @@ public Replication.Builder setMaxPendingBatchesNum(int maxPendingBatchesNum) { return this; } - public Replication.Builder setShardMaxSize(int shardMaxSize) { - this.shardMaxSize = shardMaxSize; - return this; - } - public Builder setMaxQueueLength(int maxQueueLength) { this.maxQueueLength = maxQueueLength; return this; @@ -584,8 +571,7 @@ public Replication build() { regionMigrationSpeedLimitBytesPerSecond, subscriptionWalRetentionSizeInBytes, subscriptionWalRetentionTimeMs, - snapshotTransmissionProgressLogIntervalMs, - shardMaxSize); + snapshotTransmissionProgressLogIntervalMs); } } } diff --git a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java index 1ac46f1b3d2e3..304badbd8aef3 100644 --- a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java +++ b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java @@ -24,8 +24,6 @@ import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory; import org.apache.iotdb.commons.concurrent.ThreadName; -import org.apache.iotdb.commons.request.ConsensusRequestShardSource; -import org.apache.iotdb.commons.request.ShardableConsensusRequest; import org.apache.iotdb.commons.service.metric.MetricService; import org.apache.iotdb.commons.subscription.config.SubscriptionConfig; import org.apache.iotdb.consensus.common.Peer; @@ -44,8 +42,6 @@ import org.slf4j.LoggerFactory; import java.io.IOException; -import java.nio.ByteBuffer; -import java.util.Collections; import java.util.Iterator; import java.util.LinkedList; import java.util.List; @@ -256,10 +252,6 @@ public class LogDispatcherThread implements Runnable { private final CountDownLatch runFinished = new CountDownLatch(1); private volatile long lastIdleWriterSafeTimeBarrierSentTimeMs = 0L; - // An in-flight shardable request that is still being streamed out shard by shard. While this - // is non-null, the dispatcher keeps appending its remaining shards to batches before any entry - // with a larger searchIndex, so the receiver can reassemble the group in order. - private PendingShard pendingShard; public LogDispatcherThread(Peer peer, IoTConsensusConfig config, long initialSyncIndex) { this.peer = peer; @@ -465,15 +457,6 @@ public Batch getBatch() { } Batch batches = new Batch(config); - // If a shardable request is still being streamed out, finish sending its remaining shards - // first so the group stays contiguous and is never interleaved with newer entries. - drainPendingShard(batches); - if (pendingShard != null || !batches.canAccumulate()) { - // The batch is full before the in-flight group completed (or there is nothing left to - // drain into it), so send only what has been accumulated so far. - batches.buildIndex(); - return batches; - } // This condition will be executed in several scenarios: // 1. restart // 2. The getBatch() is invoked immediately at the moment the PendingEntries are consumed @@ -671,135 +654,16 @@ private boolean constructBatchFromWAL(long currentIndex, long maxIndex, Batch lo private void constructBatchIndexedFromConsensusRequest( IndexedConsensusRequest request, Batch logBatches) { - appendSerializedRequests(request, false, logBatches); - } - - /** - * Sends the remaining shards of an in-flight shardable request into {@code batch}. Stops when - * the batch is full, leaving {@link #pendingShard} non-null so the next {@link #getBatch()} - * continues from where it stopped. - */ - private void drainPendingShard(Batch batch) { - while (pendingShard != null && pendingShard.hasNext() && batch.canAccumulate()) { - batch.addTLogEntry(pendingShard.nextEntry()); - } - if (pendingShard != null && !pendingShard.hasNext()) { - pendingShard = null; - } - } - - /** - * Appends a request's serialized bytes into {@code batch}. When the request is shardable and - * its serialized form is larger than the configured shard size, the payload is split into - * multiple {@link TLogEntry}s that share the same searchIndex, each carrying a single shard, so - * it is streamed out shard by shard instead of as one oversized message. Shards that do not fit - * into this batch are kept in {@link #pendingShard}. - * - * @return true if the whole request was consumed; false if the batch filled up before the - * request was fully appended (leftover shards are retained in {@link #pendingShard}). - */ - private boolean appendSerializedRequests( - IndexedConsensusRequest data, boolean fromWAL, Batch logBatches) { - final int shardMaxSize = config.getReplication().getShardMaxSize(); - if (isShardableAndOversized(data, shardMaxSize)) { - final long shardGroupId = - ((ShardableConsensusRequest) data.getRequests().get(0)).shardGroupId(); - final PendingShard group = - new PendingShard( - new ConsensusRequestShardSource( - toSingleBuffer(data.getSerializedRequests()), shardGroupId, shardMaxSize), - data.getSearchIndex(), - fromWAL, - data.getRoutingEpoch(), - data.getPhysicalTime()); - while (group.hasNext() && logBatches.canAccumulate()) { - logBatches.addTLogEntry(group.nextEntry()); - } - if (group.hasNext()) { - pendingShard = group; - return false; - } - return true; - } - addLogEntry(data.getSerializedRequests(), data, fromWAL, data.getMemorySize(), logBatches); - return true; - } - - private boolean isShardableAndOversized(IndexedConsensusRequest data, int shardMaxSize) { - if (data.getRequests().isEmpty() - || !(data.getRequests().get(0) instanceof ShardableConsensusRequest)) { - return false; - } - long total = 0; - for (ByteBuffer buffer : data.getSerializedRequests()) { - total += buffer.remaining(); - } - return total > shardMaxSize; - } - - private void addLogEntry( - List data, - IndexedConsensusRequest request, - boolean fromWAL, - long memorySize, - Batch logBatches) { - TLogEntry logEntry = new TLogEntry(data, request.getSearchIndex(), fromWAL, memorySize); + TLogEntry logEntry = + new TLogEntry( + request.getSerializedRequests(), + request.getSearchIndex(), + false, + request.getMemorySize()); logEntry.setRoutingEpoch(request.getRoutingEpoch()); logEntry.setPhysicalTime(request.getPhysicalTime()); logBatches.addTLogEntry(logEntry, request.containsUserData()); } - - private static ByteBuffer toSingleBuffer(List buffers) { - if (buffers.size() == 1) { - return buffers.get(0); - } - int total = 0; - for (ByteBuffer buffer : buffers) { - total += buffer.remaining(); - } - final ByteBuffer merged = ByteBuffer.allocate(total); - for (ByteBuffer buffer : buffers) { - merged.put(buffer.duplicate()); - } - merged.flip(); - return merged; - } - } - - /** In-flight metadata of a shardable request that is still being streamed out. */ - private static class PendingShard { - - private final ConsensusRequestShardSource source; - private final long searchIndex; - private final boolean fromWAL; - private final long routingEpoch; - private final long physicalTime; - - private PendingShard( - ConsensusRequestShardSource source, - long searchIndex, - boolean fromWAL, - long routingEpoch, - long physicalTime) { - this.source = source; - this.searchIndex = searchIndex; - this.fromWAL = fromWAL; - this.routingEpoch = routingEpoch; - this.physicalTime = physicalTime; - } - - private boolean hasNext() { - return source.hasNext(); - } - - private TLogEntry nextEntry() { - final ByteBuffer shard = source.next(); - final TLogEntry entry = - new TLogEntry(Collections.singletonList(shard), searchIndex, fromWAL, shard.remaining()); - entry.setRoutingEpoch(routingEpoch); - entry.setPhysicalTime(physicalTime); - return entry; - } } public static AtomicLong getReceiverMemSizeSum() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java index 96bbb6ccc16ba..5ca2d565632ab 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java @@ -20,6 +20,9 @@ package org.apache.iotdb.db.consensus.statemachine.dataregion; import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.consensus.index.ProgressIndexType; import org.apache.iotdb.commons.exception.IllegalPathException; import org.apache.iotdb.commons.exception.MetadataLeaseFencedException; import org.apache.iotdb.commons.exception.SemanticException; @@ -34,6 +37,7 @@ import org.apache.iotdb.db.i18n.DataNodeMiscMessages; import org.apache.iotdb.db.i18n.DataNodePipeMessages; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedDeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedInsertNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; @@ -51,10 +55,13 @@ import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.tsfile.exception.write.PageException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.HashMap; import java.util.Map; public class DataExecutionVisitor implements PlanVisitor { @@ -65,6 +72,46 @@ public TSStatus visitPlan(PlanNode node, DataRegion context) { return null; } + @Override + public TSStatus visitLoadTsFileConsensus( + final LoadTsFileConsensusNode node, final DataRegion dataRegion) { + switch (node.getOp()) { + case BEGIN: + return StatusUtils.OK; + case PREPARE: + return StatusUtils.OK; + case PIECE: + try { + dataRegion.writeLoadTsFilePiece(node.getLoadId(), node.getTsFileDataList()); + return StatusUtils.OK; + } catch (final IOException | PageException e) { + LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); + return RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR, e.getMessage()); + } + case ABORT: + dataRegion.rollbackLoadTsFile(node.getLoadId()); + return StatusUtils.OK; + case COMMIT: + try { + final Map progressIndexes = new HashMap<>(); + node.getTimePartition2ProgressIndex() + .forEach( + (slot, bytes) -> + progressIndexes.put( + slot, ProgressIndexType.deserializeFrom(ByteBuffer.wrap(bytes)))); + return dataRegion.commitLoadTsFile( + node.getLoadId(), node.isGeneratedByPipe(), progressIndexes) + ? StatusUtils.OK + : RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR); + } catch (final Exception e) { + LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); + return RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR, e.getMessage()); + } + default: + return RpcUtils.getStatus(TSStatusCode.UNSUPPORTED_OPERATION); + } + } + public TSStatus visitRelationalInsertRow(RelationalInsertRowNode node, DataRegion context) { return visitInsertRow(node, context); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java index 4ea0757014908..67c8b8c2ec8e0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java @@ -23,6 +23,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.CountSchemaMergeNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.DeviceSchemaFetchScanNode; @@ -173,6 +174,8 @@ public PlanNode deserializeFromWAL(DataInputStream stream) throws IOException { return RelationalDeleteDataNode.deserializeFromWAL(stream); case 2004: return ObjectNode.deserializeFromWAL(stream); + case 2010: + return LoadTsFileConsensusNode.deserializeFromWAL(stream); default: throw new IllegalArgumentException(DataNodeQueryMessages.INVALID_NODE_TYPE + nodeType); } @@ -202,6 +205,8 @@ public PlanNode deserializeFromWAL(ByteBuffer buffer) { return RelationalDeleteDataNode.deserializeFromWAL(buffer); case 2004: return ObjectNode.deserialize(buffer); + case 2010: + return LoadTsFileConsensusNode.deserializeFromWAL(buffer); default: throw new IllegalArgumentException(DataNodeQueryMessages.INVALID_NODE_TYPE + nodeType); } @@ -484,6 +489,8 @@ public PlanNode deserialize(ByteBuffer buffer, short nodeType) { return RelationalDeleteDataNode.deserialize(buffer); case 2004: return ObjectNode.deserialize(buffer); + case 2010: + return LoadTsFileConsensusNode.deserialize(buffer); default: return super.deserialize(buffer, nodeType); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java index 43afdc5077f03..788d248af6200 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.queryengine.plan.planner.plan.node; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.ICoreQueryPlanVisitor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.CountSchemaMergeNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.DeviceSchemaFetchScanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.DevicesCountNode; @@ -148,6 +149,10 @@ @SuppressWarnings("java:S6539") // suppress "Monster class" warning public interface PlanVisitor extends ICoreQueryPlanVisitor { + default R visitLoadTsFileConsensus(LoadTsFileConsensusNode node, C context) { + return visitPlan(node, context); + } + ///////////////////////////////////////////////////////////////////////////////////////////////// // Data Query Node ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java index da8ffe144f500..3fc87e2c42e04 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java @@ -27,7 +27,6 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; -import org.apache.iotdb.commons.request.ShardableConsensusRequest; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; @@ -36,7 +35,6 @@ import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.IWALByteBufferView; import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryType; import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryValue; -import org.apache.iotdb.db.storageengine.load.LoadTsFileManager; import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; import org.apache.tsfile.exception.NotImplementedException; @@ -47,10 +45,8 @@ import java.io.ByteArrayOutputStream; import java.io.DataInputStream; import java.io.DataOutputStream; -import java.io.File; import java.io.IOException; import java.io.InputStream; -import java.io.RandomAccessFile; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; @@ -63,8 +59,7 @@ * Consensus-backed LOAD request carrying Begin / Piece / Seal / Commit phases. Submitted once to * the DataRegion write peer (or Ratis leader); replicas apply via the consensus state machine. */ -public class LoadTsFileConsensusNode extends SearchNode - implements WALEntryValue, ShardableConsensusRequest { +public class LoadTsFileConsensusNode extends SearchNode implements WALEntryValue { private LoadTsFileConsensusOp op; private String loadId; @@ -73,7 +68,6 @@ public class LoadTsFileConsensusNode extends SearchNode private String database; private int expectedPieceCount = -1; private long pieceIndex = -1; - private long pieceOffset = -1; private List pieceRefs = new ArrayList<>(); private long dataSize; private long checksum; @@ -116,22 +110,15 @@ public static LoadTsFileConsensusNode piece( String loadId, String tsFileId, long pieceIndex, - long pieceOffset, - List tsFileDataList, - long checksum) { + List tsFileDataList) { final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); node.op = LoadTsFileConsensusOp.PIECE; node.loadId = loadId; node.tsFileId = tsFileId; node.pieceIndex = pieceIndex; - node.pieceOffset = pieceOffset; node.tsFileDataList = tsFileDataList == null ? new ArrayList<>() : new ArrayList<>(tsFileDataList); node.dataSize = node.tsFileDataList.stream().mapToLong(TsFileData::getDataSize).sum(); - // The checksum is part of the consensus contract and must be preserved exactly as supplied by - // the caller. Recomputing it here silently changes the value for callers that use a staged - // piece checksum (and made marker/piece validation disagree). - node.checksum = checksum; return node; } @@ -326,10 +313,6 @@ public long getPieceIndex() { return pieceIndex; } - public long getPieceOffset() { - return pieceOffset; - } - public String getRelativePath() { return pieceRefs.isEmpty() ? null : pieceRefs.get(0).relativePath; } @@ -494,7 +477,6 @@ private void serializeBody(DataOutputStream stream, boolean includeContent) thro ReadWriteIOUtils.write(database == null ? "" : database, stream); ReadWriteIOUtils.write(expectedPieceCount, stream); ReadWriteIOUtils.write(pieceIndex, stream); - ReadWriteIOUtils.write(pieceOffset, stream); ReadWriteIOUtils.write(dataSize, stream); ReadWriteIOUtils.write(checksum, stream); ReadWriteIOUtils.write(pieceCount, stream); @@ -517,7 +499,11 @@ private void serializeBody(DataOutputStream stream, boolean includeContent) thro data.serialize(stream); } ReadWriteIOUtils.write(timePartition2ProgressIndex.size(), stream); - for (Map.Entry entry : timePartition2ProgressIndex.entrySet()) { + final List> sortedProgressIndexes = + new ArrayList<>(timePartition2ProgressIndex.entrySet()); + sortedProgressIndexes.sort( + java.util.Comparator.comparingLong(entry -> entry.getKey().getStartTime())); + for (Map.Entry entry : sortedProgressIndexes) { ReadWriteIOUtils.write(entry.getKey().getStartTime(), stream); ReadWriteIOUtils.write(entry.getValue().length, stream); stream.write(entry.getValue()); @@ -554,7 +540,6 @@ private static LoadTsFileConsensusNode deserializeBody(InputStream stream, boole node.database = ReadWriteIOUtils.readString(stream); node.expectedPieceCount = ReadWriteIOUtils.readInt(stream); node.pieceIndex = ReadWriteIOUtils.readLong(stream); - node.pieceOffset = ReadWriteIOUtils.readLong(stream); node.dataSize = ReadWriteIOUtils.readLong(stream); node.checksum = ReadWriteIOUtils.readLong(stream); node.pieceCount = ReadWriteIOUtils.readInt(stream); @@ -562,6 +547,9 @@ private static LoadTsFileConsensusNode deserializeBody(InputStream stream, boole node.isGeneratedByPipe = ReadWriteIOUtils.readBool(stream); node.deleteAfterLoad = ReadWriteIOUtils.readBool(stream); final int refCount = ReadWriteIOUtils.readInt(stream); + if (refCount < 0) { + throw new IOException(); + } node.pieceRefs = new ArrayList<>(refCount); for (int i = 0; i < refCount; i++) { final String refPath = ReadWriteIOUtils.readString(stream); @@ -570,6 +558,9 @@ private static LoadTsFileConsensusNode deserializeBody(InputStream stream, boole byte[] content = null; if (readContent) { final int contentLength = ReadWriteIOUtils.readInt(stream); + if (contentLength < 0) { + throw new IOException(); + } if (contentLength > 0) { content = new byte[contentLength]; int offset = 0; @@ -587,15 +578,24 @@ private static LoadTsFileConsensusNode deserializeBody(InputStream stream, boole node.pieceRefs.add(new PieceRef(refPath, refOffset, refSize, content)); } final int dataCount = ReadWriteIOUtils.readInt(stream); + if (dataCount < 0) { + throw new IOException(); + } node.tsFileDataList = new ArrayList<>(dataCount); for (int i = 0; i < dataCount; i++) { node.tsFileDataList.add(TsFileData.deserialize(stream)); } final int progressCount = ReadWriteIOUtils.readInt(stream); + if (progressCount < 0) { + throw new IOException(); + } node.timePartition2ProgressIndex = new HashMap<>(progressCount); for (int i = 0; i < progressCount; i++) { final long startTime = ReadWriteIOUtils.readLong(stream); final int len = ReadWriteIOUtils.readInt(stream); + if (len < 0) { + throw new IOException(); + } final byte[] bytes = new byte[len]; int offset = 0; while (offset < len) { @@ -647,38 +647,12 @@ public ByteBuffer serialize() { } } - /** Stable, collision-resistant identifier of the logical request this node belongs to. */ - @Override - public long shardGroupId() { - return Objects.hash(loadId, tsFileId, pieceIndex); - } - private byte[] readPieceContent(PieceRef ref) { if (ref.content != null) { return ref.content; } - // PieceRef construction already validated offset/size bounds (non-negative, no overflow, - // within int range), so the narrowing cast below is safe. - final byte[] content = new byte[(int) ref.size]; - if (content.length == 0) { - return content; - } - final File file = - LoadTsFileManager.findLoadTsFile(ref.relativePath) - .orElseThrow( - () -> - new IllegalStateException( - DataNodeQueryMessages - .EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 - + "piece file not found: " - + ref.relativePath)); - try (RandomAccessFile raf = new RandomAccessFile(file, "r")) { - raf.seek(ref.offset); - raf.readFully(content); - return content; - } catch (IOException e) { - throw new IllegalStateException(e); - } + throw new IllegalStateException( + DataNodeQueryMessages.EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkLayout.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkLayout.java new file mode 100644 index 0000000000000..da2a5fbc95f56 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkLayout.java @@ -0,0 +1,27 @@ +/* + * 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 + */ +package org.apache.iotdb.db.queryengine.plan.scheduler.load; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.read.common.Chunk; + +/** Self-contained physical layout descriptor for one Chunk. */ +public record ChunkLayout( + IDeviceID device, + boolean aligned, + long chunkGroupIndex, + long chunkGroupHeaderOffset, + boolean firstChunkOfGroup, + Chunk chunk, + long offset, + long length, + int chunkIndexInGroup) {} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkOffsetCalculator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkOffsetCalculator.java new file mode 100644 index 0000000000000..3e3c9b7649022 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkOffsetCalculator.java @@ -0,0 +1,144 @@ +/* + * 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 + */ +package org.apache.iotdb.db.queryengine.plan.scheduler.load; + +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; + +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.file.header.ChunkGroupHeader; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.utils.BytesUtils; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.List; +import java.util.Objects; + +/** Calculates absolute TsFile data-zone offsets for independently distributed Chunks. */ +public final class ChunkOffsetCalculator { + + private static final long FILE_HEADER_SIZE = + BytesUtils.stringToBytes(TSFileConfig.MAGIC_STRING).length + Byte.BYTES; + + private long nextOffset; + private IDeviceID currentDevice; + private long chunkGroupIndex = -1; + private long chunkGroupHeaderOffset = -1; + private int nextChunkIndex; + + public ChunkOffsetCalculator() { + this.nextOffset = FILE_HEADER_SIZE; + } + + public ChunkLayout allocate(final IDeviceID device, final boolean aligned, final Chunk chunk) { + Objects.requireNonNull(device); + Objects.requireNonNull(chunk); + + if (!Objects.equals(currentDevice, device)) { + currentDevice = device; + chunkGroupIndex++; + chunkGroupHeaderOffset = nextOffset; + nextOffset += getChunkGroupHeaderSize(device); + nextChunkIndex = 0; + } + + final long chunkOffset = nextOffset; + final long chunkLength = + chunk.getHeader().getSerializedSize() + (long) chunk.getData().remaining(); + final boolean firstChunkOfGroup = nextChunkIndex == 0; + final ChunkLayout layout = + new ChunkLayout( + device, + aligned, + chunkGroupIndex, + chunkGroupHeaderOffset, + firstChunkOfGroup, + chunk, + chunkOffset, + chunkLength, + nextChunkIndex); + + nextOffset += chunkLength; + nextChunkIndex++; + return layout; + } + + public ChunkData.ChunkLayout assign(final ChunkData chunkData, final Chunk chunk) { + Objects.requireNonNull(chunkData); + Objects.requireNonNull(chunk); + + final ChunkLayout allocated = allocate(chunkData.getDevice(), chunkData.isAligned(), chunk); + final ChunkData.ChunkLayout layout = + new ChunkData.ChunkLayout( + allocated.chunkGroupIndex(), + allocated.chunkGroupHeaderOffset(), + allocated.offset(), + allocated.length(), + allocated.chunkIndexInGroup(), + allocated.firstChunkOfGroup()); + chunkData.setChunkLayout(layout); + return layout; + } + + public ChunkData.ChunkLayout assign(final ChunkData chunkData) { + Objects.requireNonNull(chunkData); + final List chunks = chunkData.getChunks(); + if (chunks.isEmpty()) { + throw new IllegalArgumentException("ChunkData contains no physical chunk"); + } + + final ChunkLayout first = allocate(chunkData.getDevice(), chunkData.isAligned(), chunks.get(0)); + long totalLength = first.length(); + for (int i = 1; i < chunks.size(); i++) { + totalLength += allocate(chunkData.getDevice(), chunkData.isAligned(), chunks.get(i)).length(); + } + final ChunkData.ChunkLayout layout = + new ChunkData.ChunkLayout( + first.chunkGroupIndex(), + first.chunkGroupHeaderOffset(), + first.offset(), + totalLength, + first.chunkIndexInGroup(), + first.firstChunkOfGroup()); + chunkData.setChunkLayout(layout); + return layout; + } + + private int getChunkGroupHeaderSize(final IDeviceID device) { + try (ByteArrayOutputStream output = new ByteArrayOutputStream()) { + return new ChunkGroupHeader(device).serializeTo(output); + } catch (final IOException e) { + throw new IllegalStateException(e); + } + } + + public long getNextOffset() { + return nextOffset; + } + + public long getChunkGroupIndex() { + return chunkGroupIndex; + } + + public long getChunkGroupHeaderOffset() { + return chunkGroupHeaderOffset; + } + + public void reset() { + nextOffset = FILE_HEADER_SIZE; + currentDevice = null; + chunkGroupIndex = -1; + chunkGroupHeaderOffset = -1; + nextChunkIndex = 0; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java index d46936d1bc93b..f8f56c7fedab7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java @@ -75,7 +75,7 @@ * | DataPartitionRouter -> MemoryBoundedBuffer * | -> PieceDispatcher * | phase2: BEGIN -> PIECE* -> PREPARE -> COMMIT - * | or ABORT, via RegionConsensusContext + + * | or ABORT, via per-region state + * | LoadConsensusSubmitter * | * +--> success -> register pending deletion (the source file is kept @@ -112,7 +112,7 @@ * +-- later pieces: PIECE(1), PIECE(2), ... * | * v - * RegionConsensusContext.accumulate(bytes, checksum) + * Accumulate per-region piece count and bytes * | * v * LoadConsensusSubmitter (submit to the partition write node; bounded retry) @@ -127,7 +127,7 @@ * |---- PIECE(0, chunks) --------------->| append chunks * |---- PIECE(1, chunks) --------------->| append chunks * |---- ... | - * |---- PREPARE(count, bytes, checksum)->| seal staged TsFile + * |---- PREPARE(count, bytes)----------->| seal staged TsFile * |---- COMMIT ------------------------->| load staged TsFile * | | * on failure: @@ -178,9 +178,9 @@ * dispatcher into the route -> buffer -> dispatch pipeline * * - * {@link RegionConsensusContext} + * Per-region maps in {@link TwoPhaseConsensusLoadStrategy} * LOAD per-region two-phase state: one context per region with load id, piece count, total - * bytes, XOR checksum and BEGIN state + * bytes and BEGIN state * * * {@link LoadConsensusSubmitter} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java index 2c5674f2dd111..b08714539e8e7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java @@ -27,10 +27,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; -import org.apache.iotdb.db.storageengine.load.splitter.ChunkEncoder; import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; -import org.apache.iotdb.db.storageengine.load.splitter.EncodedChunk; -import org.apache.iotdb.db.storageengine.load.splitter.EncodedChunkGroup; import org.apache.tsfile.utils.Pair; import org.slf4j.Logger; @@ -90,8 +87,9 @@ interface DispatchCallback { private final Map> regionId2ReplicaSetAndNode = new HashMap<>(); - private final Map partition2NextPieceIndex = new HashMap<>(); - private final Map partition2NextOffset = new HashMap<>(); + private final Map region2NextPieceIndex = new HashMap<>(); + private final Map partition2OffsetCalculator = + new HashMap<>(); /** * Max-heap of (regionId, buffered piece size at offer time) used to dispatch the largest pieces @@ -114,9 +112,8 @@ interface DispatchCallback { void offerChunk(ChunkData chunkData, TRegionReplicaSet replicaSet) throws LoadFileException, IOException { final TConsensusGroupId regionId = replicaSet.getRegionId(); - final EncodedChunkGroup encodedChunkGroup = ChunkEncoder.encode(chunkData); final PartitionKey key = - new PartitionKey(regionId, encodedChunkGroup.getTimePartitionSlot().getStartTime()); + new PartitionKey(regionId, chunkData.getTimePartitionSlot().getStartTime()); if (regionId2ReplicaSetAndNode.containsKey(key) && !Objects.equals(regionId2ReplicaSetAndNode.get(key).getLeft(), replicaSet)) { // Detected region replica set changed (maybe due to region migration), throw an exception @@ -124,22 +121,15 @@ void offerChunk(ChunkData chunkData, TRegionReplicaSet replicaSet) regionId2ReplicaSetAndNode.get(key).getLeft(), replicaSet); } - long nextOffset = partition2NextOffset.getOrDefault(key, 0L); - for (EncodedChunk encodedChunk : encodedChunkGroup.getChunks()) { - encodedChunk.setOffset(nextOffset); - nextOffset += encodedChunk.getBytes().length; - } - partition2NextOffset.put(key, nextOffset); - final long memoryDelta = encodedChunkGroup.getDataSize() - chunkData.getDataSize(); - if (memoryDelta > 0) { - memoryBuffer.add(memoryDelta); - } else if (memoryDelta < 0) { - memoryBuffer.release(-memoryDelta); - } + ChunkOffsetCalculator calculator = + partition2OffsetCalculator.computeIfAbsent(key, ignored -> new ChunkOffsetCalculator()); + + calculator.assign(chunkData); + regionId2ReplicaSetAndNode .computeIfAbsent(key, o -> new Pair<>(replicaSet, newPieceNode(key))) .getRight() - .addTsFileData(encodedChunkGroup); + .addTsFileData(chunkData); offerPieceRegion(key); } @@ -220,8 +210,10 @@ private void replacePieceNode(PartitionKey key, TRegionReplicaSet replicaSet) { } private LoadTsFilePieceNode newPieceNode(PartitionKey key) { - final long pieceIndex = partition2NextPieceIndex.getOrDefault(key, 0L); - partition2NextPieceIndex.put(key, pieceIndex + 1); + // pieceIndex is a consensus request identity, not a physical file offset. It must be unique + // across every time partition routed to the same Region/loadId. + final long pieceIndex = region2NextPieceIndex.getOrDefault(key.regionId(), 0L); + region2NextPieceIndex.put(key.regionId(), pieceIndex + 1); return new LoadTsFilePieceNode( singleTsFileNode.getPlanNodeId(), singleTsFileNode.getTsFileResource().getTsFile(), @@ -254,8 +246,8 @@ private PartitionKey pollLargestPieceRegion() { void clear() { regionId2ReplicaSetAndNode.clear(); - partition2NextPieceIndex.clear(); - partition2NextOffset.clear(); + region2NextPieceIndex.clear(); + partition2OffsetCalculator.clear(); largestPieceRegions.clear(); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java deleted file mode 100644 index eaba46488aab8..0000000000000 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java +++ /dev/null @@ -1,93 +0,0 @@ -/* - * 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.queryengine.plan.scheduler.load; - -import org.apache.iotdb.db.storageengine.load.LoadTsFileChecksumUtils; - -import java.util.UUID; - -/** - * One high-cohesion state object per target region of the LOAD consensus two-phase protocol. It - * replaces the five parallel maps ({@code regionPieceCount}, {@code regionPieceTotalBytes}, {@code - * regionPieceChecksum}, {@code begunConsensusRegions}, {@code regionLoadId}) that used to live in - * {@link LoadTsFileScheduler}, so the load id, the BEGIN state and the three accumulated counters - * can never diverge. - * - *

Lifecycle, driven by {@link TwoPhaseConsensusLoadStrategy}: - * - *

    - *
  • Created lazily when the first piece of a region is dispatched ({@code - * consensusContexts.computeIfAbsent}). - *
  • {@link #markBegun()} runs before the BEGIN command is submitted, so a region is begun at - * most once per load. - *
  • {@link #accumulate(long, long)} runs after every successful PIECE: it increments the piece - * count, adds the piece bytes and XORs the piece checksum. - *
  • Phase two reads {@link #getPieceCount()}, {@link #getTotalBytes()} and {@link - * #getChecksum()} to build the PREPARE command. - *
- * - * The load id is generated per instance, so every region of the same source file is isolated from - * the others on the write nodes. - */ -public class RegionConsensusContext { - - /** Each region gets its own load id so its staged data is isolated from other regions. */ - private final String loadId = UUID.randomUUID().toString(); - - private long pieceCount = 0; - private long totalBytes = 0; - private long checksum = 0; - - /** Whether the BEGIN command has already been sent to this region. */ - private boolean begun = false; - - public String getLoadId() { - return loadId; - } - - public long getPieceCount() { - return pieceCount; - } - - public long getTotalBytes() { - return totalBytes; - } - - public long getChecksum() { - return checksum; - } - - public boolean isBegun() { - return begun; - } - - public void markBegun() { - begun = true; - } - - /** Records one successfully applied piece: increments the count, adds bytes and XORs checksum. */ - public void accumulate(long bytes, long pieceChecksum) { - pieceCount++; - totalBytes += bytes; - // Order-sensitive aggregation (rotate-then-XOR): the aggregate must depend on the exact piece - // order, because the plain XOR used before could not detect reordered or swapped pieces. - checksum = LoadTsFileChecksumUtils.combine(checksum, pieceChecksum); - } -} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java index 2fe9e25d0a91e..34dc055f6d452 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java @@ -30,6 +30,7 @@ import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; +import java.io.IOException; import java.util.ArrayList; import java.util.List; import java.util.function.Consumer; @@ -118,7 +119,11 @@ private void routeChunkData() throws LoadFileException { final List replicaSets = router.route(nonDirectionalChunkData); for (int i = 0, size = nonDirectionalChunkData.size(); i < size; i++) { - dispatcher.offerChunk(nonDirectionalChunkData.get(i), replicaSets.get(i)); + try { + dispatcher.offerChunk(nonDirectionalChunkData.get(i), replicaSets.get(i)); + } catch (final IOException e) { + throw new LoadFileException(e); + } } nonDirectionalChunkData.clear(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java index 5a40933c4c4d7..8d51783add037 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java @@ -31,7 +31,6 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; -import org.apache.iotdb.db.storageengine.load.LoadTsFileChecksumUtils; import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; @@ -57,18 +56,15 @@ * correlation) and feeds the source TsFile through {@link TsFileSplitter} into {@link * TsFileSplitConsumer}. Every dispatched piece goes through {@code dispatchConsensusPiece}: the * first piece of a region sends BEGIN with a fresh per-region load id, then PIECE commands with a - * monotonically increasing {@code pieceIndex}; {@link RegionConsensusContext#accumulate(long, - * long)} records piece count, total bytes and the XOR checksum. Submission goes through {@link + * monotonically increasing {@code pieceIndex}. Submission goes through {@link * LoadConsensusSubmitter} with bounded retries for transient failures only. * *

Phase 2 (commit or abort). If every region received all its pieces, each touched region - * gets PREPARE (with the accumulated count/bytes/checksum) followed by COMMIT; otherwise every - * touched region gets ABORT so the staged data is dropped. + * gets PREPARE (with the accumulated count/bytes) followed by COMMIT; otherwise every touched + * region gets ABORT so the staged data is dropped. * - *

Per-file state: {@code allReplicaSets} (the regions to prepare/commit/abort), {@code - * consensusContexts} (per-region two-phase state) and {@code timePartitionSlotToProgressIndex} - * (pipe progress index per time partition, collected while splitting for the upcoming - * progress-index sync). + *

Per-file state tracks the touched regions, their load ids, BEGIN state, piece counters and + * pipe progress indexes. */ public class TwoPhaseConsensusLoadStrategy implements TsFileLoadStrategy { @@ -100,9 +96,10 @@ public class TwoPhaseConsensusLoadStrategy implements TsFileLoadStrategy { /** Regions touched by the current file; used to send ABORT/PREPARE+COMMIT in phase two. */ private final Set allReplicaSets = new HashSet<>(); - /** Per-region two-phase state of the current file; replaces the old five parallel maps. */ - private final Map consensusContexts = - new ConcurrentHashMap<>(); + private final Map regionLoadIds = new ConcurrentHashMap<>(); + private final Set begunRegions = ConcurrentHashMap.newKeySet(); + private final Map regionPieceCounts = new ConcurrentHashMap<>(); + private final Map regionTotalBytes = new ConcurrentHashMap<>(); /** * Progress index per time partition, assigned while the file is being split. Kept for the @@ -131,7 +128,10 @@ public boolean execute(LoadSingleTsFileNode node) { this.currentNode = node; dispatcher.setUuid(UUID.randomUUID().toString()); allReplicaSets.clear(); - consensusContexts.clear(); + regionLoadIds.clear(); + begunRegions.clear(); + regionPieceCounts.clear(); + regionTotalBytes.clear(); timePartitionSlotToProgressIndex.clear(); long startTime = System.nanoTime(); @@ -195,8 +195,8 @@ private boolean dispatchOnePieceNode( /** * Submits a LOAD consensus request with a bounded number of attempts. Only transient failures are - * retried; permanent rejections (checksum mismatch, missing staged writer) are returned to the - * caller immediately so the scheduler can abort. + * retried; permanent rejections are returned to the caller immediately so the scheduler can + * abort. */ private TSStatus submitConsensusWithRetry( TRegionReplicaSet replicaSet, LoadTsFileConsensusNode node) { @@ -241,12 +241,10 @@ private boolean isTransientConsensusFailure(TSStatus status) { private boolean dispatchConsensusPiece( LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { final TConsensusGroupId regionId = replicaSet.getRegionId(); - final RegionConsensusContext context = - consensusContexts.computeIfAbsent(regionId, o -> new RegionConsensusContext()); - final String loadId = context.getLoadId(); + final String loadId = + regionLoadIds.computeIfAbsent(regionId, ignored -> UUID.randomUUID().toString()); - if (!context.isBegun()) { - context.markBegun(); + if (begunRegions.add(regionId)) { final LoadTsFileConsensusNode begin = LoadTsFileConsensusNode.begin( new PlanNodeId("load-begin-" + loadId), @@ -272,20 +270,13 @@ private boolean dispatchConsensusPiece( } final long pieceIndex = pieceNode.getPieceIndex(); - // The checksum is part of the consensus contract: both the write node and every follower - // verify it, and the per-piece digest includes the piece index so reordered payloads are - // detected. Sending a constant here would silently disable checksum validation. - final long checksum = - LoadTsFileChecksumUtils.checksum(pieceIndex, pieceNode.getAllTsFileData()); final LoadTsFileConsensusNode piece = LoadTsFileConsensusNode.piece( new PlanNodeId("load-piece-" + loadId + "-" + pieceIndex), loadId, pieceNode.getTsFile() == null ? null : pieceNode.getTsFile().getName(), pieceIndex, - 0L, - pieceNode.getAllTsFileData(), - checksum); + pieceNode.getAllTsFileData()); final TSStatus pieceStatus = submitConsensusWithRetry(replicaSet, piece); if (pieceStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { LOGGER.warn( @@ -298,7 +289,8 @@ private boolean dispatchConsensusPiece( pieceNode); return false; } - context.accumulate(piece.getDataSize(), piece.getChecksum()); + regionPieceCounts.merge(regionId, 1L, Long::sum); + regionTotalBytes.merge(regionId, piece.getDataSize(), Long::sum); return true; } @@ -311,7 +303,7 @@ private boolean secondPhase(LoadSingleTsFileNode node, boolean isFirstPhaseSucce private boolean abortAllRegions() { for (TRegionReplicaSet replicaSet : allReplicaSets) { - final String loadId = consensusContexts.get(replicaSet.getRegionId()).getLoadId(); + final String loadId = regionLoadIds.get(replicaSet.getRegionId()); final LoadTsFileConsensusNode abort = LoadTsFileConsensusNode.abort( new PlanNodeId("load-abort-" + loadId), loadId, null, isGeneratedByPipe); @@ -336,16 +328,16 @@ private boolean prepareAndCommitAllRegions(LoadSingleTsFileNode node) { final Map timePartition2ProgressIndex = serializeTimePartitionProgressIndexes(); for (TRegionReplicaSet replicaSet : allReplicaSets) { - final RegionConsensusContext context = consensusContexts.get(replicaSet.getRegionId()); - final String loadId = context.getLoadId(); + final TConsensusGroupId regionId = replicaSet.getRegionId(); + final String loadId = regionLoadIds.get(regionId); final LoadTsFileConsensusNode prepare = LoadTsFileConsensusNode.prepare( new PlanNodeId("load-prepare-" + loadId), loadId, null, - (int) context.getPieceCount(), - context.getTotalBytes(), - context.getChecksum(), + Math.toIntExact(regionPieceCounts.getOrDefault(regionId, 0L)), + regionTotalBytes.getOrDefault(regionId, 0L), + 0L, timePartition2ProgressIndex); final TSStatus prepareStatus = consensusSubmitter.submit(replicaSet, prepare); if (prepareStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java index 1ac6c15488b11..f7f4a60460930 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java @@ -54,6 +54,7 @@ import org.apache.iotdb.db.exception.StorageEngineException; import org.apache.iotdb.db.exception.TsFileProcessorException; import org.apache.iotdb.db.exception.WriteProcessRejectException; +import org.apache.iotdb.db.exception.load.LoadFileException; import org.apache.iotdb.db.exception.load.LoadReadOnlyException; import org.apache.iotdb.db.exception.runtime.StorageEngineFailureException; import org.apache.iotdb.db.i18n.StorageEngineMessages; @@ -79,6 +80,7 @@ import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALException; import org.apache.iotdb.db.storageengine.dataregion.wal.recover.WALRecoverManager; import org.apache.iotdb.db.storageengine.load.LoadTsFileManager; +import org.apache.iotdb.db.storageengine.load.active.ActiveLoadAgent; import org.apache.iotdb.db.storageengine.load.limiter.LoadTsFileRateLimiter; import org.apache.iotdb.db.storageengine.rescon.disk.TierManager; import org.apache.iotdb.db.storageengine.rescon.memory.SystemInfo; @@ -108,6 +110,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.concurrent.Callable; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; @@ -169,7 +172,7 @@ public class StorageEngine implements IService { private final List customFlushListeners = new ArrayList<>(); private int recoverDataRegionNum = 0; - private final LoadTsFileManager loadTsFileManager = new LoadTsFileManager(); + private final ActiveLoadAgent activeLoadAgent = new ActiveLoadAgent(); public final AtomicLong objectFileId = new AtomicLong(0); @@ -345,7 +348,18 @@ public void start() throws StartupException { } asyncRecoverTsFileResource(); - loadTsFileManager.start(); + for (final DataRegion dataRegion : dataRegionMap.values()) { + try { + dataRegion.getLoadTsFileManager(); + } catch (final Throwable e) { + LOGGER.warn( + StorageEngineMessages.STORAGE_LOG_FAILED_TO_RECOVER_DATA_REGION_804B162D, + dataRegion.getDatabaseName(), + dataRegion.getDataRegionIdString(), + e); + } + } + activeLoadAgent.start(); } private void startTimedService() { @@ -430,7 +444,8 @@ private void asyncRecoverTsFileResource() { @Override public void stop() { - loadTsFileManager.stop(); + activeLoadAgent.stop(); + dataRegionMap.values().forEach(DataRegion::stopLoadTsFileManager); for (DataRegion dataRegion : dataRegionMap.values()) { if (dataRegion != null) { CompactionScheduleTaskManager.getInstance().unregisterDataRegion(dataRegion); @@ -449,7 +464,8 @@ public void stop() { @Override public void shutdown(long milliseconds) throws ShutdownException { - loadTsFileManager.stop(); + activeLoadAgent.stop(); + dataRegionMap.values().forEach(DataRegion::stopLoadTsFileManager); try { for (DataRegion dataRegion : dataRegionMap.values()) { if (dataRegion != null) { @@ -514,6 +530,7 @@ public DataRegion buildNewDataRegion(String databaseName, DataRegionId dataRegio /** This function is just for unit test. */ @TestOnly public synchronized void reset() { + dataRegionMap.values().forEach(DataRegion::stopLoadTsFileManager); dataRegionMap.clear(); } @@ -831,6 +848,7 @@ public TSStatus deleteDataRegion(DataRegionId regionId) { deletingDataRegionMap.computeIfAbsent(regionId, k -> dataRegionMap.remove(regionId)); if (region != null) { LOGGER.info(StorageEngineMessages.REMOVING_DATA_REGION, regionId); + region.stopLoadTsFileManager(); region.markDeleted(); try { region.abortCompaction(); @@ -952,6 +970,7 @@ public DataRegion setDataRegionForSnapshotLoad( DataRegionId regionId, Supplier newRegionSupplier) { if (dataRegionMap.containsKey(regionId)) { DataRegion oldRegion = dataRegionMap.get(regionId); + oldRegion.stopLoadTsFileManager(); oldRegion.markDeleted(); oldRegion.abortCompaction(); oldRegion.syncCloseAllWorkingTsFileProcessors(); @@ -978,6 +997,7 @@ public DataRegion setDataRegionForSnapshotLoad( public void setDataRegion(DataRegionId regionId, DataRegion newRegion) { if (dataRegionMap.containsKey(regionId)) { DataRegion oldRegion = dataRegionMap.get(regionId); + oldRegion.stopLoadTsFileManager(); oldRegion.markDeleted(); oldRegion.abortCompaction(); oldRegion.syncCloseAllWorkingTsFileProcessors(); @@ -1040,7 +1060,7 @@ public TSStatus writeLoadTsFileNode( } try { - loadTsFileManager.writeToDataRegion(dataRegion, pieceNode, uuid); + dataRegion.getLoadTsFileManager().writeToDataRegion(pieceNode, uuid); } catch (IOException | PageException e) { LOGGER.warn( StorageEngineMessages @@ -1076,7 +1096,7 @@ public TSStatus executeLoadCommand( try { switch (loadCommand) { case EXECUTE: - if (loadTsFileManager.loadAll(uuid, isGeneratedByPipe, timePartitionProgressIndexMap)) { + if (loadAllInDataRegions(uuid, isGeneratedByPipe, timePartitionProgressIndexMap)) { status = RpcUtils.SUCCESS_STATUS; } else { status.setCode(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()); @@ -1089,7 +1109,7 @@ public TSStatus executeLoadCommand( } break; case ROLLBACK: - if (loadTsFileManager.deleteAll(uuid)) { + if (deleteAllInDataRegions(uuid)) { status = RpcUtils.SUCCESS_STATUS; } else { status.setCode(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()); @@ -1114,6 +1134,32 @@ public TSStatus executeLoadCommand( return status; } + private boolean loadAllInDataRegions( + final String uuid, + final boolean isGeneratedByPipe, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + boolean loaded = false; + for (final DataRegion dataRegion : dataRegionMap.values()) { + final Optional manager = dataRegion.getLoadTsFileManagerIfPresent(); + if (manager.isPresent()) { + loaded |= manager.get().loadAll(uuid, isGeneratedByPipe, timePartitionProgressIndexMap); + } + } + return loaded; + } + + private boolean deleteAllInDataRegions(final String uuid) { + boolean deleted = false; + for (final DataRegion dataRegion : dataRegionMap.values()) { + final Optional manager = dataRegion.getLoadTsFileManagerIfPresent(); + if (manager.isPresent()) { + deleted |= manager.get().deleteAll(uuid); + } + } + return deleted; + } + /** reboot timed flush sequence/unsequence memtable thread */ public void rebootTimedService() throws ShutdownException { LOGGER.info(StorageEngineMessages.START_REBOOTING_ALL_TIMED_SERVICE); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java index 6b14b43edf5ff..f7e9ca9594b5d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java @@ -23,12 +23,14 @@ import org.apache.iotdb.calc.plan.relational.metadata.CommonMetadataUtils; import org.apache.iotdb.calc.utils.ObjectTypeUtils; import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.client.exception.ClientManagerException; import org.apache.iotdb.commons.cluster.NodeStatus; import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory; import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.conf.IoTDBConstant; import org.apache.iotdb.commons.consensus.DataRegionId; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; import org.apache.iotdb.commons.exception.DiskSpaceInsufficientException; import org.apache.iotdb.commons.exception.MetadataException; import org.apache.iotdb.commons.exception.MetadataLeaseFencedException.LeaseFencedRetryPolicy; @@ -155,8 +157,10 @@ import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALMode; import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALFlushListener; import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALRecoverListener; +import org.apache.iotdb.db.storageengine.load.LoadTsFileManager; import org.apache.iotdb.db.storageengine.load.disk.ILoadDiskSelector; import org.apache.iotdb.db.storageengine.load.limiter.LoadTsFileRateLimiter; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; import org.apache.iotdb.db.storageengine.rescon.disk.TierManager; import org.apache.iotdb.db.storageengine.rescon.memory.SystemInfo; import org.apache.iotdb.db.storageengine.rescon.memory.TimePartitionInfo; @@ -176,6 +180,7 @@ import com.github.benmanes.caffeine.cache.Caffeine; import com.google.common.io.BaseEncoding; import org.apache.thrift.TException; +import org.apache.tsfile.exception.write.PageException; import org.apache.tsfile.external.commons.io.FileUtils; import org.apache.tsfile.external.commons.lang3.tuple.Triple; import org.apache.tsfile.file.metadata.ChunkMetadata; @@ -319,6 +324,9 @@ public class DataRegion implements IDataRegionForQuery { /** manage seqFileList and unSeqFileList. */ private final TsFileManager tsFileManager; + /** LOAD state is isolated by DataRegion and initialized only when this Region handles LOAD. */ + private volatile LoadTsFileManager loadTsFileManager; + /** manage tsFileResource degrade. */ private final TsFileResourceManager tsFileResourceManager = TsFileResourceManager.getInstance(); @@ -542,6 +550,49 @@ public String getDatabaseName() { return databaseName; } + public LoadTsFileManager getLoadTsFileManager() { + LoadTsFileManager manager = loadTsFileManager; + if (manager == null) { + synchronized (this) { + manager = loadTsFileManager; + if (manager == null) { + manager = new LoadTsFileManager(this); + loadTsFileManager = manager; + } + } + } + return manager; + } + + public Optional getLoadTsFileManagerIfPresent() { + return Optional.ofNullable(loadTsFileManager); + } + + public void stopLoadTsFileManager() { + final LoadTsFileManager manager = loadTsFileManager; + if (manager != null) { + manager.stop(); + loadTsFileManager = null; + } + } + + public void writeLoadTsFilePiece(final String loadId, final List tsFileDataList) + throws IOException, PageException { + getLoadTsFileManager().writePiece(loadId, tsFileDataList); + } + + public boolean commitLoadTsFile( + final String loadId, + final boolean isGeneratedByPipe, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + return getLoadTsFileManager().loadAll(loadId, isGeneratedByPipe, timePartitionProgressIndexMap); + } + + public void rollbackLoadTsFile(final String loadId) { + getLoadTsFileManager().deleteAll(loadId); + } + public boolean isTableModel() { return isTableModel; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntryType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntryType.java index 2ad72f4a2324d..dd589f4e7713a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntryType.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntryType.java @@ -48,6 +48,7 @@ public enum WALEntryType { MEMORY_TABLE_SNAPSHOT((byte) 10), RELATIONAL_DELETE_DATA_NODE((byte) 11), OBJECT_FILE_NODE((byte) 12), + LOAD_TSFILE_CONSENSUS_NODE((byte) 13), // endregion // region signal entry type // signal wal buffer has been closed @@ -75,7 +76,8 @@ public boolean needSearch() { || this == INSERT_ROWS_NODE || this == DELETE_DATA_NODE || this == RELATIONAL_DELETE_DATA_NODE - || this == OBJECT_FILE_NODE; + || this == OBJECT_FILE_NODE + || this == LOAD_TSFILE_CONSENSUS_NODE; } public static WALEntryType valueOf(byte code) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java index 1be41760e8e30..381c870213494 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java @@ -29,13 +29,10 @@ import org.apache.iotdb.commons.disk.strategy.DirectoryStrategyType; import org.apache.iotdb.commons.exception.DiskSpaceInsufficientException; import org.apache.iotdb.commons.file.SystemFileFactory; -import org.apache.iotdb.commons.schema.table.TsFileTableSchemaUtil; -import org.apache.iotdb.commons.schema.table.TsTable; import org.apache.iotdb.commons.service.metric.MetricService; import org.apache.iotdb.commons.service.metric.enums.Metric; import org.apache.iotdb.commons.service.metric.enums.Tag; import org.apache.iotdb.commons.utils.FileUtils; -import org.apache.iotdb.commons.utils.PathUtils; import org.apache.iotdb.commons.utils.RetryUtils; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; @@ -45,8 +42,6 @@ import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; import org.apache.iotdb.db.queryengine.plan.scheduler.load.LoadTsFileScheduler.LoadCommand; -import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache; -import org.apache.iotdb.db.service.metrics.DataNodeExceptionMetrics; import org.apache.iotdb.db.storageengine.dataregion.DataRegion; import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; @@ -64,14 +59,14 @@ import org.apache.tsfile.common.constant.TsFileConstant; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.PageException; -import org.apache.tsfile.file.metadata.ChunkGroupMetadata; -import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.read.TimeValuePair; +import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.utils.Pair; import org.apache.tsfile.utils.RamUsageEstimator; import org.apache.tsfile.utils.TsPrimitiveType; -import org.apache.tsfile.write.writer.TsFileIOWriter; +import org.apache.tsfile.write.TsFilePrecalculatedChunkWriter; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -87,7 +82,6 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.PriorityBlockingQueue; import java.util.concurrent.atomic.AtomicLong; @@ -122,13 +116,15 @@ public class LoadTsFileManager { .build(); private final Map uuid2WriterManager = new ConcurrentHashMap<>(); + private final DataRegion dataRegion; private final Map uuid2CleanupTask = new ConcurrentHashMap<>(); private final PriorityBlockingQueue cleanupTaskQueue = new PriorityBlockingQueue<>(); private final ActiveLoadAgent activeLoadAgent = new ActiveLoadAgent(); - public LoadTsFileManager() { + public LoadTsFileManager(final DataRegion dataRegion) { + this.dataRegion = Objects.requireNonNull(dataRegion); registerCleanupTaskExecutor(); recover(); } @@ -226,7 +222,11 @@ private void recover() { } } - final File[] baseDirs = Arrays.stream(LOAD_BASE_DIRS.get()).map(File::new).toArray(File[]::new); + final File[] baseDirs = + Arrays.stream(LOAD_BASE_DIRS.get()) + .map(File::new) + .map(this::getDataRegionLoadDir) + .toArray(File[]::new); final File[] files = Arrays.stream(baseDirs) .filter(File::exists) @@ -251,7 +251,12 @@ private void recover() { }); } - public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNode, String uuid) + public void writeToDataRegion(LoadTsFilePieceNode pieceNode, String uuid) + throws IOException, PageException { + writePiece(uuid, pieceNode.getAllTsFileData()); + } + + public void writePiece(final String uuid, final List tsFileDataList) throws IOException, PageException { createCleanupTaskIfAbsent(uuid); @@ -265,7 +270,10 @@ public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNo o -> { try { return getFolderManager() - .getNextWithRetry(folder -> new TsFileWriterManager(new File(folder, uuid))); + .getNextWithRetry( + folder -> + new TsFileWriterManager( + new File(getDataRegionLoadDir(new File(folder)), uuid))); } catch (DiskSpaceInsufficientException e) { exception.set(e); return null; @@ -281,7 +289,7 @@ public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNo exception.get()); } - for (TsFileData tsFileData : pieceNode.getAllTsFileData()) { + for (TsFileData tsFileData : tsFileDataList) { switch (tsFileData.getType()) { case CHUNK: ChunkData chunkData = (ChunkData) tsFileData; @@ -328,6 +336,14 @@ private FolderManager getFolderManager() throws DiskSpaceInsufficientException { return FOLDER_MANAGER.get(); } + private File getDataRegionLoadDir(final File baseDir) { + return new File( + baseDir, + dataRegion.getDatabaseName() + + IoTDBConstant.FILE_NAME_SEPARATOR + + dataRegion.getDataRegionIdString()); + } + public boolean loadAll( String uuid, boolean isGeneratedByPipe, @@ -442,20 +458,16 @@ public static void cleanTsFile(final File tsFile) { private static class TsFileWriterManager { private final File taskDir; - private Map dataPartition2Writer; + private Map dataPartition2Writer; private Map dataPartition2Resource; - private Map dataPartition2LastDevice; private Map dataPartition2ModificationFile; - private Map> device2Partition; private boolean isClosed; private TsFileWriterManager(File taskDir) { this.taskDir = taskDir; this.dataPartition2Writer = new HashMap<>(); this.dataPartition2Resource = new HashMap<>(); - this.dataPartition2LastDevice = new HashMap<>(); this.dataPartition2ModificationFile = new HashMap<>(); - device2Partition = new HashMap<>(); this.isClosed = false; clearDir(taskDir); @@ -490,83 +502,32 @@ private void write(DataPartitionInfo partitionInfo, ChunkData chunkData) return; } - final long chunkMetadataMaxSizeForEachWriter = - CONFIG.getLoadChunkMetadataMemorySizeInBytes() / (dataPartition2Writer.size() + 1); - final TsFileIOWriter writer = - new TsFileIOWriter(newTsFile, chunkMetadataMaxSizeForEachWriter); - final TsFileResource resource = new TsFileResource(writer.getFile()); - writer.addFlushListener( - // Update time index by chunk groups going to be flushed to temp file - sortedChunkMetadataList -> - sortedChunkMetadataList.forEach( - pair -> { - final IDeviceID deviceId = pair.left.left; - pair.getRight() - .forEach( - chunkMetadata -> { - resource.updateStartTime(deviceId, chunkMetadata.getStartTime()); - resource.updateEndTime(deviceId, chunkMetadata.getEndTime()); - }); - })); - - // When a new writer is added, we need to reduce the metadata size limit of all existing - // writers for memory control - for (final TsFileIOWriter existingWriter : dataPartition2Writer.values()) { - existingWriter.setMaxMetadataSize(chunkMetadataMaxSizeForEachWriter); - } + final TsFilePrecalculatedChunkWriter writer = new TsFilePrecalculatedChunkWriter(newTsFile); + final TsFileResource resource = new TsFileResource(newTsFile); dataPartition2Writer.put(partitionInfo, writer); dataPartition2Resource.put(partitionInfo, resource); } - TsFileIOWriter writer = dataPartition2Writer.get(partitionInfo); - - // Table model needs to register TableSchema - final String tableName = - chunkData.getDevice() != null ? chunkData.getDevice().getTableName() : null; - if (tableName != null - && PathUtils.isTableModelDatabase(partitionInfo.getDataRegion().getDatabaseName())) { - // If the table does not exist, it means that the table is all deleted by mods - final TsTable table = - DataNodeTableCache.getInstance() - .getTable(partitionInfo.getDataRegion().getDatabaseName(), tableName, false); - if (Objects.nonNull(table)) { - writer - .getSchema() - .getTableSchemaMap() - .computeIfAbsent( - tableName, t -> TsFileTableSchemaUtil.toTsFileTableSchemaNoAttribute(table)); - } + final TsFilePrecalculatedChunkWriter writer = dataPartition2Writer.get(partitionInfo); + final ChunkData.ChunkLayout layout = chunkData.getChunkLayout(); + if (layout == null) { + throw new IOException("Chunk layout is missing"); } - - IDeviceID device = chunkData.getDevice(); - IDeviceID lastDevice = dataPartition2LastDevice.get(partitionInfo); - - if (!Objects.equals(device, lastDevice)) { - if (lastDevice != null && device2Partition.containsKey(lastDevice)) { - Set partitions = device2Partition.get(lastDevice); - for (DataPartitionInfo partition : partitions) { - TsFileIOWriter w = dataPartition2Writer.get(partition); - if (dataPartition2LastDevice.containsKey(partition) && w != null) { - w.endChunkGroup(); - w.checkMetadataSizeAndMayFlush(); - } - } - device2Partition.remove(lastDevice); - } - if (writer.isWritingChunkGroup()) { - LOGGER.warn( - StorageEngineMessages - .STORAGE_LOG_WRITER_FOR_PARTITION_IS_ALREADY_WRITING_CHUNK_GROUP_FOR_903B1D66, - writer.getFile().getAbsolutePath(), - partitionInfo, - device, - lastDevice); - } - writer.startChunkGroup(device); - dataPartition2LastDevice.put(partitionInfo, device); - device2Partition.computeIfAbsent(device, k -> new HashSet<>()).add(partitionInfo); + long chunkOffset = layout.offset(); + final List chunks = chunkData.getChunks(); + for (int i = 0; i < chunks.size(); i++) { + final Chunk chunk = chunks.get(i); + writer.writeChunk( + chunkData.getDevice(), + chunkData.isAligned(), + layout.chunkGroupHeaderOffset(), + layout.firstChunkOfGroup() && i == 0, + chunk, + chunkOffset); + chunkOffset += chunk.getHeader().getSerializedSize() + (long) chunk.getData().remaining(); + } + if (chunkOffset != layout.offset() + layout.length()) { + throw new IOException("Chunk layout length does not match encoded chunks"); } - - chunkData.writeToFileWriter(writer); } private void writeDeletion(DataRegion dataRegion, DeletionData deletionData) @@ -574,10 +535,11 @@ private void writeDeletion(DataRegion dataRegion, DeletionData deletionData) if (isClosed) { throw new IOException(String.format(MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED, taskDir)); } - for (Map.Entry entry : dataPartition2Writer.entrySet()) { + for (Map.Entry entry : + dataPartition2Writer.entrySet()) { final DataPartitionInfo partitionInfo = entry.getKey(); if (partitionInfo.getDataRegion().equals(dataRegion)) { - final TsFileIOWriter writer = entry.getValue(); + final TsFilePrecalculatedChunkWriter writer = entry.getValue(); if (!dataPartition2ModificationFile.containsKey(partitionInfo)) { File newModificationFile = ModificationFile.getExclusiveMods(writer.getFile()); if (!newModificationFile.createNewFile()) { @@ -592,7 +554,7 @@ private void writeDeletion(DataRegion dataRegion, DeletionData deletionData) partitionInfo, new ModificationFile(newModificationFile, false)); } ModificationFile modificationFile = dataPartition2ModificationFile.get(partitionInfo); - writer.flush(); + writer.getOutput().flush(); deletionData.writeToModificationFile(modificationFile); } } @@ -609,13 +571,10 @@ private void loadAll( dataPartition2ModificationFile.entrySet()) { entry.getValue().close(); } - for (final Map.Entry entry : + for (final Map.Entry entry : dataPartition2Writer.entrySet()) { - final TsFileIOWriter writer = entry.getValue(); - if (writer.isWritingChunkGroup()) { - writer.endChunkGroup(); - } - writer.endFile(); + final TsFilePrecalculatedChunkWriter writer = entry.getValue(); + writer.close(); final DataRegion dataRegion = entry.getKey().getDataRegion(); final TsFileResource tsFileResource = dataPartition2Resource.get(entry.getKey()); @@ -625,12 +584,7 @@ private void loadAll( tsFileResource, timePartitionProgressIndexMap.getOrDefault( entry.getKey().getTimePartitionSlot(), MinimumProgressIndex.INSTANCE)); - dataRegion.loadNewTsFile( - tsFileResource, - true, - isGeneratedByPipe, - false, - Optional.ofNullable(writer.getTableSizeMap())); + dataRegion.loadNewTsFile(tsFileResource, true, isGeneratedByPipe, false, Optional.empty()); // Metrics dataRegion @@ -643,7 +597,9 @@ private void loadAll( } private void endTsFileResource( - TsFileIOWriter writer, TsFileResource tsFileResource, ProgressIndex progressIndex) + TsFilePrecalculatedChunkWriter writer, + TsFileResource tsFileResource, + ProgressIndex progressIndex) throws IOException { // Update time index by chunk groups still in memory Map> deviceLastValues = null; @@ -652,9 +608,10 @@ private void endTsFileResource( } AtomicLong lastValuesMemCost = new AtomicLong(0); - for (final ChunkGroupMetadata chunkGroupMetadata : writer.getChunkGroupMetadataList()) { - final IDeviceID device = chunkGroupMetadata.getDevice(); - for (final ChunkMetadata chunkMetadata : chunkGroupMetadata.getChunkMetadataList()) { + for (final Map.Entry> entry : + writer.getChunkMetadataListMap().entrySet()) { + final IDeviceID device = entry.getKey(); + for (final IChunkMetadata chunkMetadata : entry.getValue()) { tsFileResource.updateStartTime(device, chunkMetadata.getStartTime()); tsFileResource.updateEndTime(device, chunkMetadata.getEndTime()); if (deviceLastValues != null) { @@ -727,9 +684,9 @@ private void endTsFileResource( tsFileResource.serialize(); } - private long getTsFileWritePointCount(TsFileIOWriter writer) { - return writer.getChunkGroupMetadataList().stream() - .flatMap(chunkGroupMetadata -> chunkGroupMetadata.getChunkMetadataList().stream()) + private long getTsFileWritePointCount(TsFilePrecalculatedChunkWriter writer) { + return writer.getChunkMetadataListMap().values().stream() + .flatMap(List::stream) .mapToLong(chunkMetadata -> chunkMetadata.getStatistics().getCount()) .sum(); } @@ -739,12 +696,11 @@ private void close() { return; } if (dataPartition2Writer != null) { - for (Map.Entry entry : dataPartition2Writer.entrySet()) { + for (Map.Entry entry : + dataPartition2Writer.entrySet()) { try { - final TsFileIOWriter writer = entry.getValue(); - if (writer.canWrite()) { - writer.close(); - } + final TsFilePrecalculatedChunkWriter writer = entry.getValue(); + writer.getOutput().close(); final Path writerPath = writer.getFile().toPath(); if (Files.exists(writerPath)) { RetryUtils.retryOnException( @@ -758,7 +714,6 @@ private void close() { StorageEngineMessages.CLOSE_TSFILE_IO_WRITER_ERROR, entry.getValue().getFile().getPath(), e); - DataNodeExceptionMetrics.getInstance().recordSuspiciousDiskException(e); } } } @@ -779,7 +734,6 @@ private void close() { } catch (IOException e) { LOGGER.warn( StorageEngineMessages.CLOSE_MODIFICATION_FILE_ERROR, entry.getValue().getFile(), e); - DataNodeExceptionMetrics.getInstance().recordSuspiciousDiskException(e); } } } @@ -793,13 +747,10 @@ private void close() { LOGGER.info(StorageEngineMessages.TASK_DIR_NOT_EMPTY_SKIP_DELETE, taskDir.getPath()); } catch (IOException e) { LOGGER.warn(MESSAGE_DELETE_FAIL, taskDir.getPath(), e); - DataNodeExceptionMetrics.getInstance().recordSuspiciousDiskException(e); } dataPartition2Writer = null; dataPartition2Resource = null; - dataPartition2LastDevice = null; dataPartition2ModificationFile = null; - device2Partition = null; isClosed = true; } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java index 5ca04b4d7d3af..a49dbc3a00074 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java @@ -22,8 +22,9 @@ import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.utils.TimePartitionUtils; import org.apache.iotdb.db.i18n.StorageEngineMessages; -import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.PageException; import org.apache.tsfile.file.header.ChunkHeader; @@ -35,11 +36,11 @@ import org.apache.tsfile.file.metadata.statistics.Statistics; import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.utils.Binary; -import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.utils.TsPrimitiveType; import org.apache.tsfile.write.UnSupportedDataTypeException; import org.apache.tsfile.write.chunk.AlignedChunkWriterImpl; +import org.apache.tsfile.write.chunk.ValueChunkWriter; import org.apache.tsfile.write.schema.IMeasurementSchema; import org.apache.tsfile.write.schema.MeasurementSchema; import org.apache.tsfile.write.writer.TsFileIOWriter; @@ -50,7 +51,6 @@ import java.io.DataOutputStream; import java.io.IOException; import java.io.InputStream; -import java.io.Serializable; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.LinkedList; @@ -58,6 +58,8 @@ import java.util.Queue; import static org.apache.iotdb.db.storageengine.load.LoadTsFileManager.MEASUREMENT_ID_CACHE; +import static org.apache.tsfile.common.constant.TsFileConstant.TIME_COLUMN_MASK; +import static org.apache.tsfile.common.constant.TsFileConstant.VALUE_COLUMN_MASK; public class AlignedChunkData implements ChunkData { protected static final int DEFAULT_INT32 = 0; @@ -67,32 +69,99 @@ public class AlignedChunkData implements ChunkData { protected static final boolean DEFAULT_BOOLEAN = false; protected static final Binary DEFAULT_BINARY = null; + // ---------------- Permanent state: describes this ChunkData and its final result + // ---------------- protected final TTimePartitionSlot timePartitionSlot; protected final IDeviceID device; + protected final EncryptParameter encryptParameter; protected List chunkHeaderList; - protected PublicBAOS byteStream; - protected DataOutputStream stream; - protected List timeBatch; + /** Final storage for complete Chunks; written directly when decoding is not required. */ + protected final List entireChunks = new ArrayList<>(); + protected long dataSize; protected boolean needDecodeChunk; + + // ---------------- Temporary state: used only while the splitter collects Pages ---------------- + /** Pages that are not encoded yet; built and cleared together in {@code endChunk()}. */ + protected final List pageBuffers = new ArrayList<>(); + + /** Index of the Chunk currently being written: 0 = Time Chunk, 1..n = Value Chunks. */ + protected int currentChunkIndex = -1; + + protected ChunkLayout chunkLayout; + + /** Temporary time batches used for aligned row-level splitting. */ + protected List timeBatch; + protected List pageNumbers; protected Queue satisfiedLengthQueue; - protected ByteBuffer chunkData; + protected static class PageBuffer { + final int chunkIndex; + final boolean needDecode; + final PageHeader pageHeader; + final ByteBuffer pageData; + final long[] timeBatch; + final TsPrimitiveType[] valueBatch; + final int satisfiedLength; + + PageBuffer(int chunkIndex, PageHeader pageHeader, ByteBuffer pageData) { + this.chunkIndex = chunkIndex; + this.needDecode = false; + this.pageHeader = pageHeader; + this.pageData = pageData; + this.timeBatch = null; + this.valueBatch = null; + this.satisfiedLength = 0; + } + + PageBuffer(int chunkIndex, long[] timeBatch, int satisfiedLength) { + this.chunkIndex = chunkIndex; + this.needDecode = true; + this.pageHeader = null; + this.pageData = null; + this.timeBatch = timeBatch; + this.valueBatch = null; + this.satisfiedLength = satisfiedLength; + } + + PageBuffer( + int chunkIndex, long[] timeBatch, TsPrimitiveType[] valueBatch, int satisfiedLength) { + this.chunkIndex = chunkIndex; + this.needDecode = true; + this.pageHeader = null; + this.pageData = null; + this.timeBatch = timeBatch; + this.valueBatch = valueBatch; + this.satisfiedLength = satisfiedLength; + } + } public AlignedChunkData( @Nonnull final IDeviceID device, final ChunkHeader chunkHeader, final TTimePartitionSlot timePartitionSlot) { - this(device, timePartitionSlot); + this(device, chunkHeader, timePartitionSlot, EncryptUtils.getEncryptParameter()); + } + + public AlignedChunkData( + @Nonnull final IDeviceID device, + final ChunkHeader chunkHeader, + final TTimePartitionSlot timePartitionSlot, + final EncryptParameter encryptParameter) { + this(device, timePartitionSlot, encryptParameter); chunkHeaderList.add(chunkHeader); pageNumbers.add(0); + currentChunkIndex = 0; addAttrDataSize(); } protected AlignedChunkData(final AlignedChunkData alignedChunkData) { - this(alignedChunkData.device, alignedChunkData.timePartitionSlot); + this( + alignedChunkData.device, + alignedChunkData.timePartitionSlot, + alignedChunkData.encryptParameter); this.satisfiedLengthQueue = new LinkedList<>(alignedChunkData.satisfiedLengthQueue); this.needDecodeChunk = alignedChunkData.needDecodeChunk; addAttrDataSize(); @@ -100,24 +169,30 @@ protected AlignedChunkData(final AlignedChunkData alignedChunkData) { protected AlignedChunkData( @Nonnull final IDeviceID device, final TTimePartitionSlot timePartitionSlot) { + this(device, timePartitionSlot, EncryptUtils.getEncryptParameter()); + } + + protected AlignedChunkData( + @Nonnull final IDeviceID device, + final TTimePartitionSlot timePartitionSlot, + final EncryptParameter encryptParameter) { this.dataSize = 0; this.device = device; + this.encryptParameter = encryptParameter; this.chunkHeaderList = new ArrayList<>(); this.timePartitionSlot = timePartitionSlot; this.needDecodeChunk = true; this.pageNumbers = new ArrayList<>(); this.satisfiedLengthQueue = new LinkedList<>(); - this.byteStream = new PublicBAOS(); - this.stream = new DataOutputStream(byteStream); } - private void addAttrDataSize() { // Should be init before serialize, corresponding serializeAttr - dataSize += 2 * Byte.BYTES; // isModification and isAligned - dataSize += Long.BYTES; // timePartitionSlot - dataSize += device.serializedSize(); // device - dataSize += Integer.BYTES; // chunkHeaderListSize + private void addAttrDataSize() { + dataSize += 2 * Byte.BYTES; + dataSize += Long.BYTES; + dataSize += device.serializedSize(); + dataSize += Integer.BYTES; if (!chunkHeaderList.isEmpty()) { - dataSize += chunkHeaderList.get(0).getSerializedSize(); // timeChunkHeader + dataSize += chunkHeaderList.get(0).getSerializedSize(); } } @@ -151,66 +226,56 @@ public boolean isEntireChunk() { return !needDecodeChunk; } - @Override - public void writeToFileWriter(final TsFileIOWriter writer) throws IOException, PageException { - writeTsFileData(writer); - } - public void addValueChunk(final ChunkHeader chunkHeader) { this.chunkHeaderList.add(chunkHeader); this.pageNumbers.add(0); + currentChunkIndex++; dataSize += chunkHeader.getSerializedSize(); if (needDecodeChunk) { - dataSize += Integer.BYTES; // pageNumber + dataSize += Integer.BYTES; } } @Override - public void serialize(final DataOutputStream stream) throws IOException { - ReadWriteIOUtils.write(getType().ordinal(), stream); - ReadWriteIOUtils.write(isAligned(), stream); - serializeAttr(stream); - ReadWriteIOUtils.write(byteStream.size(), stream); - byteStream.writeTo(stream); - close(); + public List getChunks() { + if (!pageBuffers.isEmpty()) { + endChunk(); + } + return new ArrayList<>(entireChunks); } - private void serializeAttr(final DataOutputStream stream) throws IOException { - ReadWriteIOUtils.write(timePartitionSlot.getStartTime(), stream); - - ReadWriteIOUtils.write(device instanceof StringArrayDeviceID, stream); - device.serialize(stream); + @Override + public ChunkLayout getChunkLayout() { + return chunkLayout; + } - ReadWriteIOUtils.write(dataSize, stream); - ReadWriteIOUtils.write(needDecodeChunk, stream); - ReadWriteIOUtils.write(chunkHeaderList.size(), stream); - for (final ChunkHeader chunkHeader : chunkHeaderList) { - chunkHeader.serializeTo(stream); // Chunk header already serialize chunk type - } - if (needDecodeChunk) { - for (final Integer pageNumber : pageNumbers) { - ReadWriteIOUtils.write(pageNumber, stream); - } - } + @Override + public void setChunkLayout(final ChunkLayout chunkLayout) { + this.chunkLayout = chunkLayout; } + // ----------------------- Data Collection (In-Memory) ----------------------- + @Override public void writeEntireChunk(final ByteBuffer chunkData, final IChunkMetadata chunkMetadata) throws IOException { - dataSize += ReadWriteIOUtils.write(chunkData, stream); - dataSize += chunkMetadata.getStatistics().serialize(stream); + if (chunkHeaderList.isEmpty()) { + throw new IOException("Chunk header list is empty."); + } + final ChunkHeader currentChunkHeader = chunkHeaderList.get(chunkHeaderList.size() - 1); + + // Store complete chunk data directly into memory without decoding + entireChunks.add(new Chunk(currentChunkHeader, chunkData, null, chunkMetadata.getStatistics())); + needDecodeChunk = false; + dataSize += chunkData.remaining() + chunkMetadata.getStatistics().getSerializedSize(); } @Override public void writeEntirePage(final PageHeader pageHeader, final ByteBuffer pageData) throws IOException { pageNumbers.set(pageNumbers.size() - 1, pageNumbers.get(pageNumbers.size() - 1) + 1); - // serialize needDecode==false - dataSize += ReadWriteIOUtils.write(false, stream); - // serialize pageHeader - dataSize += pageHeader.serializeTo(stream); - // serialize pageData - dataSize += ReadWriteIOUtils.write(pageData, stream); + // Wrap the un-decoded page and store it into the buffer + pageBuffers.add(new PageBuffer(currentChunkIndex, pageHeader, pageData)); } @Override @@ -218,19 +283,15 @@ public void writeDecodePage(final long[] times, final Object[] values, final int throws IOException { pageNumbers.set(pageNumbers.size() - 1, pageNumbers.get(pageNumbers.size() - 1) + 1); satisfiedLengthQueue.offer(satisfiedLength); - final long startTime = timePartitionSlot.getStartTime(); - final long endTime = TimePartitionUtils.getTimePartitionEndTime(startTime); - // serialize needDecode==true - dataSize += ReadWriteIOUtils.write(true, stream); - dataSize += ReadWriteIOUtils.write(satisfiedLength, stream); - - for (final long time : times) { - if (time > endTime) { - break; - } - if (time >= startTime) { - dataSize += ReadWriteIOUtils.write(time, stream); + + if (currentChunkIndex == 0 || values == null) { + pageBuffers.add(new PageBuffer(currentChunkIndex, times, satisfiedLength)); + } else { + final TsPrimitiveType[] alignedValues = new TsPrimitiveType[values.length]; + for (int i = 0; i < values.length; i++) { + alignedValues[i] = (TsPrimitiveType) values[i]; } + pageBuffers.add(new PageBuffer(currentChunkIndex, times, alignedValues, satisfiedLength)); } } @@ -238,224 +299,222 @@ public void writeDecodeValuePage( final long[] times, final TsPrimitiveType[] values, final TSDataType dataType) throws IOException { pageNumbers.set(pageNumbers.size() - 1, pageNumbers.get(pageNumbers.size() - 1) + 1); - final long startTime = timePartitionSlot.getStartTime(); - final long endTime = TimePartitionUtils.getTimePartitionEndTime(startTime); final int satisfiedLength = satisfiedLengthQueue.poll(); - // serialize needDecode==true - dataSize += ReadWriteIOUtils.write(true, stream); - dataSize += ReadWriteIOUtils.write(satisfiedLength, stream); satisfiedLengthQueue.offer(satisfiedLength); - - for (int i = 0; i < times.length; i++) { - if (times[i] > endTime) { - break; - } - if (times[i] >= startTime) { - if (values.length == 0 || values[i] == null) { - dataSize += ReadWriteIOUtils.write(true, stream); - } else { - dataSize += ReadWriteIOUtils.write(false, stream); - switch (dataType) { - case INT32: - case DATE: - dataSize += ReadWriteIOUtils.write(values[i].getInt(), stream); - break; - case INT64: - case TIMESTAMP: - dataSize += ReadWriteIOUtils.write(values[i].getLong(), stream); - break; - case FLOAT: - dataSize += ReadWriteIOUtils.write(values[i].getFloat(), stream); - break; - case DOUBLE: - dataSize += ReadWriteIOUtils.write(values[i].getDouble(), stream); - break; - case BOOLEAN: - dataSize += ReadWriteIOUtils.write(values[i].getBoolean(), stream); - break; - case TEXT: - case BLOB: - case STRING: - dataSize += ReadWriteIOUtils.write(values[i].getBinary(), stream); - break; - default: - throw new UnSupportedDataTypeException( - String.format( - StorageEngineMessages.STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, - dataType)); - } - } - } - } + pageBuffers.add(new PageBuffer(currentChunkIndex, times, values, satisfiedLength)); } - protected void writeTsFileData(TsFileIOWriter writer) throws IOException, PageException { - ensureDataReadyForWriting(); - final InputStream stream = new LoadTsFilePieceNode.ByteBufferInputStream(chunkData); - if (needDecodeChunk) { - writeChunkToWriter(stream, writer); - } else { - writeEntireChunkToWriter(stream, writer); - } - } + // ----------------------- Unified Encoder Trigger ----------------------- - private void ensureDataReadyForWriting() throws IOException { - if (chunkData != null) { - chunkData.rewind(); + @Override + public void endChunk() { + if (pageBuffers.isEmpty()) { return; } - chunkData = ByteBuffer.wrap(byteStream.getBuf(), 0, byteStream.size()); + try { + // Force encoding of buffered pages into standard Chunk objects + encodeAndBuildChunks(); + + // Clear the temporary buffers + pageBuffers.clear(); + // Mark as fully processed standard chunks + needDecodeChunk = false; + } catch (final IOException | PageException e) { + throw new IllegalStateException("Failed to encode chunk using AlignedChunkWriterImpl", e); + } } - protected void deserializeTsFileDataByte(final InputStream stream) throws IOException { - final int size = ReadWriteIOUtils.readInt(stream); - if (stream instanceof LoadTsFilePieceNode.ByteBufferInputStream) { - this.chunkData = ((LoadTsFilePieceNode.ByteBufferInputStream) stream).read(size); + protected void encodeAndBuildChunks() throws IOException, PageException { + final AlignedChunkWriterImpl writer = createAlignedChunkWriter(); + + writePageBuffersToWriter(writer); + + // endChunk() seals one physical Chunk. Building all columns here would emit the time Chunk + // again whenever a later value Chunk is split and encoded. + final ChunkHeader sourceHeader = chunkHeaderList.get(currentChunkIndex); + final ByteBuffer encodedData; + final Statistics statistics; + final int pageCount; + final int mask; + if (currentChunkIndex == 0) { + encodedData = writer.getTimeChunkWriter().getByteBuffer(); + statistics = writer.getTimeChunkWriter().getStatistics(); + pageCount = writer.getTimeChunkWriter().getNumOfPages(); + mask = TIME_COLUMN_MASK; } else { - byte[] data = new byte[size]; - new DataInputStream(stream).readFully(data); - this.chunkData = ByteBuffer.wrap(data); + final ValueChunkWriter valueWriter = + writer.getValueChunkWriterList().get(currentChunkIndex - 1); + encodedData = valueWriter.getByteBuffer(); + statistics = valueWriter.getStatistics(); + pageCount = valueWriter.getNumOfPages(); + mask = VALUE_COLUMN_MASK; } + + final ChunkHeader encodedHeader = + createEncodedChunkHeader(sourceHeader, encodedData.remaining(), pageCount, mask); + entireChunks.add(new Chunk(encodedHeader, encodedData, null, statistics)); + dataSize += encodedData.remaining() + statistics.getSerializedSize(); } - private void writeEntireChunkToWriter(final InputStream stream, final TsFileIOWriter writer) - throws IOException { - for (final ChunkHeader chunkHeader : chunkHeaderList) { - final ByteBuffer chunkData = - ByteBuffer.wrap(ReadWriteIOUtils.readBytesWithSelfDescriptionLength(stream)); - final Statistics statistics = - Statistics.deserialize(stream, chunkHeader.getDataType()); - writer.writeChunk(new Chunk(chunkHeader, chunkData, null, statistics)); - } + protected static ChunkHeader createEncodedChunkHeader( + final ChunkHeader sourceHeader, + final int encodedDataSize, + final int pageCount, + final int mask) { + return new ChunkHeader( + sourceHeader.getMeasurementID(), + encodedDataSize, + sourceHeader.getDataType(), + sourceHeader.getCompressionType(), + sourceHeader.getEncodingType(), + pageCount, + mask); } - protected void writeChunkToWriter(final InputStream stream, final TsFileIOWriter writer) - throws IOException, PageException { - final List measurementSchemaList = new ArrayList<>(); + private AlignedChunkWriterImpl createAlignedChunkWriter() { IMeasurementSchema timeSchema = null; - for (final ChunkHeader chunkHeader : chunkHeaderList) { - if (TSDataType.VECTOR.equals(chunkHeader.getDataType())) { - timeSchema = - new MeasurementSchema( - chunkHeader.getMeasurementID(), - chunkHeader.getDataType(), - chunkHeader.getEncodingType(), - chunkHeader.getCompressionType()); - continue; - } - measurementSchemaList.add( + final List valueSchemaList = new ArrayList<>(); + for (final ChunkHeader header : chunkHeaderList) { + final MeasurementSchema schema = new MeasurementSchema( - chunkHeader.getMeasurementID(), - chunkHeader.getDataType(), - chunkHeader.getEncodingType(), - chunkHeader.getCompressionType())); - } - AlignedChunkWriterImpl chunkWriter = - new AlignedChunkWriterImpl(timeSchema, measurementSchemaList); - timeBatch = new ArrayList<>(); - final int chunkHeaderSize = chunkHeaderList.size(); - for (int i = 0; i < chunkHeaderSize; i++) { - buildChunk(chunkWriter, stream, chunkHeaderList.get(i), pageNumbers.get(i), i - 1, i == 0); + header.getMeasurementID(), + header.getDataType(), + header.getEncodingType(), + header.getCompressionType()); + if (header.getDataType() == TSDataType.VECTOR) { + timeSchema = schema; + } else { + valueSchemaList.add(schema); + } } - timeBatch = null; - chunkWriter.writeToFileWriter(writer); + return new AlignedChunkWriterImpl(timeSchema, valueSchemaList, encryptParameter); } - @SuppressWarnings({"squid:S6541", "squid:S3776"}) - private void buildChunk( - final AlignedChunkWriterImpl chunkWriter, - final InputStream stream, - final ChunkHeader chunkHeader, - final int pageNumber, - final int valueChunkIndex, - final boolean isTimeChunk) + private void writePageBuffersToWriter(final AlignedChunkWriterImpl writer) throws IOException, PageException { - boolean needDecode; - int decodePageIndex = 0; - for (int j = 0; j < pageNumber; j++) { - needDecode = ReadWriteIOUtils.readBool(stream); - if (needDecode) { - final int length = ReadWriteIOUtils.readInt(stream); - long[] timePageBatch = new long[length]; - if (!isTimeChunk) { - timePageBatch = timeBatch.get(decodePageIndex); - } - for (int i = 0; i < length; i++) { - if (isTimeChunk) { - final long time = ReadWriteIOUtils.readLong(stream); - timePageBatch[i] = time; - chunkWriter.writeTime(time); - } else { - final boolean isNull = ReadWriteIOUtils.readBool(stream); - switch (chunkHeader.getDataType()) { - case INT32: - case DATE: - final int int32Value = isNull ? DEFAULT_INT32 : ReadWriteIOUtils.readInt(stream); - chunkWriter.write(timePageBatch[i], int32Value, isNull, valueChunkIndex); - break; - case INT64: - case TIMESTAMP: - final long int64Value = isNull ? DEFAULT_INT64 : ReadWriteIOUtils.readLong(stream); - chunkWriter.write(timePageBatch[i], int64Value, isNull, valueChunkIndex); - break; - case FLOAT: - final float floatValue = - isNull ? DEFAULT_FLOAT : ReadWriteIOUtils.readFloat(stream); - chunkWriter.write(timePageBatch[i], floatValue, isNull, valueChunkIndex); - break; - case DOUBLE: - final double doubleValue = - isNull ? DEFAULT_DOUBLE : ReadWriteIOUtils.readDouble(stream); - chunkWriter.write(timePageBatch[i], doubleValue, isNull, valueChunkIndex); - break; - case BOOLEAN: - final boolean boolValue = - isNull ? DEFAULT_BOOLEAN : ReadWriteIOUtils.readBool(stream); - chunkWriter.write(timePageBatch[i], boolValue, isNull, valueChunkIndex); - break; - case TEXT: - case BLOB: - case STRING: - final Binary binaryValue = - isNull ? DEFAULT_BINARY : ReadWriteIOUtils.readBinary(stream); - chunkWriter.write(timePageBatch[i], binaryValue, isNull, valueChunkIndex); - break; - default: - throw new UnSupportedDataTypeException( - String.format( - StorageEngineMessages - .STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, - chunkHeader.getDataType())); - } - } - } - if (isTimeChunk) { - timeBatch.add(timePageBatch); - } - decodePageIndex += 1; - - if (isTimeChunk) { - chunkWriter.sealCurrentTimePage(); + for (final PageBuffer page : pageBuffers) { + if (!page.needDecode) { + if (page.chunkIndex == 0) { + writer.writePageHeaderAndDataIntoTimeBuff(page.pageData, page.pageHeader); } else { - chunkWriter.sealCurrentValuePage(valueChunkIndex); + writer.writePageHeaderAndDataIntoValueBuff( + page.pageData, page.pageHeader, page.chunkIndex - 1); } + continue; + } + if (page.valueBatch == null) { + for (int i = 0; i < page.satisfiedLength; i++) { + writer.writeTime(page.timeBatch[i]); + } + writer.sealCurrentTimePage(); } else { - final PageHeader pageHeader = - PageHeader.deserializeFrom(stream, chunkHeader.getDataType(), true); - if (isTimeChunk) { - chunkWriter.writePageHeaderAndDataIntoTimeBuff( - ByteBuffer.wrap(ReadWriteIOUtils.readBytesWithSelfDescriptionLength(stream)), - pageHeader); - } else { - chunkWriter.writePageHeaderAndDataIntoValueBuff( - ByteBuffer.wrap(ReadWriteIOUtils.readBytesWithSelfDescriptionLength(stream)), - pageHeader, - valueChunkIndex); + final TSDataType dataType = chunkHeaderList.get(page.chunkIndex).getDataType(); + for (int i = 0; i < page.satisfiedLength; i++) { + writeValueToAlignedChunkWriter( + writer, page.timeBatch[i], page.valueBatch[i], dataType, page.chunkIndex - 1); } + writer.sealCurrentValuePage(page.chunkIndex - 1); } } } + private void writeValueToAlignedChunkWriter( + final AlignedChunkWriterImpl writer, + final long time, + final TsPrimitiveType value, + final TSDataType dataType, + final int valueChunkIndex) { + final boolean isNull = value == null; + switch (dataType) { + case INT32: + case DATE: + writer.write(time, isNull ? DEFAULT_INT32 : value.getInt(), isNull, valueChunkIndex); + break; + case INT64: + case TIMESTAMP: + writer.write(time, isNull ? DEFAULT_INT64 : value.getLong(), isNull, valueChunkIndex); + break; + case FLOAT: + writer.write(time, isNull ? DEFAULT_FLOAT : value.getFloat(), isNull, valueChunkIndex); + break; + case DOUBLE: + writer.write(time, isNull ? DEFAULT_DOUBLE : value.getDouble(), isNull, valueChunkIndex); + break; + case BOOLEAN: + writer.write(time, isNull ? DEFAULT_BOOLEAN : value.getBoolean(), isNull, valueChunkIndex); + break; + case TEXT: + case BLOB: + case STRING: + writer.write(time, isNull ? DEFAULT_BINARY : value.getBinary(), isNull, valueChunkIndex); + break; + default: + throw new UnSupportedDataTypeException( + String.format( + StorageEngineMessages.STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, + dataType)); + } + } + + // ----------------------- Serialization / Deserialization ----------------------- + + @Override + public void serialize(final DataOutputStream stream) throws IOException { + ReadWriteIOUtils.write(getType().ordinal(), stream); + + // Always false now since we ensure all outputs are encoded Chunk objects + ReadWriteIOUtils.write(isAligned(), stream); + serializeAttr(stream); + writeChunkLayout(stream); + + // Fallback trigger: ensure un-processed buffers are encoded before serialization + if (!pageBuffers.isEmpty()) { + endChunk(); + } + + // Write all standard chunks out natively + ReadWriteIOUtils.write(entireChunks.size(), stream); + for (Chunk chunk : entireChunks) { + chunk.getHeader().serializeTo(stream); + chunk.getChunkStatistic().serialize(stream); + + ByteBuffer data = chunk.getData(); + ReadWriteIOUtils.write(data.remaining(), stream); + stream.write(data.array(), data.arrayOffset() + data.position(), data.remaining()); + } + } + + private void serializeAttr(final DataOutputStream stream) throws IOException { + ReadWriteIOUtils.write(timePartitionSlot.getStartTime(), stream); + ReadWriteIOUtils.write(device instanceof StringArrayDeviceID, stream); + device.serialize(stream); + ReadWriteIOUtils.write(dataSize, stream); + + // needDecodeChunk will be false by the time we reach here if endChunk is triggered + ReadWriteIOUtils.write(needDecodeChunk, stream); + ReadWriteIOUtils.write(chunkHeaderList.size(), stream); + for (final ChunkHeader chunkHeader : chunkHeaderList) { + chunkHeader.serializeTo(stream); + } + if (needDecodeChunk) { + for (final Integer pageNumber : pageNumbers) { + ReadWriteIOUtils.write(pageNumber, stream); + } + } + } + + private void writeChunkLayout(final DataOutputStream stream) throws IOException { + ReadWriteIOUtils.write(chunkLayout != null, stream); + if (chunkLayout != null) { + ReadWriteIOUtils.write(chunkLayout.chunkGroupIndex(), stream); + ReadWriteIOUtils.write(chunkLayout.chunkGroupHeaderOffset(), stream); + ReadWriteIOUtils.write(chunkLayout.offset(), stream); + ReadWriteIOUtils.write(chunkLayout.length(), stream); + ReadWriteIOUtils.write(chunkLayout.chunkIndexInGroup(), stream); + ReadWriteIOUtils.write(chunkLayout.firstChunkOfGroup(), stream); + } + } + public static AlignedChunkData deserialize(final InputStream stream) throws IOException, PageException { final TTimePartitionSlot timePartitionSlot = @@ -468,6 +527,7 @@ public static AlignedChunkData deserialize(final InputStream stream) final long dataSize = ReadWriteIOUtils.readLong(stream); final boolean needDecodeChunk = ReadWriteIOUtils.readBool(stream); final int chunkHeaderListSize = ReadWriteIOUtils.readInt(stream); + final List chunkHeaderList = new ArrayList<>(); for (int i = 0; i < chunkHeaderListSize; i++) { final byte chunkType = ReadWriteIOUtils.readByte(stream); @@ -476,6 +536,7 @@ public static AlignedChunkData deserialize(final InputStream stream) chunkHeader.setMeasurementID(MEASUREMENT_ID_CACHE.get(measurementID, m -> m)); chunkHeaderList.add(chunkHeader); } + final List pageNumbers = new ArrayList<>(); if (needDecodeChunk) { for (int i = 0; i < chunkHeaderListSize; i++) { @@ -492,15 +553,54 @@ public static AlignedChunkData deserialize(final InputStream stream) chunkData.needDecodeChunk = needDecodeChunk; chunkData.chunkHeaderList = chunkHeaderList; chunkData.pageNumbers = pageNumbers; - chunkData.deserializeTsFileDataByte(stream); chunkData.dataSize = dataSize; - chunkData.close(); + + if (ReadWriteIOUtils.readBool(stream)) { + chunkData.chunkLayout = + new ChunkLayout( + ReadWriteIOUtils.readLong(stream), + ReadWriteIOUtils.readLong(stream), + ReadWriteIOUtils.readLong(stream), + ReadWriteIOUtils.readLong(stream), + ReadWriteIOUtils.readInt(stream), + ReadWriteIOUtils.readBool(stream)); + } + + chunkData.deserializeTsFileData(stream); + return chunkData; } - private void close() throws IOException { - byteStream.close(); - stream.close(); + protected void deserializeTsFileData(final InputStream stream) throws IOException { + int chunkCount = ReadWriteIOUtils.readInt(stream); + for (int i = 0; i < chunkCount; i++) { + byte chunkType = ReadWriteIOUtils.readByte(stream); + ChunkHeader header = ChunkHeader.deserializeFrom(stream, chunkType); + Statistics statistics = Statistics.deserialize(stream, header.getDataType()); + + int dataBytesSize = ReadWriteIOUtils.readInt(stream); + byte[] dataBytes = new byte[dataBytesSize]; + new DataInputStream(stream).readFully(dataBytes); + + entireChunks.add(new Chunk(header, ByteBuffer.wrap(dataBytes), null, statistics)); + } + // Deserialized chunks are now fully standard + this.needDecodeChunk = false; + } + + // ----------------------- Write to TsFile ----------------------- + + @Override + public void writeToFileWriter(final TsFileIOWriter writer) throws IOException, PageException { + // Ensure all internal pages are encoded to chunks if not already + if (!pageBuffers.isEmpty()) { + endChunk(); + } + + // Everything is now standardized as complete chunks in memory, directly write to disk + for (final Chunk chunk : entireChunks) { + writer.writeChunk(chunk); + } } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java index 998441da3ae9a..57fb7a01a23b6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java @@ -23,25 +23,22 @@ import org.apache.iotdb.commons.utils.TimePartitionUtils; import org.apache.iotdb.db.i18n.StorageEngineMessages; -import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.PageException; import org.apache.tsfile.file.header.ChunkHeader; -import org.apache.tsfile.file.header.PageHeader; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.file.metadata.statistics.Statistics; -import org.apache.tsfile.utils.Binary; -import org.apache.tsfile.utils.ReadWriteIOUtils; +import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.utils.TsPrimitiveType; import org.apache.tsfile.write.UnSupportedDataTypeException; import org.apache.tsfile.write.chunk.ValueChunkWriter; import org.apache.tsfile.write.schema.MeasurementSchema; -import org.apache.tsfile.write.writer.TsFileIOWriter; import java.io.IOException; -import java.io.InputStream; import java.io.Serializable; import java.nio.ByteBuffer; +import static org.apache.tsfile.common.constant.TsFileConstant.VALUE_COLUMN_MASK; + /** * This class is used to be compatible with the new distribution of aligned series in chunk group. * In past versions, a time chunk and all its corresponding value chunks were continuous within the @@ -60,80 +57,28 @@ public BatchedAlignedValueChunkData(IDeviceID device, TTimePartitionSlot timePar super(device, timePartitionSlot); } - @Override - public void writeDecodeValuePage(long[] times, TsPrimitiveType[] values, TSDataType dataType) - throws IOException { - pageNumbers.set(pageNumbers.size() - 1, pageNumbers.get(pageNumbers.size() - 1) + 1); + /** + * Override the encoding process to handle standalone Value Chunks without a Time Chunk. It uses + * ValueChunkWriter directly, writes a dummy time (0), and manually corrects the start/end time in + * the statistics. + */ + protected void encodeAndBuildChunks() throws IOException, PageException { final long startTime = timePartitionSlot.getStartTime(); - final long endTime = TimePartitionUtils.getTimePartitionEndTime(startTime); - final int satisfiedLength = satisfiedLengthQueue.poll(); - // serialize needDecode==true - dataSize += ReadWriteIOUtils.write(true, stream); - dataSize += ReadWriteIOUtils.write(satisfiedLength, stream); - satisfiedLengthQueue.offer(satisfiedLength); - - // The time column data is not included after serialization, so each value page needs to record - // its start time and end time. - long pageStartTime = Long.MAX_VALUE, pageEndTime = Long.MIN_VALUE; - for (int i = 0; i < times.length; i++) { - if (times[i] > endTime) { - break; - } - if (times[i] >= startTime) { - if (values.length == 0 || values[i] == null) { - dataSize += ReadWriteIOUtils.write(true, stream); - } else { - pageStartTime = Math.min(pageStartTime, times[i]); - pageEndTime = Math.max(pageEndTime, times[i]); - dataSize += ReadWriteIOUtils.write(false, stream); - switch (dataType) { - case INT32: - case DATE: - dataSize += ReadWriteIOUtils.write(values[i].getInt(), stream); - break; - case INT64: - case TIMESTAMP: - dataSize += ReadWriteIOUtils.write(values[i].getLong(), stream); - break; - case FLOAT: - dataSize += ReadWriteIOUtils.write(values[i].getFloat(), stream); - break; - case DOUBLE: - dataSize += ReadWriteIOUtils.write(values[i].getDouble(), stream); - break; - case BOOLEAN: - dataSize += ReadWriteIOUtils.write(values[i].getBoolean(), stream); - break; - case TEXT: - case BLOB: - case OBJECT: - case STRING: - dataSize += ReadWriteIOUtils.write(values[i].getBinary(), stream); - break; - default: - throw new UnSupportedDataTypeException( - String.format( - StorageEngineMessages.STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, - dataType)); - } - } - } + long endTime = startTime + TimePartitionUtils.getTimePartitionInterval() - 1; + if (endTime <= startTime) { + endTime = Long.MAX_VALUE; } - dataSize += ReadWriteIOUtils.write(pageStartTime, stream); - dataSize += ReadWriteIOUtils.write(pageEndTime, stream); - } - @Override - protected void writeChunkToWriter(final InputStream stream, final TsFileIOWriter writer) - throws IOException, PageException { + // Process each standalone value chunk for (int i = 0; i < chunkHeaderList.size(); i++) { - ChunkHeader chunkHeader = chunkHeaderList.get(i); + final ChunkHeader chunkHeader = chunkHeaderList.get(i); MeasurementSchema measurementSchema = new MeasurementSchema( chunkHeader.getMeasurementID(), chunkHeader.getDataType(), chunkHeader.getEncodingType(), chunkHeader.getCompressionType()); + ValueChunkWriter valueChunkWriter = new ValueChunkWriter( measurementSchema.getMeasurementName(), @@ -141,87 +86,102 @@ protected void writeChunkToWriter(final InputStream stream, final TsFileIOWriter measurementSchema.getType(), measurementSchema.getEncodingType(), measurementSchema.getValueEncoder()); - buildValueChunkWriter(stream, chunkHeader, pageNumbers.get(i), valueChunkWriter); - valueChunkWriter.writeToFileWriter(writer); - } - } - private void buildValueChunkWriter( - final InputStream stream, - final ChunkHeader chunkHeader, - final int pageNumber, - final ValueChunkWriter valueChunkWriter) - throws IOException, PageException { - boolean needDecode; - for (int i = 0; i < pageNumber; i++) { - needDecode = ReadWriteIOUtils.readBool(stream); - if (!needDecode) { - final PageHeader pageHeader = - PageHeader.deserializeFrom(stream, chunkHeader.getDataType(), true); - valueChunkWriter.writePageHeaderAndDataIntoBuff( - ByteBuffer.wrap(ReadWriteIOUtils.readBytesWithSelfDescriptionLength(stream)), - pageHeader); - continue; - } + // Filter and process the pages that belong to the current value chunk (index 'i') + for (PageBuffer page : pageBuffers) { + if (page.chunkIndex != i) { + continue; + } - // Since there is no data in the time column, we cannot provide the corresponding time for - // each value. - // However, if we only need to construct ValueChunk, we only need to provide the start time - // and end time of each page to construct the statistic. Therefore, when writing data to - // ValueChunkWriter, first use the wrong time. Before sealing each page, modify the statistic - // in the page writer and set the correct start time and end time. - final int length = ReadWriteIOUtils.readInt(stream); - for (int j = 0; j < length; j++) { - final boolean isNull = ReadWriteIOUtils.readBool(stream); - switch (chunkHeader.getDataType()) { - case INT32: - case DATE: - final int int32Value = isNull ? DEFAULT_INT32 : ReadWriteIOUtils.readInt(stream); - valueChunkWriter.write(0, int32Value, isNull); - break; - case INT64: - case TIMESTAMP: - final long int64Value = isNull ? DEFAULT_INT64 : ReadWriteIOUtils.readLong(stream); - valueChunkWriter.write(0, int64Value, isNull); - break; - case FLOAT: - final float floatValue = isNull ? DEFAULT_FLOAT : ReadWriteIOUtils.readFloat(stream); - valueChunkWriter.write(0, floatValue, isNull); - break; - case DOUBLE: - final double doubleValue = - isNull ? DEFAULT_DOUBLE : ReadWriteIOUtils.readDouble(stream); - valueChunkWriter.write(0, doubleValue, isNull); - break; - case BOOLEAN: - final boolean boolValue = isNull ? DEFAULT_BOOLEAN : ReadWriteIOUtils.readBool(stream); - valueChunkWriter.write(0, boolValue, isNull); - break; - case TEXT: - case BLOB: - case OBJECT: - case STRING: - final Binary binaryValue = - isNull ? DEFAULT_BINARY : ReadWriteIOUtils.readBinary(stream); - valueChunkWriter.write(0, binaryValue, isNull); + if (!page.needDecode) { + valueChunkWriter.writePageHeaderAndDataIntoBuff(page.pageData, page.pageHeader); + continue; + } + + // Keep track of actual bounds since we will write '0' as dummy time + long pageStartTime = Long.MAX_VALUE; + long pageEndTime = Long.MIN_VALUE; + + for (int j = 0; j < page.satisfiedLength; j++) { + long time = page.timeBatch[j]; + + // Apply time partition bounds check + if (time > endTime) { break; - default: - throw new UnSupportedDataTypeException( - String.format( - StorageEngineMessages.STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, - chunkHeader.getDataType())); + } + + if (time >= startTime) { + boolean isNull = + (page.valueBatch == null + || page.valueBatch.length == 0 + || page.valueBatch[j] == null); + if (!isNull) { + pageStartTime = Math.min(pageStartTime, time); + pageEndTime = Math.max(pageEndTime, time); + + TsPrimitiveType val = page.valueBatch[j]; + // Write 0 as a dummy time, valueChunkWriter will ignore it anyway for value columns + switch (chunkHeader.getDataType()) { + case INT32: + case DATE: + valueChunkWriter.write(0, val.getInt(), false); + break; + case INT64: + case TIMESTAMP: + valueChunkWriter.write(0, val.getLong(), false); + break; + case FLOAT: + valueChunkWriter.write(0, val.getFloat(), false); + break; + case DOUBLE: + valueChunkWriter.write(0, val.getDouble(), false); + break; + case BOOLEAN: + valueChunkWriter.write(0, val.getBoolean(), false); + break; + case TEXT: + case BLOB: + case STRING: + valueChunkWriter.write(0, val.getBinary(), false); + break; + default: + throw new UnSupportedDataTypeException( + String.format( + StorageEngineMessages + .STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, + chunkHeader.getDataType())); + } + } else { + valueChunkWriter.write(0, 0L, true); // Dummy write for NULL + } + } } + + // Correct the page statistics with the actual Start/End time + Statistics statistics = + valueChunkWriter.getPageWriter().getStatistics(); + if (pageStartTime <= pageEndTime) { + statistics.setStartTime(pageStartTime); + statistics.setEndTime(pageEndTime); + } + + valueChunkWriter.sealCurrentPage(); } - Statistics statistics = - valueChunkWriter.getPageWriter().getStatistics(); - // set the correct start time and end time in statistics - long startTime = ReadWriteIOUtils.readLong(stream); - long endTime = ReadWriteIOUtils.readLong(stream); - statistics.setStartTime(startTime); - statistics.setEndTime(endTime); + // Directly assemble the final standard Chunk into memory + final ByteBuffer valueData = valueChunkWriter.getByteBuffer(); + final Statistics valueStat = valueChunkWriter.getStatistics(); + final ChunkHeader encodedHeader = + createEncodedChunkHeader( + chunkHeader, + valueData.remaining(), + valueChunkWriter.getNumOfPages(), + VALUE_COLUMN_MASK); + + entireChunks.add(new Chunk(encodedHeader, valueData, null, valueStat)); - valueChunkWriter.sealCurrentPage(); + // Update global size (optional, since it will be fully recalculated during serialize anyway) + dataSize += valueData.remaining() + valueStat.getSerializedSize(); } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java index cddd35deafa57..342511ad6d4e7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java @@ -26,12 +26,15 @@ import org.apache.tsfile.file.header.PageHeader; import org.apache.tsfile.file.metadata.IChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.writer.TsFileIOWriter; import java.io.IOException; import java.io.InputStream; import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.List; public interface ChunkData extends TsFileData { IDeviceID getDevice(); @@ -52,6 +55,24 @@ public interface ChunkData extends TsFileData { void writeToFileWriter(TsFileIOWriter writer) throws IOException, PageException; + default void endChunk() {} + + default List getChunks() { + return Collections.emptyList(); + } + + record ChunkLayout( + long chunkGroupIndex, + long chunkGroupHeaderOffset, + long offset, + long length, + int chunkIndexInGroup, + boolean firstChunkOfGroup) {} + + ChunkLayout getChunkLayout(); + + void setChunkLayout(ChunkLayout layout); + @Override default TsFileDataType getType() { return TsFileDataType.CHUNK; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkEncoder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkEncoder.java deleted file mode 100644 index 6299c77a8a078..0000000000000 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkEncoder.java +++ /dev/null @@ -1,54 +0,0 @@ -/* - * 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 - */ -package org.apache.iotdb.db.storageengine.load.splitter; - -import org.apache.tsfile.file.metadata.ChunkMetadata; -import org.apache.tsfile.write.writer.TsFileIOWriter; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.List; - -/** Encodes an undecoded source chunk into its physical TsFile chunk bytes. */ -public final class ChunkEncoder { - - private ChunkEncoder() {} - - public static EncodedChunkGroup encode(final ChunkData chunkData) throws IOException { - final MemoryTsFileOutput output = new MemoryTsFileOutput(); - final TsFileIOWriter writer = new TsFileIOWriter(output); - writer.startChunkGroup(chunkData.getDevice()); - final int chunkStart = (int) writer.getPos(); - try { - chunkData.writeToFileWriter(writer); - } catch (Exception e) { - throw new IOException(e); - } - final int chunkEnd = (int) writer.getPos(); - final List metadata = writer.getChunkMetadataListOfCurrentDeviceInMemory(); - final List chunks = new ArrayList<>(metadata.size()); - for (int i = 0; i < metadata.size(); i++) { - final int start = (int) metadata.get(i).getOffsetOfChunkHeader(); - final int end = - i + 1 < metadata.size() ? (int) metadata.get(i + 1).getOffsetOfChunkHeader() : chunkEnd; - if (end > start && start >= chunkStart) { - chunks.add( - new EncodedChunk( - chunkData.getDevice(), - chunkData.getTimePartitionSlot(), - output.copyRange(start, end - start), - new ChunkMetadata(metadata.get(i)))); - } - } - return new EncodedChunkGroup(chunkData.getDevice(), chunkData.getTimePartitionSlot(), chunks); - } -} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunk.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunk.java deleted file mode 100644 index 512ea5cef55f1..0000000000000 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunk.java +++ /dev/null @@ -1,62 +0,0 @@ -/* - * 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 - */ -package org.apache.iotdb.db.storageengine.load.splitter; - -import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; - -import org.apache.tsfile.file.metadata.IChunkMetadata; -import org.apache.tsfile.file.metadata.IDeviceID; - -/** Physical Chunk bytes and the metadata required to seal the TsFile footer. */ -public final class EncodedChunk { - private final IDeviceID device; - private final TTimePartitionSlot timePartitionSlot; - private final byte[] bytes; - private final IChunkMetadata metadata; - private long offset; - - public EncodedChunk( - IDeviceID device, - TTimePartitionSlot timePartitionSlot, - byte[] bytes, - IChunkMetadata metadata) { - this.device = device; - this.timePartitionSlot = timePartitionSlot; - this.bytes = bytes; - this.metadata = metadata; - this.offset = 0; - } - - public IDeviceID getDevice() { - return device; - } - - public TTimePartitionSlot getTimePartitionSlot() { - return timePartitionSlot; - } - - public byte[] getBytes() { - return bytes; - } - - public IChunkMetadata getMetadata() { - return metadata; - } - - public long getOffset() { - return offset; - } - - public void setOffset(long offset) { - this.offset = offset; - } -} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunkGroup.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunkGroup.java deleted file mode 100644 index 6761b312a049a..0000000000000 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/EncodedChunkGroup.java +++ /dev/null @@ -1,68 +0,0 @@ -/* - * 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. - * The ASF licenses this file to you under the Apache License, Version 2.0. - */ -package org.apache.iotdb.db.storageengine.load.splitter; - -import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; - -import org.apache.tsfile.file.metadata.IDeviceID; -import org.apache.tsfile.file.metadata.StringArrayDeviceID; -import org.apache.tsfile.utils.ReadWriteIOUtils; - -import java.io.DataOutputStream; -import java.io.IOException; -import java.util.List; - -public final class EncodedChunkGroup implements TsFileData { - private final IDeviceID device; - private final TTimePartitionSlot timePartitionSlot; - private final List chunks; - - public EncodedChunkGroup( - IDeviceID device, TTimePartitionSlot timePartitionSlot, List chunks) { - this.device = device; - this.timePartitionSlot = timePartitionSlot; - this.chunks = List.copyOf(chunks); - } - - public IDeviceID getDevice() { - return device; - } - - public TTimePartitionSlot getTimePartitionSlot() { - return timePartitionSlot; - } - - public List getChunks() { - return chunks; - } - - @Override - public long getDataSize() { - return chunks.stream().mapToLong(chunk -> chunk.getBytes().length).sum(); - } - - @Override - public TsFileDataType getType() { - return TsFileDataType.ENCODED_CHUNK; - } - - @Override - public void serialize(DataOutputStream stream) throws IOException { - ReadWriteIOUtils.write(getType().ordinal(), stream); - ReadWriteIOUtils.write(timePartitionSlot.getStartTime(), stream); - ReadWriteIOUtils.write(device instanceof StringArrayDeviceID, stream); - device.serialize(stream); - ReadWriteIOUtils.write(chunks.size(), stream); - for (EncodedChunk chunk : chunks) { - ReadWriteIOUtils.write(chunk.getMetadata().getMeasurementUid(), stream); - ReadWriteIOUtils.write(chunk.getMetadata().getDataType().ordinal(), stream); - ReadWriteIOUtils.write(chunk.getOffset(), stream); - ReadWriteIOUtils.write(chunk.getBytes().length, stream); - stream.write(chunk.getBytes()); - chunk.getMetadata().serializeTo(stream, true); - } - } -} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/MemoryTsFileOutput.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/MemoryTsFileOutput.java deleted file mode 100644 index cdcf2365a4f9b..0000000000000 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/MemoryTsFileOutput.java +++ /dev/null @@ -1,72 +0,0 @@ -/* - * 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. - * The ASF licenses this file to you under the Apache License, Version 2.0. - */ -package org.apache.iotdb.db.storageengine.load.splitter; - -import org.apache.tsfile.utils.PublicBAOS; -import org.apache.tsfile.write.writer.TsFileOutput; - -import java.io.IOException; -import java.io.OutputStream; -import java.nio.ByteBuffer; - -final class MemoryTsFileOutput implements TsFileOutput { - private final PublicBAOS output = new PublicBAOS(); - - @Override - public void write(byte[] bytes) throws IOException { - output.write(bytes); - } - - @Override - public void write(byte value) { - output.write(value); - } - - @Override - public void write(ByteBuffer buffer) { - final int length = buffer.remaining(); - if (buffer.hasArray()) { - output.write(buffer.array(), buffer.arrayOffset() + buffer.position(), length); - buffer.position(buffer.position() + length); - } else { - final byte[] bytes = new byte[length]; - buffer.get(bytes); - output.write(bytes, 0, bytes.length); - } - } - - @Override - public long getPosition() { - return output.size(); - } - - @Override - public void close() throws IOException { - output.close(); - } - - @Override - public OutputStream wrapAsStream() { - return output; - } - - @Override - public void flush() {} - - @Override - public void truncate(long size) { - output.reset(); - } - - @Override - public void force() {} - - byte[] copyRange(int offset, int length) { - final byte[] result = new byte[length]; - System.arraycopy(output.getBuf(), offset, result, 0, length); - return result; - } -} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java index fdba8fa3d88ca..05dec3cc622ae 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java @@ -22,8 +22,9 @@ import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.utils.TimePartitionUtils; import org.apache.iotdb.db.i18n.StorageEngineMessages; -import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.exception.write.PageException; import org.apache.tsfile.file.header.ChunkHeader; import org.apache.tsfile.file.header.PageHeader; @@ -34,7 +35,6 @@ import org.apache.tsfile.file.metadata.statistics.Statistics; import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.utils.Binary; -import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.UnSupportedDataTypeException; import org.apache.tsfile.write.chunk.ChunkWriterImpl; @@ -43,46 +43,84 @@ import javax.annotation.Nonnull; +import java.io.DataInputStream; import java.io.DataOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.Serializable; import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; import static org.apache.iotdb.db.storageengine.load.LoadTsFileManager.MEASUREMENT_ID_CACHE; +import static org.apache.tsfile.common.constant.TsFileConstant.TIME_COLUMN_MASK; +import static org.apache.tsfile.common.constant.TsFileConstant.VALUE_COLUMN_MASK; public class NonAlignedChunkData implements ChunkData { private final TTimePartitionSlot timePartitionSlot; private final IDeviceID device; private final ChunkHeader chunkHeader; + private final EncryptParameter encryptParameter; - private final PublicBAOS byteStream; - private final DataOutputStream stream; private long dataSize; private boolean needDecodeChunk; - private int pageNumber; - private ChunkWriterImpl chunkWriter; + // Buffer for unencoded pages + private final List pageBuffers = new ArrayList<>(); + // The final standard Chunk object private Chunk chunk; + private ChunkLayout chunkLayout; + + private static class PageBuffer { + final boolean needDecode; + final PageHeader pageHeader; + final ByteBuffer pageData; + final long[] timeBatch; + final Object[] valueBatch; + final int satisfiedLength; + + PageBuffer(PageHeader pageHeader, ByteBuffer pageData) { + this.needDecode = false; + this.pageHeader = pageHeader; + this.pageData = pageData; + this.timeBatch = null; + this.valueBatch = null; + this.satisfiedLength = 0; + } + + PageBuffer(long[] timeBatch, Object[] valueBatch, int satisfiedLength) { + this.needDecode = true; + this.pageHeader = null; + this.pageData = null; + this.timeBatch = timeBatch; + this.valueBatch = valueBatch; + this.satisfiedLength = satisfiedLength; + } + } public NonAlignedChunkData( @Nonnull final IDeviceID device, final ChunkHeader chunkHeader, final TTimePartitionSlot timePartitionSlot) { + this(device, chunkHeader, timePartitionSlot, EncryptUtils.getEncryptParameter()); + } + + public NonAlignedChunkData( + @Nonnull final IDeviceID device, + final ChunkHeader chunkHeader, + final TTimePartitionSlot timePartitionSlot, + final EncryptParameter encryptParameter) { this.dataSize = 0; this.device = device; this.chunkHeader = chunkHeader; this.timePartitionSlot = timePartitionSlot; + this.encryptParameter = encryptParameter; this.needDecodeChunk = true; - this.pageNumber = 0; - this.byteStream = new PublicBAOS(); - this.stream = new DataOutputStream(byteStream); - addAttrDataSize(); } - private void addAttrDataSize() { // should be init before serialize, corresponding serializeAttr + private void addAttrDataSize() { dataSize += 2 * Byte.BYTES; // isModification and isAligned dataSize += Long.BYTES; // timePartitionSlot dataSize += device.serializedSize(); // device @@ -119,188 +157,176 @@ public boolean isEntireChunk() { return !needDecodeChunk; } + // ----------------------- Data Collection (In-Memory) ----------------------- + @Override - public void writeToFileWriter(final TsFileIOWriter writer) throws IOException { - ensureDataReadyForWriting(); - if (chunk != null) { - writer.writeChunk(chunk); - } else { - chunkWriter.writeToFileWriter(writer); - } + public void writeEntireChunk(final ByteBuffer chunkData, final IChunkMetadata chunkMetadata) { + this.chunk = new Chunk(chunkHeader, chunkData, null, chunkMetadata.getStatistics()); + this.needDecodeChunk = false; + this.dataSize += chunkData.remaining() + chunkMetadata.getStatistics().getSerializedSize(); } - private void ensureDataReadyForWriting() throws IOException { - if (chunk != null || chunkWriter != null) { - return; - } - - try { - deserializeTsFileData( - new LoadTsFilePieceNode.ByteBufferInputStream( - ByteBuffer.wrap(byteStream.getBuf(), 0, byteStream.size()))); - } catch (final PageException e) { - throw new IOException(e); - } + @Override + public void writeEntirePage(final PageHeader pageHeader, final ByteBuffer pageData) { + pageBuffers.add(new PageBuffer(pageHeader, pageData)); } @Override - public void serialize(final DataOutputStream stream) throws IOException { - ReadWriteIOUtils.write(getType().ordinal(), stream); - ReadWriteIOUtils.write(isAligned(), stream); - serializeAttr(stream); - byteStream.writeTo(stream); - close(); + public void writeDecodePage( + final long[] times, final Object[] values, final int satisfiedLength) { + pageBuffers.add(new PageBuffer(times, values, satisfiedLength)); } - private void serializeAttr(final DataOutputStream stream) throws IOException { - ReadWriteIOUtils.write(timePartitionSlot.getStartTime(), stream); - - ReadWriteIOUtils.write(device instanceof StringArrayDeviceID, stream); - device.serialize(stream); - - ReadWriteIOUtils.write(dataSize, stream); - ReadWriteIOUtils.write(needDecodeChunk, stream); - chunkHeader.serializeTo(stream); // chunk header already serialize chunk type - if (needDecodeChunk) { - ReadWriteIOUtils.write(pageNumber, stream); + @Override + public List getChunks() { + if (chunk == null) { + endChunk(); } + return chunk == null ? List.of() : List.of(chunk); } @Override - public void writeEntireChunk(final ByteBuffer chunkData, final IChunkMetadata chunkMetadata) - throws IOException { - dataSize += ReadWriteIOUtils.write(chunkData, stream); - dataSize += chunkMetadata.getStatistics().serialize(stream); + public ChunkLayout getChunkLayout() { + return chunkLayout; } @Override - public void writeEntirePage(final PageHeader pageHeader, final ByteBuffer pageData) - throws IOException { - pageNumber += 1; - dataSize += ReadWriteIOUtils.write(false, stream); - dataSize += pageHeader.serializeTo(stream); - dataSize += ReadWriteIOUtils.write(pageData, stream); + public void setChunkLayout(final ChunkLayout chunkLayout) { + this.chunkLayout = chunkLayout; } + // ----------------------- Unified Encoder Trigger ----------------------- + @Override - public void writeDecodePage(final long[] times, final Object[] values, final int satisfiedLength) - throws IOException { - pageNumber += 1; - final long startTime = timePartitionSlot.getStartTime(); - final long endTime = TimePartitionUtils.getTimePartitionEndTime(startTime); - dataSize += ReadWriteIOUtils.write(true, stream); - dataSize += ReadWriteIOUtils.write(satisfiedLength, stream); - - for (int i = 0; i < times.length; i++) { - if (times[i] > endTime) { - break; - } - if (times[i] >= startTime) { - dataSize += ReadWriteIOUtils.write(times[i], stream); - switch (chunkHeader.getDataType()) { - case INT32: - case DATE: - dataSize += ReadWriteIOUtils.write((int) values[i], stream); - break; - case INT64: - case TIMESTAMP: - dataSize += ReadWriteIOUtils.write((long) values[i], stream); - break; - case FLOAT: - dataSize += ReadWriteIOUtils.write((float) values[i], stream); - break; - case DOUBLE: - dataSize += ReadWriteIOUtils.write((double) values[i], stream); - break; - case BOOLEAN: - dataSize += ReadWriteIOUtils.write((boolean) values[i], stream); - break; - case TEXT: - case BLOB: - case STRING: - dataSize += ReadWriteIOUtils.write((Binary) values[i], stream); - break; - default: - throw new UnSupportedDataTypeException( - String.format( - StorageEngineMessages.STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, - chunkHeader.getDataType())); - } - } + public void endChunk() { + if (pageBuffers.isEmpty()) { + return; } - } - - private void deserializeTsFileData(final InputStream stream) throws IOException, PageException { - if (needDecodeChunk) { - buildChunkWriter(stream); - } else { - deserializeEntireChunk(stream); + try { + encodeAndBuildChunk(); + pageBuffers.clear(); + needDecodeChunk = false; + } catch (final IOException | PageException e) { + throw new IllegalStateException("Failed to encode chunk using ChunkWriterImpl", e); } } - private void deserializeEntireChunk(final InputStream stream) throws IOException { - final ByteBuffer chunkData = - ByteBuffer.wrap(ReadWriteIOUtils.readBytesWithSelfDescriptionLength(stream)); - final Statistics statistics = - Statistics.deserialize(stream, chunkHeader.getDataType()); - chunk = new Chunk(chunkHeader, chunkData, null, statistics); - } - - private void buildChunkWriter(final InputStream stream) throws IOException, PageException { - chunkWriter = + private void encodeAndBuildChunk() throws IOException, PageException { + ChunkWriterImpl chunkWriter = new ChunkWriterImpl( new MeasurementSchema( chunkHeader.getMeasurementID(), chunkHeader.getDataType(), chunkHeader.getEncodingType(), - chunkHeader.getCompressionType())); - boolean needDecode; - for (int j = 0; j < pageNumber; j++) { - needDecode = ReadWriteIOUtils.readBool(stream); - if (needDecode) { - final int length = ReadWriteIOUtils.readInt(stream); - for (int i = 0; i < length; i++) { - final long time = ReadWriteIOUtils.readLong(stream); - switch (chunkHeader.getDataType()) { - case INT32: - case DATE: - chunkWriter.write(time, ReadWriteIOUtils.readInt(stream)); - break; - case INT64: - case TIMESTAMP: - chunkWriter.write(time, ReadWriteIOUtils.readLong(stream)); - break; - case FLOAT: - chunkWriter.write(time, ReadWriteIOUtils.readFloat(stream)); - break; - case DOUBLE: - chunkWriter.write(time, ReadWriteIOUtils.readDouble(stream)); - break; - case BOOLEAN: - chunkWriter.write(time, ReadWriteIOUtils.readBool(stream)); - break; - case TEXT: - case BLOB: - case STRING: - chunkWriter.write(time, ReadWriteIOUtils.readBinary(stream)); - break; - default: - throw new UnSupportedDataTypeException( - String.format( - StorageEngineMessages.STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, - chunkHeader.getDataType())); - } - } + chunkHeader.getCompressionType()), + encryptParameter); + for (PageBuffer page : pageBuffers) { + if (page.needDecode) { + for (int j = 0; j < page.satisfiedLength; j++) { + writePointToChunkWriter(chunkWriter, page.timeBatch[j], page.valueBatch[j]); + } chunkWriter.sealCurrentPage(); } else { - final PageHeader pageHeader = - PageHeader.deserializeFrom(stream, chunkHeader.getDataType(), true); - chunkWriter.writePageHeaderAndDataIntoBuff( - ByteBuffer.wrap(ReadWriteIOUtils.readBytesWithSelfDescriptionLength(stream)), - pageHeader); + chunkWriter.writePageHeaderAndDataIntoBuff(page.pageData, page.pageHeader); } } + + // The buffered pages may have been split and re-encoded, so the source ChunkHeader no longer + // describes the resulting payload. Rebuild it from the actual encoded bytes and page count. + final ByteBuffer chunkData = chunkWriter.getByteBuffer(); + final Statistics statistics = chunkWriter.getStatistics(); + final ChunkHeader encodedHeader = + new ChunkHeader( + chunkHeader.getMeasurementID(), + chunkData.remaining(), + chunkHeader.getDataType(), + chunkHeader.getCompressionType(), + chunkHeader.getEncodingType(), + chunkWriter.getNumOfPages(), + chunkHeader.getChunkType() & (TIME_COLUMN_MASK | VALUE_COLUMN_MASK)); + this.chunk = new Chunk(encodedHeader, chunkData, null, statistics); + + // Update global data size + this.dataSize += chunkData.remaining() + statistics.getSerializedSize(); + } + + private void writePointToChunkWriter( + ChunkWriterImpl chunkWriter, final long time, final Object value) { + switch (chunkHeader.getDataType()) { + case INT32: + case DATE: + chunkWriter.write(time, (int) value); + break; + case INT64: + case TIMESTAMP: + chunkWriter.write(time, (long) value); + break; + case FLOAT: + chunkWriter.write(time, (float) value); + break; + case DOUBLE: + chunkWriter.write(time, (double) value); + break; + case BOOLEAN: + chunkWriter.write(time, (boolean) value); + break; + case TEXT: + case BLOB: + case STRING: + chunkWriter.write(time, (Binary) value); + break; + default: + throw new UnSupportedDataTypeException( + String.format( + StorageEngineMessages.STORAGE_EXCEPTION_DATA_TYPE_S_IS_NOT_SUPPORTED_5D5C02E4, + chunkHeader.getDataType())); + } + } + + // ----------------------- Serialization / Deserialization ----------------------- + + @Override + public void serialize(final DataOutputStream stream) throws IOException { + ReadWriteIOUtils.write(getType().ordinal(), stream); + ReadWriteIOUtils.write(isAligned(), stream); + serializeAttr(stream); + writeChunkLayout(stream); + + // Ensure all internal pages are encoded to chunks if not already + if (!pageBuffers.isEmpty()) { + endChunk(); + } + + // Write the standard chunk out natively + if (chunk != null) { + chunk.getChunkStatistic().serialize(stream); + ByteBuffer data = chunk.getData(); + ReadWriteIOUtils.write(data.remaining(), stream); + stream.write(data.array(), data.arrayOffset() + data.position(), data.remaining()); + } + } + + private void serializeAttr(final DataOutputStream stream) throws IOException { + ReadWriteIOUtils.write(timePartitionSlot.getStartTime(), stream); + ReadWriteIOUtils.write(device instanceof StringArrayDeviceID, stream); + device.serialize(stream); + ReadWriteIOUtils.write(dataSize, stream); + // Data has been standardized, so needDecodeChunk is always false on the wire + ReadWriteIOUtils.write(false, stream); + chunkHeader.serializeTo(stream); + } + + private void writeChunkLayout(final DataOutputStream stream) throws IOException { + ReadWriteIOUtils.write(chunkLayout != null, stream); + if (chunkLayout != null) { + ReadWriteIOUtils.write(chunkLayout.chunkGroupIndex(), stream); + ReadWriteIOUtils.write(chunkLayout.chunkGroupHeaderOffset(), stream); + ReadWriteIOUtils.write(chunkLayout.offset(), stream); + ReadWriteIOUtils.write(chunkLayout.length(), stream); + ReadWriteIOUtils.write(chunkLayout.chunkIndexInGroup(), stream); + ReadWriteIOUtils.write(chunkLayout.firstChunkOfGroup(), stream); + } } public static NonAlignedChunkData deserialize(final InputStream stream) @@ -313,29 +339,57 @@ public static NonAlignedChunkData deserialize(final InputStream stream) ? StringArrayDeviceID.deserialize(stream) : PlainDeviceID.deserialize(stream).convertToStringArrayDeviceId(); final long dataSize = ReadWriteIOUtils.readLong(stream); - final boolean needDecodeChunk = ReadWriteIOUtils.readBool(stream); + + // Ignore needDecodeChunk flag from stream since it's now always standardized to false + ReadWriteIOUtils.readBool(stream); + final byte chunkType = ReadWriteIOUtils.readByte(stream); final ChunkHeader chunkHeader = ChunkHeader.deserializeFrom(stream, chunkType); String measurementID = chunkHeader.getMeasurementID(); chunkHeader.setMeasurementID(MEASUREMENT_ID_CACHE.get(measurementID, m -> m)); - int pageNumber = 0; - if (needDecodeChunk) { - pageNumber = ReadWriteIOUtils.readInt(stream); - } final NonAlignedChunkData chunkData = new NonAlignedChunkData(device, chunkHeader, timePartitionSlot); - chunkData.needDecodeChunk = needDecodeChunk; - chunkData.pageNumber = pageNumber; - chunkData.deserializeTsFileData(stream); chunkData.dataSize = dataSize; - chunkData.close(); + chunkData.needDecodeChunk = false; // Always false after deserialization + + if (ReadWriteIOUtils.readBool(stream)) { + chunkData.chunkLayout = + new ChunkLayout( + ReadWriteIOUtils.readLong(stream), + ReadWriteIOUtils.readLong(stream), + ReadWriteIOUtils.readLong(stream), + ReadWriteIOUtils.readLong(stream), + ReadWriteIOUtils.readInt(stream), + ReadWriteIOUtils.readBool(stream)); + } + + chunkData.deserializeTsFileData(stream); + return chunkData; } - private void close() throws IOException { - byteStream.close(); - stream.close(); + private void deserializeTsFileData(final InputStream stream) throws IOException { + Statistics statistics = + Statistics.deserialize(stream, chunkHeader.getDataType()); + int dataBytesSize = ReadWriteIOUtils.readInt(stream); + byte[] dataBytes = new byte[dataBytesSize]; + new DataInputStream(stream).readFully(dataBytes); + + this.chunk = new Chunk(chunkHeader, ByteBuffer.wrap(dataBytes), null, statistics); + } + + // ----------------------- Write to TsFile ----------------------- + + @Override + public void writeToFileWriter(final TsFileIOWriter writer) throws IOException, PageException { + // Fallback trigger for encoding + if (!pageBuffers.isEmpty()) { + endChunk(); + } + if (chunk != null) { + writer.writeChunk(chunk); + } } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileData.java index e3fd43085087f..21b7c2046c424 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileData.java @@ -38,7 +38,11 @@ public interface TsFileData { static TsFileData deserialize(InputStream stream) throws IOException, PageException, IllegalPathException { - final TsFileDataType type = TsFileDataType.values()[ReadWriteIOUtils.readInt(stream)]; + final int typeOrdinal = ReadWriteIOUtils.readInt(stream); + if (typeOrdinal < 0 || typeOrdinal >= TsFileDataType.values().length) { + throw new IOException(); + } + final TsFileDataType type = TsFileDataType.values()[typeOrdinal]; switch (type) { case CHUNK: return ChunkData.deserialize(stream); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileDataType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileDataType.java index bf983ec923d6a..9e3f73f724dd5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileDataType.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileDataType.java @@ -21,6 +21,5 @@ public enum TsFileDataType { CHUNK, - ENCODED_CHUNK, DELETION } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java index 67c49bf42a5ac..59e6a39fa849a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java @@ -59,7 +59,6 @@ import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -73,7 +72,7 @@ public class TsFileSplitter { private final TsFileDataConsumer consumer; private Map offset2ChunkMetadata = new HashMap<>(); private List deletions = new ArrayList<>(); - private Map> pageIndex2ChunkData = new LinkedHashMap<>(); + private Map> pageIndex2ChunkData = new HashMap<>(); private Map pageIndex2Times = new HashMap<>(); private boolean isTimeChunkNeedDecode = true; private IDeviceID curDevice = null; @@ -95,6 +94,14 @@ public TsFileSplitter(File tsFile, TsFileDataConsumer consumer) { this.consumer = consumer; } + private ChunkData createChunkData( + final boolean aligned, + final IDeviceID device, + final ChunkHeader header, + final TTimePartitionSlot timePartitionSlot) { + return ChunkData.createChunkData(aligned, device, header, timePartitionSlot); + } + @SuppressWarnings({"squid:S3776", "squid:S6541"}) public void splitTsFileByDataPartition() throws IOException, LoadFileException, IllegalStateException { @@ -110,6 +117,7 @@ public void splitTsFileByDataPartition() tsFile.getPath())); } + reader.readFileMetadata(); reader.position((long) TSFileConfig.MAGIC_STRING.getBytes().length + 1); getChunkMetadata(reader, offset2ChunkMetadata); byte marker; @@ -177,7 +185,7 @@ private void processTimeChunkOrNonAlignedChunk(TsFileSequenceReader reader, byte == TsFileConstant.TIME_COLUMN_MASK); if (isAligned) { pageIndex2Times = new HashMap<>(); - pageIndex2ChunkData = new LinkedHashMap<>(); + pageIndex2ChunkData = new HashMap<>(); isTimeChunkNeedDecode = true; } @@ -190,8 +198,7 @@ private void processTimeChunkOrNonAlignedChunk(TsFileSequenceReader reader, byte } TTimePartitionSlot timePartitionSlot = TimePartitionUtils.getTimePartitionSlot(chunkMetadata.getStartTime()); - ChunkData chunkData = - ChunkData.createChunkData(isAligned, curDevice, header, timePartitionSlot); + ChunkData chunkData = createChunkData(isAligned, curDevice, header, timePartitionSlot); if (!needDecodeChunk(chunkMetadata)) { chunkData.setNotDecode(); @@ -245,10 +252,13 @@ private void decodeAndWriteTimeChunkOrNonAlignedChunk( TimePartitionUtils.getTimePartitionSlot(startTime); if (!timePartitionSlot.equals(pageTimePartitionSlot)) { if (!isAligned) { + chunkData.endChunk(); consumeChunkData(measurementId, chunkOffset, chunkData); + } else { + chunkData.endChunk(); } timePartitionSlot = pageTimePartitionSlot; - chunkData = ChunkData.createChunkData(isAligned, curDevice, header, timePartitionSlot); + chunkData = createChunkData(isAligned, curDevice, header, timePartitionSlot); } if (isAligned) { pageIndex2ChunkData @@ -268,24 +278,32 @@ private void decodeAndWriteTimeChunkOrNonAlignedChunk( int satisfiedLength = 0; long endTime = - TimePartitionUtils.getTimePartitionUpperBound(timePartitionSlot.getStartTime()); + timePartitionSlot.getStartTime() + TimePartitionUtils.getTimePartitionInterval(); + // beware of overflow + if (endTime <= timePartitionSlot.getStartTime()) { + endTime = Long.MAX_VALUE; + } for (int i = 0; i < times.length; i++) { - if (TimePartitionUtils.isAfterOrEqualToTimePartitionUpperBound( - times[i], timePartitionSlot.getStartTime(), endTime)) { + if (times[i] >= endTime) { chunkData.writeDecodePage(times, values, satisfiedLength); if (isAligned) { + chunkData.endChunk(); pageIndex2ChunkData .computeIfAbsent(pageIndex, o -> new ArrayList<>()) .add((AlignedChunkData) chunkData); } else { + chunkData.endChunk(); consumeChunkData(measurementId, chunkOffset, chunkData); } timePartitionSlot = TimePartitionUtils.getTimePartitionSlot(times[i]); satisfiedLength = 0; endTime = - TimePartitionUtils.getTimePartitionUpperBound(timePartitionSlot.getStartTime()); - chunkData = ChunkData.createChunkData(isAligned, curDevice, header, timePartitionSlot); + timePartitionSlot.getStartTime() + TimePartitionUtils.getTimePartitionInterval(); + if (endTime <= timePartitionSlot.getStartTime()) { + endTime = Long.MAX_VALUE; + } + chunkData = createChunkData(isAligned, curDevice, header, timePartitionSlot); } satisfiedLength += 1; } @@ -302,6 +320,7 @@ private void decodeAndWriteTimeChunkOrNonAlignedChunk( } if (!isAligned) { + chunkData.endChunk(); consumeChunkData(measurementId, chunkOffset, chunkData); } } @@ -448,7 +467,7 @@ private void consumeAllAlignedChunkData( return; } - Map chunkDataMap = new LinkedHashMap<>(); + Map chunkDataMap = new HashMap<>(); for (Map.Entry> entry : pageIndex2ChunkData.entrySet()) { List alignedChunkDataList = entry.getValue(); for (int i = 0; i < alignedChunkDataList.size(); i++) { @@ -459,6 +478,7 @@ private void consumeAllAlignedChunkData( } } for (AlignedChunkData chunkData : chunkDataMap.keySet()) { + chunkData.endChunk(); timePartitionSlots.add(chunkData.getTimePartitionSlot()); if (deletions.isEmpty() && timePartitionSlots.size() > CONFIG.getLoadTsFileSpiltPartitionMaxSize()) { @@ -477,7 +497,7 @@ private void consumeAllAlignedChunkData( chunkData)); } } - this.pageIndex2ChunkData = new LinkedHashMap<>(); + this.pageIndex2ChunkData = new HashMap<>(); } private void consumeChunkData(String measurement, long offset, ChunkData chunkData) diff --git a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java new file mode 100644 index 0000000000000..49eb75cf43289 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java @@ -0,0 +1,294 @@ +/* + * 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 + */ +package org.apache.tsfile.write; + +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.apache.tsfile.common.constant.TsFileConstant; +import org.apache.tsfile.file.MetaMarker; +import org.apache.tsfile.file.header.ChunkGroupHeader; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.MeasurementMetadataIndexEntry; +import org.apache.tsfile.file.metadata.MetadataIndexNode; +import org.apache.tsfile.file.metadata.TimeseriesMetadata; +import org.apache.tsfile.file.metadata.TsFileMetadata; +import org.apache.tsfile.file.metadata.enums.MetadataIndexNodeType; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.utils.BloomFilter; +import org.apache.tsfile.utils.BytesUtils; +import org.apache.tsfile.utils.ReadWriteIOUtils; +import org.apache.tsfile.write.writer.LocalTsFileOutput; +import org.apache.tsfile.write.writer.TsFileOutput; +import org.apache.tsfile.write.writer.tsmiterator.TSMIterator; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.TreeMap; + +import static org.apache.tsfile.file.metadata.MetadataIndexConstructor.checkAndBuildLevelIndex; +import static org.apache.tsfile.file.metadata.MetadataIndexConstructor.splitDeviceByTable; + +/** Low-level Chunk writer that writes precalculated physical Chunks directly by absolute offset. */ +public class TsFilePrecalculatedChunkWriter implements AutoCloseable { + + private static final Logger LOGGER = + LoggerFactory.getLogger(TsFilePrecalculatedChunkWriter.class); + + private final TsFileOutput out; + private final File file; + private final Map>> device2MetadataMap = + new TreeMap<>(); + + public TsFilePrecalculatedChunkWriter(File file) throws IOException { + this.file = file; + this.out = new LocalTsFileOutput(new FileOutputStream(file)); + startFile(); + } + + public TsFilePrecalculatedChunkWriter(TsFileOutput out) throws IOException { + this.file = null; + this.out = out; + startFile(); + } + + private void startFile() throws IOException { + out.write(BytesUtils.stringToBytes(TSFileConfig.MAGIC_STRING)); + out.write(new byte[] {TSFileConfig.VERSION_NUMBER}); + LOGGER.info( + "Started precalculated TsFile writer: file={}, headerEndOffset={}", + getFileForLog(), + out.getPosition()); + } + + public void writeChunk( + IDeviceID device, + boolean isAligned, + long chunkGroupHeaderOffset, + boolean isFirstChunkOfGroup, + Chunk chunk, + long chunkOffset) + throws IOException { + long currentPos = out.getPosition(); + OutputStream stream = out.wrapAsStream(); + final int chunkHeaderSize = chunk.getHeader().getSerializedSize(); + final int chunkDataSize = chunk.getData().remaining(); + final long chunkLength = chunkHeaderSize + (long) chunkDataSize; + + LOGGER.info( + "Writing precalculated Chunk: file={}, device={}, aligned={}, measurement={}, dataType={}, " + + "chunkType={}, pageCount={}, firstChunkOfGroup={}, expectedGroupHeaderOffset={}, " + + "expectedChunkOffset={}, actualOffsetBeforeWrite={}, chunkHeaderSize={}, " + + "chunkDataSize={}, chunkLength={}", + getFileForLog(), + device, + isAligned, + chunk.getHeader().getMeasurementID(), + chunk.getHeader().getDataType(), + chunk.getHeader().getChunkType(), + chunk.getHeader().getNumOfPages(), + isFirstChunkOfGroup, + chunkGroupHeaderOffset, + chunkOffset, + currentPos, + chunkHeaderSize, + chunkDataSize, + chunkLength); + + if (isFirstChunkOfGroup) { + checkOffsetMismatch("chunk group header", chunkGroupHeaderOffset, currentPos, device, chunk); + new ChunkGroupHeader(device).serializeTo(stream); + currentPos = out.getPosition(); + LOGGER.info( + "Wrote ChunkGroupHeader: file={}, device={}, headerOffset={}, headerEndOffset={}, " + + "headerLength={}", + getFileForLog(), + device, + chunkGroupHeaderOffset, + currentPos, + currentPos - chunkGroupHeaderOffset); + } + + checkOffsetMismatch("chunk", chunkOffset, currentPos, device, chunk); + chunk.getHeader().serializeTo(stream); + out.write(chunk.getData().duplicate()); + + LOGGER.info( + "Wrote precalculated Chunk: file={}, device={}, measurement={}, chunkOffset={}, " + + "chunkEndOffset={}, expectedChunkEndOffset={}, metadataMask={}", + getFileForLog(), + device, + chunk.getHeader().getMeasurementID(), + chunkOffset, + out.getPosition(), + chunkOffset + chunkLength, + chunk.getHeader().getChunkType() + & (TsFileConstant.TIME_COLUMN_MASK | TsFileConstant.VALUE_COLUMN_MASK)); + + ChunkMetadata metadata = + new ChunkMetadata( + chunk.getHeader().getMeasurementID(), + chunk.getHeader().getDataType(), + chunk.getHeader().getEncodingType(), + chunk.getHeader().getCompressionType(), + chunkOffset, + chunk.getChunkStatistic()); + metadata.setMask( + (byte) + (chunk.getHeader().getChunkType() + & (TsFileConstant.TIME_COLUMN_MASK | TsFileConstant.VALUE_COLUMN_MASK))); + device2MetadataMap + .computeIfAbsent(device, k -> new TreeMap<>()) + .computeIfAbsent(chunk.getHeader().getMeasurementID(), k -> new ArrayList<>()) + .add(metadata); + } + + private void checkOffsetMismatch( + final String positionType, + final long expectedOffset, + final long actualOffset, + final IDeviceID device, + final Chunk chunk) + throws IOException { + if (expectedOffset != actualOffset) { + LOGGER.error( + "Precalculated offset mismatch: file={}, positionType={}, device={}, measurement={}, " + + "expectedOffset={}, actualOffset={}, delta={}", + getFileForLog(), + positionType, + device, + chunk.getHeader().getMeasurementID(), + expectedOffset, + actualOffset, + actualOffset - expectedOffset); + throw new IOException( + String.format( + "Precalculated %s offset mismatch for %s/%s: expected=%d, actual=%d", + positionType, + device, + chunk.getHeader().getMeasurementID(), + expectedOffset, + actualOffset)); + } + } + + @Override + public void close() throws IOException { + final long metaOffset = out.getPosition(); + LOGGER.info( + "Closing precalculated TsFile writer: file={}, dataEndOffset={}, deviceCount={}, " + + "seriesCount={}, chunkCount={}", + getFileForLog(), + metaOffset, + device2MetadataMap.size(), + device2MetadataMap.values().stream().mapToInt(Map::size).sum(), + device2MetadataMap.values().stream() + .flatMap(measurements -> measurements.values().stream()) + .mapToInt(List::size) + .sum()); + final OutputStream stream = out.wrapAsStream(); + ReadWriteIOUtils.write(MetaMarker.SEPARATOR, stream); + final Map deviceMetadataIndexMap = new TreeMap<>(); + + int seriesCount = device2MetadataMap.values().stream().mapToInt(Map::size).sum(); + BloomFilter bloomFilter = + BloomFilter.getEmptyBloomFilter( + TSFileDescriptor.getInstance().getConfig().getBloomFilterErrorRate(), + Math.max(seriesCount, 1)); + + for (Map.Entry>> deviceEntry : + device2MetadataMap.entrySet()) { + IDeviceID deviceId = deviceEntry.getKey(); + Map> measurementMap = deviceEntry.getValue(); + + MetadataIndexNode measurementNode = + new MetadataIndexNode(MetadataIndexNodeType.LEAF_MEASUREMENT); + for (Map.Entry> measEntry : measurementMap.entrySet()) { + String measurementId = measEntry.getKey(); + List chunkMetadatas = measEntry.getValue(); + + final TimeseriesMetadata tsMetadata = + TSMIterator.constructOneTimeseriesMetadata(measurementId, chunkMetadatas); + final long tsMetadataOffset = out.getPosition(); + measurementNode.addEntry( + new MeasurementMetadataIndexEntry(measurementId, tsMetadataOffset)); + tsMetadata.serializeTo(stream); + bloomFilter.add(deviceId + "." + measurementId); + } + + measurementNode.setEndOffset(out.getPosition()); + deviceMetadataIndexMap.put(deviceId, measurementNode); + } + + final TsFileMetadata tsFileMetadata = new TsFileMetadata(); + final Map tableIndexNodeMap = new TreeMap<>(); + for (final Map.Entry> tableEntry : + splitDeviceByTable(deviceMetadataIndexMap).entrySet()) { + tableIndexNodeMap.put( + tableEntry.getKey(), checkAndBuildLevelIndex(tableEntry.getValue(), out)); + } + tsFileMetadata.setTableMetadataIndexNodeMap(tableIndexNodeMap); + tsFileMetadata.setMetaOffset(metaOffset); + tsFileMetadata.setBloomFilter(bloomFilter); + + long tsFileMetadataOffset = out.getPosition(); + tsFileMetadata.serializeTo(stream); + int tsFileMetadataSize = (int) (out.getPosition() - tsFileMetadataOffset); + + ReadWriteIOUtils.write(tsFileMetadataSize, stream); + out.write(BytesUtils.stringToBytes(TSFileConfig.MAGIC_STRING)); + LOGGER.info( + "Finished precalculated TsFile writer: file={}, metaOffset={}, tsFileMetadataOffset={}, " + + "tsFileMetadataSize={}, fileEndOffset={}", + getFileForLog(), + metaOffset, + tsFileMetadataOffset, + tsFileMetadataSize, + out.getPosition()); + out.close(); + } + + private String getFileForLog() { + return file == null ? "" : file.getAbsolutePath(); + } + + public TsFileOutput getOutput() { + return out; + } + + public File getFile() { + if (file == null) { + throw new IllegalStateException("This writer is not backed by a file"); + } + return file; + } + + public Map> getChunkMetadataListMap() { + final Map> result = new TreeMap<>(); + device2MetadataMap.forEach( + (device, measurementMap) -> + result.put( + device, + measurementMap.values().stream() + .flatMap(List::stream) + .collect(java.util.stream.Collectors.toList()))); + return result; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java index 7173f5cdc5884..e3d05f5f7535f 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java @@ -61,21 +61,6 @@ public void testPieceSerializeRoundTrip() { Assert.assertTrue(node.isGeneratedByPipe()); } - @Test - public void testPiecePreservesCallerChecksum() { - final LoadTsFileConsensusNode piece = - LoadTsFileConsensusNode.piece( - new PlanNodeId("checksum"), - "load", - "file", - 0L, - 0L, - Collections.emptyList(), - 987654321L); - - Assert.assertEquals(987654321L, piece.getChecksum()); - } - @Test public void testPieceRefSerializeRoundTrip() { final LoadTsFileConsensusNode.PieceRef ref = diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java index c824a6c0ce94b..f42d43775b830 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java @@ -30,11 +30,17 @@ import org.apache.tsfile.file.metadata.enums.TSEncoding; import org.apache.tsfile.file.metadata.statistics.Statistics; import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.utils.TsPrimitiveType; import org.apache.tsfile.write.writer.TsFileIOWriter; import org.junit.Test; import org.mockito.Mockito; import java.nio.ByteBuffer; +import java.util.List; + +import static org.apache.tsfile.common.constant.TsFileConstant.TIME_COLUMN_MASK; +import static org.apache.tsfile.common.constant.TsFileConstant.VALUE_COLUMN_MASK; +import static org.junit.Assert.assertEquals; public class ChunkDataDirectWriteTest { @@ -66,6 +72,53 @@ public void testAlignedChunkDataCanWriteWithoutSerdeRoundTrip() throws Exception Mockito.verify(writer).writeChunk(Mockito.any(Chunk.class)); } + @Test + public void testAlignedChunkDataAcceptsObjectArrayFromTimePageDecode() throws Exception { + final AlignedChunkData chunkData = createAlignedTimeChunkData(); + + chunkData.writeDecodePage(new long[] {1L, 2L}, new Object[] {null, null}, 2); + } + + @Test + public void testReencodedNonAlignedChunkRebuildsHeader() { + final NonAlignedChunkData chunkData = createNonAlignedChunkData(); + chunkData.writeDecodePage(new long[] {1L, 2L}, new Object[] {1, 2}, 2); + chunkData.endChunk(); + + final Chunk chunk = chunkData.getChunks().get(0); + assertEquals(chunk.getData().remaining(), chunk.getHeader().getDataSize()); + assertEquals(1, chunk.getHeader().getNumOfPages()); + } + + @Test + public void testAlignedChunkEndBuildsOnlyCurrentPhysicalChunk() throws Exception { + final AlignedChunkData chunkData = createAlignedTimeChunkData(); + final long[] times = new long[] {1L, 2L}; + chunkData.writeDecodePage(times, new Object[] {null, null}, times.length); + chunkData.endChunk(); + + final Chunk timeChunk = chunkData.getChunks().get(0); + assertEquals(1, chunkData.getChunks().size()); + assertEquals(timeChunk.getData().remaining(), timeChunk.getHeader().getDataSize()); + assertEquals(TIME_COLUMN_MASK, timeChunk.getHeader().getChunkType() & TIME_COLUMN_MASK); + + chunkData.addValueChunk(createChunkHeader()); + chunkData.writeDecodeValuePage( + times, + new TsPrimitiveType[] { + TsPrimitiveType.getByType(TSDataType.INT32, 1), + TsPrimitiveType.getByType(TSDataType.INT32, 2) + }, + TSDataType.INT32); + chunkData.endChunk(); + + final List chunks = chunkData.getChunks(); + assertEquals(2, chunks.size()); + final Chunk valueChunk = chunks.get(1); + assertEquals(valueChunk.getData().remaining(), valueChunk.getHeader().getDataSize()); + assertEquals(VALUE_COLUMN_MASK, valueChunk.getHeader().getChunkType() & VALUE_COLUMN_MASK); + } + private static Statistics createInt32Statistics() { final Statistics statistics = Statistics.getStatsByType(TSDataType.INT32); statistics.update(1L, 1); @@ -84,6 +137,21 @@ private static AlignedChunkData createAlignedChunkData() { ChunkData.createChunkData(true, device, createChunkHeader(), new TTimePartitionSlot(0L)); } + private static AlignedChunkData createAlignedTimeChunkData() { + final IDeviceID device = new StringArrayDeviceID("root", "sg", "d1"); + final ChunkHeader timeHeader = + new ChunkHeader( + "", + 1024, + TSDataType.VECTOR, + CompressionType.UNCOMPRESSED, + TSEncoding.PLAIN, + 2, + TIME_COLUMN_MASK); + return (AlignedChunkData) + ChunkData.createChunkData(true, device, timeHeader, new TTimePartitionSlot(0L)); + } + private static ChunkHeader createChunkHeader() { return new ChunkHeader( "temperature", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0); From 6fb805d59ec683a6b73e128d35986f6337962700 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Fri, 11 Sep 2026 10:26:15 +0800 Subject: [PATCH 3/8] fix(load): preserve chunk data across time partitions --- .../load/splitter/AlignedChunkData.java | 84 ++++++++++++------- .../BatchedAlignedValueChunkData.java | 10 +-- .../load/splitter/NonAlignedChunkData.java | 21 ++++- .../splitter/ChunkDataDirectWriteTest.java | 41 ++++++++- 4 files changed, 115 insertions(+), 41 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java index a49dbc3a00074..186e05c6b9f50 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java @@ -53,9 +53,11 @@ import java.io.InputStream; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.LinkedHashSet; import java.util.LinkedList; import java.util.List; import java.util.Queue; +import java.util.Set; import static org.apache.iotdb.db.storageengine.load.LoadTsFileManager.MEASUREMENT_ID_CACHE; import static org.apache.tsfile.common.constant.TsFileConstant.TIME_COLUMN_MASK; @@ -329,31 +331,37 @@ protected void encodeAndBuildChunks() throws IOException, PageException { writePageBuffersToWriter(writer); - // endChunk() seals one physical Chunk. Building all columns here would emit the time Chunk - // again whenever a later value Chunk is split and encoded. - final ChunkHeader sourceHeader = chunkHeaderList.get(currentChunkIndex); - final ByteBuffer encodedData; - final Statistics statistics; - final int pageCount; - final int mask; - if (currentChunkIndex == 0) { - encodedData = writer.getTimeChunkWriter().getByteBuffer(); - statistics = writer.getTimeChunkWriter().getStatistics(); - pageCount = writer.getTimeChunkWriter().getNumOfPages(); - mask = TIME_COLUMN_MASK; - } else { - final ValueChunkWriter valueWriter = - writer.getValueChunkWriterList().get(currentChunkIndex - 1); - encodedData = valueWriter.getByteBuffer(); - statistics = valueWriter.getStatistics(); - pageCount = valueWriter.getNumOfPages(); - mask = VALUE_COLUMN_MASK; + // One aligned ChunkData may accumulate pages from several physical Chunks before it is + // consumed. Encode every Chunk represented by the current buffers, while leaving Chunks that + // were sealed by an earlier endChunk() untouched. + final Set bufferedChunkIndexes = new LinkedHashSet<>(); + for (final PageBuffer pageBuffer : pageBuffers) { + bufferedChunkIndexes.add(pageBuffer.chunkIndex); } + for (final int chunkIndex : bufferedChunkIndexes) { + final ChunkHeader sourceHeader = chunkHeaderList.get(chunkIndex); + final ByteBuffer encodedData; + final Statistics statistics; + final int pageCount; + final int mask; + if (chunkIndex == 0) { + encodedData = writer.getTimeChunkWriter().getByteBuffer(); + statistics = writer.getTimeChunkWriter().getStatistics(); + pageCount = writer.getTimeChunkWriter().getNumOfPages(); + mask = TIME_COLUMN_MASK; + } else { + final ValueChunkWriter valueWriter = writer.getValueChunkWriterList().get(chunkIndex - 1); + encodedData = valueWriter.getByteBuffer(); + statistics = valueWriter.getStatistics(); + pageCount = valueWriter.getNumOfPages(); + mask = VALUE_COLUMN_MASK; + } - final ChunkHeader encodedHeader = - createEncodedChunkHeader(sourceHeader, encodedData.remaining(), pageCount, mask); - entireChunks.add(new Chunk(encodedHeader, encodedData, null, statistics)); - dataSize += encodedData.remaining() + statistics.getSerializedSize(); + final ChunkHeader encodedHeader = + createEncodedChunkHeader(sourceHeader, encodedData.remaining(), pageCount, mask); + entireChunks.add(new Chunk(encodedHeader, encodedData, null, statistics)); + dataSize += encodedData.remaining() + statistics.getSerializedSize(); + } } protected static ChunkHeader createEncodedChunkHeader( @@ -403,21 +411,41 @@ private void writePageBuffersToWriter(final AlignedChunkWriterImpl writer) continue; } if (page.valueBatch == null) { - for (int i = 0; i < page.satisfiedLength; i++) { - writer.writeTime(page.timeBatch[i]); + boolean hasPoint = false; + for (final long time : page.timeBatch) { + if (!isInCurrentTimePartition(time)) { + continue; + } + writer.writeTime(time); + hasPoint = true; + } + if (hasPoint) { + writer.sealCurrentTimePage(); } - writer.sealCurrentTimePage(); } else { final TSDataType dataType = chunkHeaderList.get(page.chunkIndex).getDataType(); - for (int i = 0; i < page.satisfiedLength; i++) { + boolean hasPoint = false; + for (int i = 0; i < page.timeBatch.length; i++) { + if (!isInCurrentTimePartition(page.timeBatch[i])) { + continue; + } writeValueToAlignedChunkWriter( writer, page.timeBatch[i], page.valueBatch[i], dataType, page.chunkIndex - 1); + hasPoint = true; + } + if (hasPoint) { + writer.sealCurrentValuePage(page.chunkIndex - 1); } - writer.sealCurrentValuePage(page.chunkIndex - 1); } } } + protected boolean isInCurrentTimePartition(final long time) { + final long partitionStart = timePartitionSlot.getStartTime(); + final long partitionEnd = partitionStart + TimePartitionUtils.getTimePartitionInterval() - 1; + return time >= partitionStart && (partitionEnd <= partitionStart || time <= partitionEnd); + } + private void writeValueToAlignedChunkWriter( final AlignedChunkWriterImpl writer, final long time, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java index 57fb7a01a23b6..4578bb3e2eb52 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java @@ -102,15 +102,9 @@ protected void encodeAndBuildChunks() throws IOException, PageException { long pageStartTime = Long.MAX_VALUE; long pageEndTime = Long.MIN_VALUE; - for (int j = 0; j < page.satisfiedLength; j++) { + for (int j = 0; j < page.timeBatch.length; j++) { long time = page.timeBatch[j]; - - // Apply time partition bounds check - if (time > endTime) { - break; - } - - if (time >= startTime) { + if (time >= startTime && time <= endTime) { boolean isNull = (page.valueBatch == null || page.valueBatch.length == 0 diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java index 05dec3cc622ae..54c47b513e0bf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java @@ -223,10 +223,20 @@ private void encodeAndBuildChunk() throws IOException, PageException { for (PageBuffer page : pageBuffers) { if (page.needDecode) { - for (int j = 0; j < page.satisfiedLength; j++) { - writePointToChunkWriter(chunkWriter, page.timeBatch[j], page.valueBatch[j]); + final long partitionStart = timePartitionSlot.getStartTime(); + final long partitionEnd = getTimePartitionEnd(partitionStart); + boolean hasPoint = false; + for (int j = 0; j < page.timeBatch.length; j++) { + final long time = page.timeBatch[j]; + if (time < partitionStart || time > partitionEnd) { + continue; + } + writePointToChunkWriter(chunkWriter, time, page.valueBatch[j]); + hasPoint = true; + } + if (hasPoint) { + chunkWriter.sealCurrentPage(); } - chunkWriter.sealCurrentPage(); } else { chunkWriter.writePageHeaderAndDataIntoBuff(page.pageData, page.pageHeader); } @@ -251,6 +261,11 @@ private void encodeAndBuildChunk() throws IOException, PageException { this.dataSize += chunkData.remaining() + statistics.getSerializedSize(); } + private static long getTimePartitionEnd(final long partitionStart) { + final long partitionEnd = partitionStart + TimePartitionUtils.getTimePartitionInterval() - 1; + return partitionEnd <= partitionStart ? Long.MAX_VALUE : partitionEnd; + } + private void writePointToChunkWriter( ChunkWriterImpl chunkWriter, final long time, final Object value) { switch (chunkHeader.getDataType()) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java index f42d43775b830..4e85ef36e1878 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.storageengine.load.splitter; import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.utils.TimePartitionUtils; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.header.ChunkHeader; @@ -119,6 +120,32 @@ public void testAlignedChunkEndBuildsOnlyCurrentPhysicalChunk() throws Exception assertEquals(VALUE_COLUMN_MASK, valueChunk.getHeader().getChunkType() & VALUE_COLUMN_MASK); } + @Test + public void testNonAlignedChunkSelectsPointsFromCurrentTimePartition() { + final long partitionStart = TimePartitionUtils.getTimePartitionInterval(); + final NonAlignedChunkData chunkData = + createNonAlignedChunkData(new TTimePartitionSlot(partitionStart)); + chunkData.writeDecodePage(new long[] {1L, partitionStart + 1}, new Object[] {1, 2}, 1); + chunkData.endChunk(); + + final Statistics statistics = chunkData.getChunks().get(0).getChunkStatistic(); + assertEquals(1, statistics.getCount()); + assertEquals(partitionStart + 1, statistics.getStartTime()); + } + + @Test + public void testAlignedChunkSelectsPointsFromCurrentTimePartition() throws Exception { + final long partitionStart = TimePartitionUtils.getTimePartitionInterval(); + final AlignedChunkData chunkData = + createAlignedTimeChunkData(new TTimePartitionSlot(partitionStart)); + chunkData.writeDecodePage(new long[] {1L, partitionStart + 1}, new Object[] {null, null}, 1); + chunkData.endChunk(); + + final Statistics statistics = chunkData.getChunks().get(0).getChunkStatistic(); + assertEquals(1, statistics.getCount()); + assertEquals(partitionStart + 1, statistics.getStartTime()); + } + private static Statistics createInt32Statistics() { final Statistics statistics = Statistics.getStatsByType(TSDataType.INT32); statistics.update(1L, 1); @@ -126,9 +153,14 @@ private static Statistics createInt32Statistics() { } private static NonAlignedChunkData createNonAlignedChunkData() { + return createNonAlignedChunkData(new TTimePartitionSlot(0L)); + } + + private static NonAlignedChunkData createNonAlignedChunkData( + final TTimePartitionSlot timePartitionSlot) { final IDeviceID device = new StringArrayDeviceID("root", "sg", "d1"); return (NonAlignedChunkData) - ChunkData.createChunkData(false, device, createChunkHeader(), new TTimePartitionSlot(0L)); + ChunkData.createChunkData(false, device, createChunkHeader(), timePartitionSlot); } private static AlignedChunkData createAlignedChunkData() { @@ -138,6 +170,11 @@ private static AlignedChunkData createAlignedChunkData() { } private static AlignedChunkData createAlignedTimeChunkData() { + return createAlignedTimeChunkData(new TTimePartitionSlot(0L)); + } + + private static AlignedChunkData createAlignedTimeChunkData( + final TTimePartitionSlot timePartitionSlot) { final IDeviceID device = new StringArrayDeviceID("root", "sg", "d1"); final ChunkHeader timeHeader = new ChunkHeader( @@ -149,7 +186,7 @@ private static AlignedChunkData createAlignedTimeChunkData() { 2, TIME_COLUMN_MASK); return (AlignedChunkData) - ChunkData.createChunkData(true, device, timeHeader, new TTimePartitionSlot(0L)); + ChunkData.createChunkData(true, device, timeHeader, timePartitionSlot); } private static ChunkHeader createChunkHeader() { From 62d9b45d8b30a614d90c996423cc5637d645d318 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Tue, 15 Sep 2026 14:43:25 +0800 Subject: [PATCH 4/8] Fix LOAD consensus WAL handling and add cluster IT --- .../iotdb/db/it/IoTDBLoadTsFileClusterIT.java | 138 ++++++ .../iotdb/db/i18n/StorageEngineMessages.java | 4 + .../iotdb/db/i18n/StorageEngineMessages.java | 4 + .../dataregion/DataExecutionVisitor.java | 35 +- .../node/DataNodePlanNodeDeserializer.java | 2 +- .../node/load/LoadTsFileConsensusNode.java | 33 +- .../load/TwoPhaseConsensusLoadStrategy.java | 42 +- .../storageengine/dataregion/DataRegion.java | 65 ++- .../dataregion/wal/buffer/WALEntry.java | 6 + .../dataregion/wal/buffer/WALInfoEntry.java | 6 + .../dataregion/wal/node/IWALNode.java | 3 + .../dataregion/wal/node/WALFakeNode.java | 6 + .../dataregion/wal/node/WALNode.java | 54 ++- .../storageengine/load/LoadTsFileManager.java | 450 ++++++++---------- 14 files changed, 530 insertions(+), 318 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java new file mode 100644 index 0000000000000..e2fdf1069c660 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java @@ -0,0 +1,138 @@ +/* + * 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.it; + +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.TsFileGenerator; +import org.apache.iotdb.itbase.category.ClusterIT; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.io.File; +import java.nio.file.Files; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.Arrays; +import java.util.List; + +@RunWith(IoTDBTestRunner.class) +@Category({ClusterIT.class}) +public class IoTDBLoadTsFileClusterIT { + + private static final long PARTITION_INTERVAL = 10_000L; + private static final String DATABASE = "root.load_cluster"; + private static final List DEVICES = + Arrays.asList( + DATABASE + ".d1", DATABASE + ".d2", DATABASE + ".d3", DATABASE + ".d4", DATABASE + ".d5"); + private static final List MEASUREMENTS = Arrays.asList("s1", "s2", "s3", "s4"); + private static final int POINT_COUNT_PER_DEVICE = 20_000; + private static final int REPLICATION_FACTOR = 3; + private static final int CONFIG_NODE_NUM = 3; + private static final int DATA_NODE_NUM = 3; + + private File tmpDir; + + @Before + public void setUp() throws Exception { + tmpDir = new File(Files.createTempDirectory("load-cluster-it").toUri()); + EnvFactory.getEnv().getConfig().getCommonConfig().setTimePartitionInterval(PARTITION_INTERVAL); + EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false); + EnvFactory.getEnv().getConfig().getCommonConfig().setPipeMemoryManagementEnabled(false); + EnvFactory.getEnv().getConfig().getCommonConfig().setDatanodeMemoryProportion("1:10:1:1:1:0"); + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setConfigNodeConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS) + .setSchemaReplicationFactor(REPLICATION_FACTOR) + .setDataReplicationFactor(REPLICATION_FACTOR); + EnvFactory.getEnv().initClusterEnvironment(CONFIG_NODE_NUM, DATA_NODE_NUM); + } + + @After + public void tearDown() throws Exception { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute("delete database " + DATABASE); + } catch (final Exception ignored) { + } + + EnvFactory.getEnv().cleanClusterEnvironment(); + + final File[] files = tmpDir.listFiles(); + if (files != null) { + for (final File file : files) { + Assert.assertTrue(file.delete()); + } + } + Assert.assertTrue(tmpDir.delete()); + } + + @Test + public void testLoadTsFileInCluster() throws Exception { + final long writtenPointCount; + try (final TsFileGenerator generator = + new TsFileGenerator(new File(tmpDir, "load-cluster-1-0-0-0.tsfile"))) { + final List schemas = + Arrays.asList( + new MeasurementSchema(MEASUREMENTS.get(0), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(1), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(2), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(3), TSDataType.INT64, TSEncoding.PLAIN)); + for (final String device : DEVICES) { + generator.registerTimeseries(device, schemas); + generator.generateData(device, POINT_COUNT_PER_DEVICE, 1L, false); + } + writtenPointCount = generator.getTotalNumber(); + } + + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute("create database " + DATABASE); + for (final String device : DEVICES) { + for (final String measurement : MEASUREMENTS) { + statement.execute( + "create timeseries " + device + "." + measurement + " " + TSDataType.INT64.name()); + } + } + statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); + + try (final ResultSet resultSet = + statement.executeQuery("select count(*) from " + DATABASE + ".**")) { + Assert.assertTrue(resultSet.next()); + final long count = resultSet.getLong(1); + Assert.assertEquals(writtenPointCount, count); + } + } + } +} diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java index fccc3a53dbd37..2e3bba6cd3e57 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -526,6 +526,10 @@ private StorageEngineMessages() {} public static final String CANNOT_CREATE_TSFILE_FOR_WRITING = "Can not create TsFile {} for writing."; public static final String CLOSE_TSFILE_IO_WRITER_ERROR = "Close TsFileIOWriter {} error."; public static final String CLOSE_MODIFICATION_FILE_ERROR = "Close ModificationFile {} error."; + public static final String LOG_PREPARING_LOAD_TSFILE_ARG_SEALING_STAGED_RESOURCES_1FDF1866 = + "Preparing LOAD TsFile {}: sealing staged resources."; + public static final String LOG_COMMITTING_LOAD_TSFILE_ARG_LOADING_PREPARED_RESOURCES_INTO_DATAREGION_EA1D6335 = + "Committing LOAD TsFile {}: loading prepared resources into DataRegion."; public static final String EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93 = "Table '%s.%s' does not exist when applying LOAD chunk data. It may have been dropped after the LOAD was analyzed."; public static final String TASK_DIR_NOT_EMPTY_SKIP_DELETE = "Task dir {} is not empty, skip deleting."; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java index f0de191a06ad9..44c3126c6d446 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -524,6 +524,10 @@ private StorageEngineMessages() {} public static final String CANNOT_CREATE_TSFILE_FOR_WRITING = "无法创建 TsFile {} 用于写入。"; public static final String CLOSE_TSFILE_IO_WRITER_ERROR = "关闭 TsFileIOWriter {} 出错。"; public static final String CLOSE_MODIFICATION_FILE_ERROR = "关闭修改文件 {} 出错。"; + public static final String LOG_PREPARING_LOAD_TSFILE_ARG_SEALING_STAGED_RESOURCES_1FDF1866 = + "正在准备 LOAD TsFile {}：封存暂存资源。"; + public static final String LOG_COMMITTING_LOAD_TSFILE_ARG_LOADING_PREPARED_RESOURCES_INTO_DATAREGION_EA1D6335 = + "正在提交 LOAD TsFile {}：将已准备资源加载到 DataRegion。"; public static final String EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93 = "应用 LOAD chunk 数据时表 '%s.%s' 不存在，可能在 LOAD 分析之后被删除了。"; public static final String TASK_DIR_NOT_EMPTY_SKIP_DELETE = "任务目录 {} 非空，跳过删除。"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java index 5ca2d565632ab..c143d048baa10 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java @@ -78,19 +78,28 @@ public TSStatus visitLoadTsFileConsensus( switch (node.getOp()) { case BEGIN: return StatusUtils.OK; - case PREPARE: - return StatusUtils.OK; case PIECE: try { - dataRegion.writeLoadTsFilePiece(node.getLoadId(), node.getTsFileDataList()); + dataRegion.writeLoadTsFilePiece(node); return StatusUtils.OK; } catch (final IOException | PageException e) { LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); return RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR, e.getMessage()); } - case ABORT: - dataRegion.rollbackLoadTsFile(node.getLoadId()); - return StatusUtils.OK; + case PREPARE: + try { + final Map progressIndexes = new HashMap<>(); + node.getTimePartition2ProgressIndex() + .forEach( + (slot, bytes) -> + progressIndexes.put( + slot, ProgressIndexType.deserializeFrom(ByteBuffer.wrap(bytes)))); + final boolean prepared = dataRegion.writeLoadTsFilePrepare(node, progressIndexes); + return prepared ? StatusUtils.OK : RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR); + } catch (final Exception e) { + LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); + return RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR, e.getMessage()); + } case COMMIT: try { final Map progressIndexes = new HashMap<>(); @@ -99,10 +108,16 @@ public TSStatus visitLoadTsFileConsensus( (slot, bytes) -> progressIndexes.put( slot, ProgressIndexType.deserializeFrom(ByteBuffer.wrap(bytes)))); - return dataRegion.commitLoadTsFile( - node.getLoadId(), node.isGeneratedByPipe(), progressIndexes) - ? StatusUtils.OK - : RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR); + final boolean committed = dataRegion.writeLoadTsFileCommit(node, progressIndexes); + return committed ? StatusUtils.OK : RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR); + } catch (final Exception e) { + LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); + return RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR, e.getMessage()); + } + case ABORT: + try { + dataRegion.writeLoadTsFileAbort(node); + return StatusUtils.OK; } catch (final Exception e) { LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); return RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR, e.getMessage()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java index 67c8b8c2ec8e0..93a916aee74dd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java @@ -206,7 +206,7 @@ public PlanNode deserializeFromWAL(ByteBuffer buffer) { case 2004: return ObjectNode.deserialize(buffer); case 2010: - return LoadTsFileConsensusNode.deserializeFromWAL(buffer); + return LoadTsFileConsensusNode.deserialize(buffer); default: throw new IllegalArgumentException(DataNodeQueryMessages.INVALID_NODE_TYPE + nodeType); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java index 3fc87e2c42e04..b90f67d1a3937 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java @@ -27,6 +27,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; +import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; @@ -45,8 +46,10 @@ import java.io.ByteArrayOutputStream; import java.io.DataInputStream; import java.io.DataOutputStream; +import java.io.File; import java.io.IOException; import java.io.InputStream; +import java.io.RandomAccessFile; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; @@ -240,6 +243,7 @@ public static LoadTsFileConsensusNode prepare( int pieceCount, long totalBytes, long checksum, + boolean isGeneratedByPipe, Map timePartition2ProgressIndex) { final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); node.op = LoadTsFileConsensusOp.PREPARE; @@ -248,6 +252,7 @@ public static LoadTsFileConsensusNode prepare( node.pieceCount = pieceCount; node.totalBytes = totalBytes; node.checksum = checksum; + node.isGeneratedByPipe = isGeneratedByPipe; node.timePartition2ProgressIndex = timePartition2ProgressIndex == null ? new HashMap<>() @@ -651,8 +656,32 @@ private byte[] readPieceContent(PieceRef ref) { if (ref.content != null) { return ref.content; } - throw new IllegalStateException( - DataNodeQueryMessages.EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2); + File file = new File(ref.relativePath); + if (!file.isAbsolute()) { + File[] bases = + java.util.Arrays.stream(IoTDBDescriptor.getInstance().getConfig().getLoadTsFileDirs()) + .map(File::new) + .map(base -> new File(base, ref.relativePath)) + .toArray(File[]::new); + for (File candidate : bases) { + if (candidate.isFile()) { + file = candidate; + break; + } + } + } + if (!file.isFile()) { + throw new IllegalStateException( + DataNodeQueryMessages.EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2); + } + try (RandomAccessFile input = new RandomAccessFile(file, "r")) { + byte[] content = new byte[Math.toIntExact(ref.size)]; + input.seek(ref.offset); + input.readFully(content); + return content; + } catch (IOException e) { + throw new IllegalStateException(e); + } } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java index 8d51783add037..fd12732957377 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java @@ -54,17 +54,17 @@ *

Phase 1 (split & stream). {@link #execute(LoadSingleTsFileNode)} resets the * per-file state, assigns a fresh uuid to {@link LoadTsFileDispatcherImpl} (executor naming / log * correlation) and feeds the source TsFile through {@link TsFileSplitter} into {@link - * TsFileSplitConsumer}. Every dispatched piece goes through {@code dispatchConsensusPiece}: the - * first piece of a region sends BEGIN with a fresh per-region load id, then PIECE commands with a - * monotonically increasing {@code pieceIndex}. Submission goes through {@link - * LoadConsensusSubmitter} with bounded retries for transient failures only. + * TsFileSplitConsumer}. Every dispatched piece goes through {@code dispatchConsensusPiece}: each + * piece is submitted as PIECE directly with a fresh per-region load id and a monotonically + * increasing {@code pieceIndex}. Submission goes through {@link LoadConsensusSubmitter} with + * bounded retries for transient failures only. * *

Phase 2 (commit or abort). If every region received all its pieces, each touched region * gets PREPARE (with the accumulated count/bytes) followed by COMMIT; otherwise every touched * region gets ABORT so the staged data is dropped. * - *

Per-file state tracks the touched regions, their load ids, BEGIN state, piece counters and - * pipe progress indexes. + *

Per-file state tracks the touched regions, their load ids, piece counters and pipe progress + * indexes. */ public class TwoPhaseConsensusLoadStrategy implements TsFileLoadStrategy { @@ -90,14 +90,13 @@ public class TwoPhaseConsensusLoadStrategy implements TsFileLoadStrategy { private final String userName; private final boolean isGeneratedByPipe; - /** The source file being loaded, kept for the phase-two commands and BEGIN metadata. */ + /** The source file being loaded, kept for the phase-two commands. */ private LoadSingleTsFileNode currentNode; /** Regions touched by the current file; used to send ABORT/PREPARE+COMMIT in phase two. */ private final Set allReplicaSets = new HashSet<>(); private final Map regionLoadIds = new ConcurrentHashMap<>(); - private final Set begunRegions = ConcurrentHashMap.newKeySet(); private final Map regionPieceCounts = new ConcurrentHashMap<>(); private final Map regionTotalBytes = new ConcurrentHashMap<>(); @@ -129,7 +128,6 @@ public boolean execute(LoadSingleTsFileNode node) { dispatcher.setUuid(UUID.randomUUID().toString()); allReplicaSets.clear(); regionLoadIds.clear(); - begunRegions.clear(); regionPieceCounts.clear(); regionTotalBytes.clear(); timePartitionSlotToProgressIndex.clear(); @@ -244,31 +242,6 @@ private boolean dispatchConsensusPiece( final String loadId = regionLoadIds.computeIfAbsent(regionId, ignored -> UUID.randomUUID().toString()); - if (begunRegions.add(regionId)) { - final LoadTsFileConsensusNode begin = - LoadTsFileConsensusNode.begin( - new PlanNodeId("load-begin-" + loadId), - loadId, - pieceNode.getTsFile() == null ? null : pieceNode.getTsFile().getName(), - currentNode.isTableModel(), - currentNode.getDatabase(), - // The total piece count is only known after phase one; PREPARE carries the real - // count, BEGIN keeps the "unknown" sentinel so the two never disagree. - -1); - final TSStatus beginStatus = submitConsensusWithRetry(replicaSet, begin); - if (beginStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { - LOGGER.warn( - DataNodeQueryMessages.DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG - + DataNodeQueryMessages - .RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG, - replicaSet, - TSStatusCode.representOf(beginStatus.getCode()).name(), - beginStatus.getMessage(), - pieceNode); - return false; - } - } - final long pieceIndex = pieceNode.getPieceIndex(); final LoadTsFileConsensusNode piece = LoadTsFileConsensusNode.piece( @@ -338,6 +311,7 @@ private boolean prepareAndCommitAllRegions(LoadSingleTsFileNode node) { Math.toIntExact(regionPieceCounts.getOrDefault(regionId, 0L)), regionTotalBytes.getOrDefault(regionId, 0L), 0L, + isGeneratedByPipe, timePartition2ProgressIndex); final TSStatus prepareStatus = consensusSubmitter.submit(replicaSet, prepare); if (prepareStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java index f7e9ca9594b5d..0ef1d76440d5d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java @@ -576,9 +576,68 @@ public void stopLoadTsFileManager() { } } - public void writeLoadTsFilePiece(final String loadId, final List tsFileDataList) - throws IOException, PageException { - getLoadTsFileManager().writePiece(loadId, tsFileDataList); + public List< + org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode + .PieceRef> + writeLoadTsFilePiece(final String loadId, final List tsFileDataList) + throws IOException, PageException { + return getLoadTsFileManager().writePiece(loadId, tsFileDataList); + } + + public List< + org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode + .PieceRef> + writeLoadTsFilePiece( + final org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode + node) + throws IOException, PageException { + return getLoadTsFileManager().writePiece(node); + } + + public boolean writeLoadTsFilePrepare( + final String loadId, + final int pieceCount, + final long totalBytes, + final boolean isGeneratedByPipe, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + return getLoadTsFileManager() + .prepare(loadId, pieceCount, totalBytes, isGeneratedByPipe, timePartitionProgressIndexMap); + } + + public boolean writeLoadTsFilePrepare( + final org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode + node, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + return getLoadTsFileManager().prepare(node, timePartitionProgressIndexMap); + } + + public boolean writeLoadTsFileCommit( + final String loadId, + final boolean isGeneratedByPipe, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + return getLoadTsFileManager().loadAll(loadId, isGeneratedByPipe, timePartitionProgressIndexMap); + } + + public boolean writeLoadTsFileCommit( + final org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode + node, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + return getLoadTsFileManager().loadAll(node, timePartitionProgressIndexMap); + } + + public void writeLoadTsFileAbort(final String loadId) { + getLoadTsFileManager().deleteAll(loadId); + } + + public void writeLoadTsFileAbort( + final org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode + node) + throws IOException { + getLoadTsFileManager().deleteAll(node); } public boolean commitLoadTsFile( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntry.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntry.java index 8c0a5e353291e..bb3a18fad6199 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntry.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntry.java @@ -22,6 +22,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -81,6 +82,8 @@ protected WALEntry(long memTableId, WALEntryValue value, boolean wait) { this.type = WALEntryType.RELATIONAL_DELETE_DATA_NODE; } else if (value instanceof ObjectNode) { this.type = WALEntryType.OBJECT_FILE_NODE; + } else if (value instanceof LoadTsFileConsensusNode) { + this.type = WALEntryType.LOAD_TSFILE_CONSENSUS_NODE; } else { throw new RuntimeException(StorageEngineMessages.UNKNOWN_WAL_ENTRY_TYPE); } @@ -141,6 +144,9 @@ public static WALEntry deserialize(DataInputStream stream) throws IOException { case OBJECT_FILE_NODE: value = (ObjectNode) PlanNodeType.deserializeFromWAL(stream); break; + case LOAD_TSFILE_CONSENSUS_NODE: + value = (LoadTsFileConsensusNode) PlanNodeType.deserializeFromWAL(stream); + break; default: throw new RuntimeException(StorageEngineMessages.UNKNOWN_WAL_ENTRY_TYPE_WITH_VALUE + type); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALInfoEntry.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALInfoEntry.java index 957cf03d89237..38a445bc8fcef 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALInfoEntry.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALInfoEntry.java @@ -22,6 +22,7 @@ import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -116,6 +117,9 @@ public void serialize(IWALByteBufferView buffer) { case OBJECT_FILE_NODE: ((ObjectNode) value).serializeToWAL(buffer, encodedSearchIndex); break; + case LOAD_TSFILE_CONSENSUS_NODE: + ((LoadTsFileConsensusNode) value).serializeToWAL(buffer, encodedSearchIndex); + break; case MEMORY_TABLE_SNAPSHOT: case CONTINUOUS_SAME_SEARCH_INDEX_SEPARATOR_NODE: value.serializeToWAL(buffer); @@ -210,6 +214,8 @@ public long getMemorySize() { return RamUsageEstimator.sizeOfObject(value); case OBJECT_FILE_NODE: return ((ObjectNode) value).serializedSize(); + case LOAD_TSFILE_CONSENSUS_NODE: + return ((LoadTsFileConsensusNode) value).serializedSize(); default: throw new RuntimeException(StorageEngineMessages.UNSUPPORTED_WAL_ENTRY_TYPE + type); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/IWALNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/IWALNode.java index a8fbbee0dc4f5..2a933dc1208d7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/IWALNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/IWALNode.java @@ -21,6 +21,7 @@ import org.apache.iotdb.consensus.common.DataSet; import org.apache.iotdb.consensus.iot.log.ConsensusReqReader; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -57,6 +58,8 @@ public interface IWALNode extends FlushListener, AutoCloseable, ConsensusReqRead WALFlushListener log(long memTableId, ObjectNode objectNode); + WALFlushListener log(long memTableId, LoadTsFileConsensusNode node); + /** Callback when memTable created. */ void onMemTableCreated(IMemTable memTable, String targetTsFile); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALFakeNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALFakeNode.java index d24d1a1d82348..1be4a98e5c465 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALFakeNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALFakeNode.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.storageengine.dataregion.wal.node; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -89,6 +90,11 @@ public WALFlushListener log(long memTableId, ObjectNode objectNode) { return getResult(); } + @Override + public WALFlushListener log(long memTableId, LoadTsFileConsensusNode node) { + return getResult(); + } + private WALFlushListener getResult() { switch (status) { case SUCCESS: diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALNode.java index 0e66509d35099..fa17a59975227 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALNode.java @@ -29,6 +29,7 @@ import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -289,6 +290,11 @@ public WALFlushListener log(long memTableId, ObjectNode objectNode) { return log(walEntry); } + @Override + public WALFlushListener log(long memTableId, LoadTsFileConsensusNode node) { + return log(new WALInfoEntry(memTableId, node)); + } + private WALFlushListener log(WALEntry walEntry) { buffer.write(walEntry); @@ -824,19 +830,23 @@ public boolean hasNext() { currentEntryNodeId.set(walByteBufReader.getCurrentEntryNodeId()); currentEntryContainsUserData.set( currentEntryContainsUserData.get() || containsUserData(type)); - if (type == WALEntryType.OBJECT_FILE_NODE) { + if (type == WALEntryType.OBJECT_FILE_NODE + || type == WALEntryType.LOAD_TSFILE_CONSENSUS_NODE) { WALEntry walEntry = WALEntry.deserialize( new DataInputStream(new ByteArrayInputStream(buffer.array()))); - // only be called by leader read from wal - // wal only has relativePath, offset, eof, length - // need to add WALEntryType + memtableId + relativePath, offset, eof, length + - // content - // need to add IoTConsensusRequest instead of ObjectNode - tmpNodes - .get() - .add(new IoTConsensusRequest(((ObjectNode) walEntry.getValue()).serialize())); - memorySize += ((ObjectNode) walEntry.getValue()).getMemorySize(); + // WAL only has relativePath, offset, eof and length. For downstream V1 sync, + // read the content from the referenced file and send an IoTConsensusRequest + // carrying the expanded bytes instead of the raw WAL slice. + ByteBuffer request = + type == WALEntryType.OBJECT_FILE_NODE + ? ((ObjectNode) walEntry.getValue()).serialize() + : ((LoadTsFileConsensusNode) walEntry.getValue()).serialize(); + tmpNodes.get().add(new IoTConsensusRequest(request)); + memorySize += + type == WALEntryType.OBJECT_FILE_NODE + ? ((ObjectNode) walEntry.getValue()).getMemorySize() + : request.remaining(); } else { tmpNodes.get().add(new IoTConsensusRequest(buffer)); memorySize += buffer.remaining(); @@ -861,19 +871,23 @@ public boolean hasNext() { currentEntryNodeId.set(walByteBufReader.getCurrentEntryNodeId()); currentEntryContainsUserData.set( currentEntryContainsUserData.get() || containsUserData(type)); - if (type == WALEntryType.OBJECT_FILE_NODE) { + if (type == WALEntryType.OBJECT_FILE_NODE + || type == WALEntryType.LOAD_TSFILE_CONSENSUS_NODE) { WALEntry walEntry = WALEntry.deserialize( new DataInputStream(new ByteArrayInputStream(buffer.array()))); - // only be called by leader read from wal - // wal only has relativePath, offset, eof, length - // need to add WALEntryType + memtableId + relativePath, offset, eof, length + - // content - // need to add IoTConsensusRequest instead of ObjectNode - tmpNodes - .get() - .add(new IoTConsensusRequest(((ObjectNode) walEntry.getValue()).serialize())); - memorySize += ((ObjectNode) walEntry.getValue()).getMemorySize(); + // WAL only has relativePath, offset, eof and length. For downstream V1 sync, + // read the content from the referenced file and send an IoTConsensusRequest + // carrying the expanded bytes instead of the raw WAL slice. + ByteBuffer request = + type == WALEntryType.OBJECT_FILE_NODE + ? ((ObjectNode) walEntry.getValue()).serialize() + : ((LoadTsFileConsensusNode) walEntry.getValue()).serialize(); + tmpNodes.get().add(new IoTConsensusRequest(request)); + memorySize += + type == WALEntryType.OBJECT_FILE_NODE + ? ((ObjectNode) walEntry.getValue()).getMemorySize() + : request.remaining(); } else { tmpNodes.get().add(new IoTConsensusRequest(buffer)); memorySize += buffer.remaining(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java index 381c870213494..fae8e6975f767 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java @@ -32,23 +32,23 @@ import org.apache.iotdb.commons.service.metric.MetricService; import org.apache.iotdb.commons.service.metric.enums.Metric; import org.apache.iotdb.commons.service.metric.enums.Tag; -import org.apache.iotdb.commons.utils.FileUtils; import org.apache.iotdb.commons.utils.RetryUtils; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.consensus.DataRegionConsensusImpl; import org.apache.iotdb.db.exception.load.LoadFileException; import org.apache.iotdb.db.i18n.StorageEngineMessages; -import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; import org.apache.iotdb.db.queryengine.plan.scheduler.load.LoadTsFileScheduler.LoadCommand; import org.apache.iotdb.db.storageengine.dataregion.DataRegion; import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; +import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor; import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; import org.apache.iotdb.db.storageengine.dataregion.modification.v1.ModificationFileV1; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus; -import org.apache.iotdb.db.storageengine.load.active.ActiveLoadAgent; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALFlushListener; import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; @@ -83,11 +83,9 @@ import java.util.Objects; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.PriorityBlockingQueue; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; -import java.util.stream.Stream; /** * {@link LoadTsFileManager} is used for dealing with {@link LoadTsFilePieceNode} and {@link @@ -118,195 +116,122 @@ public class LoadTsFileManager { private final Map uuid2WriterManager = new ConcurrentHashMap<>(); private final DataRegion dataRegion; - private final Map uuid2CleanupTask = new ConcurrentHashMap<>(); - private final PriorityBlockingQueue cleanupTaskQueue = new PriorityBlockingQueue<>(); - - private final ActiveLoadAgent activeLoadAgent = new ActiveLoadAgent(); - public LoadTsFileManager(final DataRegion dataRegion) { this.dataRegion = Objects.requireNonNull(dataRegion); - registerCleanupTaskExecutor(); - recover(); } - public void start() { - activeLoadAgent.start(); - } + public void start() {} public void stop() { - activeLoadAgent.stop(); - synchronized (uuid2CleanupTask) { - uuid2CleanupTask.values().forEach(CleanupTask::cancel); - uuid2CleanupTask.clear(); - cleanupTaskQueue.clear(); - } new HashSet<>(uuid2WriterManager.keySet()).forEach(this::forceCloseWriterManager); } - private long getCleanupTaskDelayInMs() { - return CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() * 1000L; + public void writeToDataRegion(LoadTsFilePieceNode pieceNode, String uuid) + throws IOException, PageException { + writePiece(uuid, pieceNode.getAllTsFileData()); } - private void createCleanupTaskIfAbsent(final String uuid) { - synchronized (uuid2CleanupTask) { - if (uuid2CleanupTask.containsKey(uuid)) { - return; - } - - final CleanupTask cleanupTask = new CleanupTask(uuid, getCleanupTaskDelayInMs()); - uuid2CleanupTask.put(uuid, cleanupTask); - cleanupTaskQueue.add(cleanupTask); - } - } + public List writePiece( + final String uuid, final List tsFileDataList) throws IOException, PageException { + final AtomicReference exception = new AtomicReference<>(); + final TsFileWriterManager writerManager = + uuid2WriterManager.computeIfAbsent( + uuid, + o -> { + try { + return getFolderManager() + .getNextWithRetry( + folder -> + new TsFileWriterManager( + new File(getDataRegionLoadDir(new File(folder)), uuid))); + } catch (DiskSpaceInsufficientException e) { + exception.set(e); + return null; + } + }); - private void rescheduleCleanupTask(final CleanupTask cleanupTask) { - synchronized (uuid2CleanupTask) { - if (uuid2CleanupTask.get(cleanupTask.uuid) != cleanupTask) { - return; + if (exception.get() != null || writerManager == null) { + throw new IOException( + String.format( + StorageEngineMessages + .STORAGE_EXCEPTION_FAILED_TO_CREATE_TSFILEWRITERMANAGER_FOR_UUID_S_BECAUSE_A0D68950, + uuid), + exception.get()); + } + + final Map previousLengths = snapshotFileLengths(writerManager.taskDir); + for (TsFileData tsFileData : tsFileDataList) { + switch (tsFileData.getType()) { + case CHUNK: + ChunkData chunkData = (ChunkData) tsFileData; + writerManager.write( + new DataPartitionInfo(dataRegion, chunkData.getTimePartitionSlot()), chunkData); + break; + case DELETION: + writerManager.writeDeletion(dataRegion, (DeletionData) tsFileData); + break; + default: + throw new IOException( + StorageEngineMessages.UNSUPPORTED_TSFILE_DATA_TYPE + tsFileData.getType()); } - - cleanupTaskQueue.remove(cleanupTask); - cleanupTask.resetScheduledTime(); - cleanupTaskQueue.add(cleanupTask); } + writerManager.flush(); + return createPieceRefs(writerManager.taskDir, previousLengths); } - private void registerCleanupTaskExecutor() { - PipeDataNodeAgent.runtime() - .registerPeriodicalJob( - "LoadTsFileManager#cleanupTasks", - this::cleanupTasks, - CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() >> 2); + public List writePiece(final LoadTsFileConsensusNode node) + throws IOException, PageException { + final List refs = + writePiece(node.getLoadId(), node.getTsFileDataList()); + if (node.getPieceRefs().isEmpty()) { + final LoadTsFileConsensusNode walNode = + LoadTsFileConsensusNode.pieceRefs( + node.getPlanNodeId(), + node.getLoadId(), + node.getTsFileId(), + node.getPieceIndex(), + refs, + node.getChecksum(), + node.getDataSize(), + null); + walNode.setSearchIndex(node.getSearchIndex()); + logLoadNodeToWAL(walNode); + dataRegion.insertSeparatorToWAL(node); + } + return refs; } - private void cleanupTasks() { - while (!cleanupTaskQueue.isEmpty()) { - synchronized (uuid2CleanupTask) { - if (cleanupTaskQueue.isEmpty()) { - continue; - } - - final CleanupTask cleanupTask = cleanupTaskQueue.peek(); - if (cleanupTask.scheduledTime <= System.currentTimeMillis()) { - if (cleanupTask.isLoadTaskRunning) { - cleanupTaskQueue.poll(); - cleanupTask.resetScheduledTime(); - cleanupTaskQueue.add(cleanupTask); - continue; - } - - cleanupTask.run(); - - uuid2CleanupTask.remove(cleanupTask.uuid); - cleanupTaskQueue.poll(); - } else { - final long waitTimeInMs = cleanupTask.scheduledTime - System.currentTimeMillis(); - LOGGER.info( - StorageEngineMessages - .STORAGE_LOG_NEXT_LOAD_CLEANUP_TASK_IS_NOT_READY_TO_RUN_WAIT_FOR_AT_LEAST_CBE0023F, - cleanupTask.uuid, - waitTimeInMs, - waitTimeInMs / 1000.0); - return; + private static Map snapshotFileLengths(final File taskDir) { + final Map result = new HashMap<>(); + final File[] files = taskDir.listFiles(); + if (files != null) { + for (File file : files) { + if (file.isFile()) { + result.put(file, file.length()); } } } + return result; } - private void recover() { - if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) { - synchronized (FOLDER_MANAGER) { - if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) { - LOAD_BASE_DIRS.set(CONFIG.getLoadTsFileDirs()); + private static List createPieceRefs( + final File taskDir, final Map previousLengths) { + final List refs = new java.util.ArrayList<>(); + final File[] files = taskDir.listFiles(); + if (files != null) { + for (File file : files) { + if (!file.isFile()) { + continue; } - } - } - - final File[] baseDirs = - Arrays.stream(LOAD_BASE_DIRS.get()) - .map(File::new) - .map(this::getDataRegionLoadDir) - .toArray(File[]::new); - final File[] files = - Arrays.stream(baseDirs) - .filter(File::exists) - .flatMap( - dir -> { - final File[] listedFiles = dir.listFiles(); - return listedFiles != null ? Arrays.stream(listedFiles) : Stream.empty(); - }) - .toArray(File[]::new); - - Arrays.stream(files) - .parallel() - .forEach( - taskDir -> { - final String uuid = taskDir.getName(); - final TsFileWriterManager writerManager = new TsFileWriterManager(taskDir); - - uuid2WriterManager.put(uuid, writerManager); - writerManager.close(); - - createCleanupTaskIfAbsent(uuid); - }); - } - - public void writeToDataRegion(LoadTsFilePieceNode pieceNode, String uuid) - throws IOException, PageException { - writePiece(uuid, pieceNode.getAllTsFileData()); - } - - public void writePiece(final String uuid, final List tsFileDataList) - throws IOException, PageException { - createCleanupTaskIfAbsent(uuid); - - final Optional cleanupTask = Optional.ofNullable(uuid2CleanupTask.get(uuid)); - cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning); - try { - final AtomicReference exception = new AtomicReference<>(); - final TsFileWriterManager writerManager = - uuid2WriterManager.computeIfAbsent( - uuid, - o -> { - try { - return getFolderManager() - .getNextWithRetry( - folder -> - new TsFileWriterManager( - new File(getDataRegionLoadDir(new File(folder)), uuid))); - } catch (DiskSpaceInsufficientException e) { - exception.set(e); - return null; - } - }); - - if (exception.get() != null || writerManager == null) { - throw new IOException( - String.format( - StorageEngineMessages - .STORAGE_EXCEPTION_FAILED_TO_CREATE_TSFILEWRITERMANAGER_FOR_UUID_S_BECAUSE_A0D68950, - uuid), - exception.get()); - } - - for (TsFileData tsFileData : tsFileDataList) { - switch (tsFileData.getType()) { - case CHUNK: - ChunkData chunkData = (ChunkData) tsFileData; - writerManager.write( - new DataPartitionInfo(dataRegion, chunkData.getTimePartitionSlot()), chunkData); - break; - case DELETION: - writerManager.writeDeletion(dataRegion, (DeletionData) tsFileData); - break; - default: - throw new IOException( - StorageEngineMessages.UNSUPPORTED_TSFILE_DATA_TYPE + tsFileData.getType()); + final long previousLength = previousLengths.getOrDefault(file, 0L); + final long size = file.length() - previousLength; + if (size > 0) { + refs.add( + new LoadTsFileConsensusNode.PieceRef(file.getAbsolutePath(), previousLength, size)); } } - } finally { - cleanupTask.ifPresent(CleanupTask::markLoadTaskNotRunning); } + return refs; } private FolderManager getFolderManager() throws DiskSpaceInsufficientException { @@ -344,6 +269,41 @@ private File getDataRegionLoadDir(final File baseDir) { + dataRegion.getDataRegionIdString()); } + public boolean prepare( + final String uuid, + final int pieceCount, + final long totalBytes, + final boolean isGeneratedByPipe, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + if (!uuid2WriterManager.containsKey(uuid)) { + return false; + } + LOGGER.info( + StorageEngineMessages.LOG_PREPARING_LOAD_TSFILE_ARG_SEALING_STAGED_RESOURCES_1FDF1866, + uuid); + uuid2WriterManager.get(uuid).prepare(isGeneratedByPipe, timePartitionProgressIndexMap); + return true; + } + + public boolean prepare( + final LoadTsFileConsensusNode node, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + if (!uuid2WriterManager.containsKey(node.getLoadId())) { + return false; + } + LOGGER.info( + StorageEngineMessages.LOG_PREPARING_LOAD_TSFILE_ARG_SEALING_STAGED_RESOURCES_1FDF1866, + node.getLoadId()); + uuid2WriterManager + .get(node.getLoadId()) + .prepare(node.isGeneratedByPipe(), timePartitionProgressIndexMap); + logLoadNodeToWAL(node); + dataRegion.insertSeparatorToWAL(node); + return true; + } + public boolean loadAll( String uuid, boolean isGeneratedByPipe, @@ -353,17 +313,30 @@ public boolean loadAll( return false; } - createCleanupTaskIfAbsent(uuid); + LOGGER.info( + StorageEngineMessages + .LOG_COMMITTING_LOAD_TSFILE_ARG_LOADING_PREPARED_RESOURCES_INTO_DATAREGION_EA1D6335, + uuid); + uuid2WriterManager.get(uuid).loadAll(isGeneratedByPipe); + forceCloseWriterManager(uuid); + return true; + } - final Optional cleanupTask = Optional.ofNullable(uuid2CleanupTask.get(uuid)); - cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning); - try { - uuid2WriterManager.get(uuid).loadAll(isGeneratedByPipe, timePartitionProgressIndexMap); - } finally { - cleanupTask.ifPresent(CleanupTask::markLoadTaskNotRunning); + public boolean loadAll( + final LoadTsFileConsensusNode node, + final Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + if (!uuid2WriterManager.containsKey(node.getLoadId())) { + return false; } - - clean(uuid); + LOGGER.info( + StorageEngineMessages + .LOG_COMMITTING_LOAD_TSFILE_ARG_LOADING_PREPARED_RESOURCES_INTO_DATAREGION_EA1D6335, + node.getLoadId()); + uuid2WriterManager.get(node.getLoadId()).loadAll(node.isGeneratedByPipe()); + forceCloseWriterManager(node.getLoadId()); + logLoadNodeToWAL(node); + dataRegion.insertSeparatorToWAL(node); return true; } @@ -371,20 +344,36 @@ public boolean deleteAll(String uuid) { if (!uuid2WriterManager.containsKey(uuid)) { return false; } - clean(uuid); + forceCloseWriterManager(uuid); return true; } - private void clean(String uuid) { - synchronized (uuid2CleanupTask) { - final CleanupTask cleanupTask = uuid2CleanupTask.remove(uuid); - if (cleanupTask != null) { - cleanupTask.cancel(); - cleanupTaskQueue.remove(cleanupTask); - } + public boolean deleteAll(final LoadTsFileConsensusNode node) throws IOException { + if (!uuid2WriterManager.containsKey(node.getLoadId())) { + return false; } + forceCloseWriterManager(node.getLoadId()); + logLoadNodeToWAL(node); + dataRegion.insertSeparatorToWAL(node); + return true; + } - forceCloseWriterManager(uuid); + private void logLoadNodeToWAL(final LoadTsFileConsensusNode node) throws IOException { + final AtomicReference failure = new AtomicReference<>(); + dataRegion + .getWALNode() + .ifPresent( + wal -> { + final WALFlushListener listener = wal.log(TsFileProcessor.MEMTABLE_NOT_EXIST, node); + if (listener.waitForResult() == WALFlushListener.Status.FAILURE) { + final Exception cause = listener.getCause(); + failure.set( + cause instanceof IOException ? (IOException) cause : new IOException(cause)); + } + }); + if (failure.get() != null) { + throw failure.get(); + } } private void forceCloseWriterManager(String uuid) { @@ -462,6 +451,7 @@ private static class TsFileWriterManager { private Map dataPartition2Resource; private Map dataPartition2ModificationFile; private boolean isClosed; + private boolean isPrepared; private TsFileWriterManager(File taskDir) { this.taskDir = taskDir; @@ -469,15 +459,13 @@ private TsFileWriterManager(File taskDir) { this.dataPartition2Resource = new HashMap<>(); this.dataPartition2ModificationFile = new HashMap<>(); this.isClosed = false; + this.isPrepared = false; - clearDir(taskDir); + ensureDir(taskDir); } - private void clearDir(File dir) { - if (dir.exists()) { - FileUtils.deleteFileOrDirectoryWithRetry(dir); - } - if (dir.mkdirs()) { + private void ensureDir(File dir) { + if (!dir.exists() && dir.mkdirs()) { LOGGER.info(StorageEngineMessages.LOAD_TSFILE_DIR_CREATED, dir.getPath()); } } @@ -560,13 +548,22 @@ private void writeDeletion(DataRegion dataRegion, DeletionData deletionData) } } - private void loadAll( + private void flush() throws IOException { + for (TsFilePrecalculatedChunkWriter writer : dataPartition2Writer.values()) { + writer.getOutput().flush(); + } + } + + private void prepare( boolean isGeneratedByPipe, Map timePartitionProgressIndexMap) throws IOException, LoadFileException { if (isClosed) { throw new IOException(String.format(MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED, taskDir)); } + if (isPrepared) { + return; + } for (final Map.Entry entry : dataPartition2ModificationFile.entrySet()) { entry.getValue().close(); @@ -576,7 +573,6 @@ private void loadAll( final TsFilePrecalculatedChunkWriter writer = entry.getValue(); writer.close(); - final DataRegion dataRegion = entry.getKey().getDataRegion(); final TsFileResource tsFileResource = dataPartition2Resource.get(entry.getKey()); tsFileResource.setGeneratedByPipe(isGeneratedByPipe); endTsFileResource( @@ -584,6 +580,15 @@ private void loadAll( tsFileResource, timePartitionProgressIndexMap.getOrDefault( entry.getKey().getTimePartitionSlot(), MinimumProgressIndex.INSTANCE)); + } + isPrepared = true; + } + + private void loadAll(boolean isGeneratedByPipe) throws LoadFileException { + for (final Map.Entry entry : + dataPartition2Resource.entrySet()) { + final DataRegion dataRegion = entry.getKey().getDataRegion(); + final TsFileResource tsFileResource = entry.getValue(); dataRegion.loadNewTsFile(tsFileResource, true, isGeneratedByPipe, false, Optional.empty()); // Metrics @@ -592,7 +597,10 @@ private void loadAll( .ifPresent( databaseName -> updateWritePointCountMetrics( - dataRegion, databaseName, getTsFileWritePointCount(writer), false)); + dataRegion, + databaseName, + getTsFileWritePointCount(dataPartition2Writer.get(entry.getKey())), + false)); } } @@ -755,60 +763,6 @@ private void close() { } } - private class CleanupTask implements Runnable, Comparable { - - private final String uuid; - - private final long delayInMs; - private long scheduledTime; - - private volatile boolean isLoadTaskRunning = false; - private volatile boolean isCanceled = false; - - private CleanupTask(String uuid, long delayInMs) { - this.uuid = uuid; - this.delayInMs = delayInMs; - resetScheduledTime(); - } - - public void markLoadTaskRunning() { - isLoadTaskRunning = true; - rescheduleCleanupTask(this); - } - - public void markLoadTaskNotRunning() { - isLoadTaskRunning = false; - rescheduleCleanupTask(this); - } - - public void resetScheduledTime() { - scheduledTime = System.currentTimeMillis() + delayInMs; - } - - public void cancel() { - isCanceled = true; - } - - @Override - public void run() { - if (isCanceled) { - LOGGER.info(StorageEngineMessages.LOAD_CLEANUP_TASK_CANCELED, uuid); - } else { - LOGGER.info(StorageEngineMessages.LOAD_CLEANUP_TASK_STARTS, uuid); - try { - forceCloseWriterManager(uuid); - } catch (Exception e) { - LOGGER.warn(StorageEngineMessages.LOAD_CLEANUP_TASK_ERROR, uuid, e); - } - } - } - - @Override - public int compareTo(CleanupTask that) { - return Long.compare(this.scheduledTime, that.scheduledTime); - } - } - private static class DataPartitionInfo { private final DataRegion dataRegion; From bbaea5d4652efe906cc0c76bf69a745f6110d9b2 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Tue, 15 Sep 2026 17:08:33 +0800 Subject: [PATCH 5/8] Handle LOAD TsFile physical holes during V1 sync --- .../scheduler/load/ChunkOffsetCalculator.java | 16 +- .../storageengine/load/LoadTsFileManager.java | 52 ++- .../load/LoadTsFileProgress.java | 404 ++++++++++++++++++ .../write/TsFilePrecalculatedChunkWriter.java | 53 ++- 4 files changed, 496 insertions(+), 29 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkOffsetCalculator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkOffsetCalculator.java index 3e3c9b7649022..59c9eea76ee40 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkOffsetCalculator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/ChunkOffsetCalculator.java @@ -15,7 +15,9 @@ import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.file.header.ChunkGroupHeader; +import org.apache.tsfile.file.header.ChunkHeader; import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.utils.BytesUtils; @@ -54,7 +56,7 @@ public ChunkLayout allocate(final IDeviceID device, final boolean aligned, final final long chunkOffset = nextOffset; final long chunkLength = - chunk.getHeader().getSerializedSize() + (long) chunk.getData().remaining(); + getChunkHeaderSize(chunk.getHeader()) + (long) chunk.getData().remaining(); final boolean firstChunkOfGroup = nextChunkIndex == 0; final ChunkLayout layout = new ChunkLayout( @@ -115,8 +117,18 @@ public ChunkData.ChunkLayout assign(final ChunkData chunkData) { } private int getChunkGroupHeaderSize(final IDeviceID device) { + final IDeviceID normalizedDevice = + device instanceof StringArrayDeviceID ? device : new StringArrayDeviceID(device.toString()); try (ByteArrayOutputStream output = new ByteArrayOutputStream()) { - return new ChunkGroupHeader(device).serializeTo(output); + return new ChunkGroupHeader(normalizedDevice).serializeTo(output); + } catch (final IOException e) { + throw new IllegalStateException(e); + } + } + + private int getChunkHeaderSize(final ChunkHeader chunkHeader) { + try (ByteArrayOutputStream output = new ByteArrayOutputStream()) { + return chunkHeader.serializeTo(output); } catch (final IOException e) { throw new IllegalStateException(e); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java index fae8e6975f767..40671567fb7c0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java @@ -59,6 +59,7 @@ import org.apache.tsfile.common.constant.TsFileConstant; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.PageException; +import org.apache.tsfile.file.header.ChunkHeader; import org.apache.tsfile.file.metadata.IChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.read.TimeValuePair; @@ -449,6 +450,7 @@ private static class TsFileWriterManager { private final File taskDir; private Map dataPartition2Writer; private Map dataPartition2Resource; + private Map dataPartition2Progress; private Map dataPartition2ModificationFile; private boolean isClosed; private boolean isPrepared; @@ -457,6 +459,7 @@ private TsFileWriterManager(File taskDir) { this.taskDir = taskDir; this.dataPartition2Writer = new HashMap<>(); this.dataPartition2Resource = new HashMap<>(); + this.dataPartition2Progress = new HashMap<>(); this.dataPartition2ModificationFile = new HashMap<>(); this.isClosed = false; this.isPrepared = false; @@ -492,8 +495,10 @@ private void write(DataPartitionInfo partitionInfo, ChunkData chunkData) final TsFilePrecalculatedChunkWriter writer = new TsFilePrecalculatedChunkWriter(newTsFile); final TsFileResource resource = new TsFileResource(newTsFile); + final LoadTsFileProgress progress = new LoadTsFileProgress(newTsFile); dataPartition2Writer.put(partitionInfo, writer); dataPartition2Resource.put(partitionInfo, resource); + dataPartition2Progress.put(partitionInfo, progress); } final TsFilePrecalculatedChunkWriter writer = dataPartition2Writer.get(partitionInfo); final ChunkData.ChunkLayout layout = chunkData.getChunkLayout(); @@ -504,6 +509,8 @@ private void write(DataPartitionInfo partitionInfo, ChunkData chunkData) final List chunks = chunkData.getChunks(); for (int i = 0; i < chunks.size(); i++) { final Chunk chunk = chunks.get(i); + final long chunkLength = + getChunkHeaderSerializedSize(chunk.getHeader()) + (long) chunk.getData().remaining(); writer.writeChunk( chunkData.getDevice(), chunkData.isAligned(), @@ -511,10 +518,17 @@ private void write(DataPartitionInfo partitionInfo, ChunkData chunkData) layout.firstChunkOfGroup() && i == 0, chunk, chunkOffset); - chunkOffset += chunk.getHeader().getSerializedSize() + (long) chunk.getData().remaining(); - } - if (chunkOffset != layout.offset() + layout.length()) { - throw new IOException("Chunk layout length does not match encoded chunks"); + chunkOffset += chunkLength; + dataPartition2Progress + .get(partitionInfo) + .recordChunk( + chunkData.getDevice().toString(), + chunkData.isAligned(), + layout.chunkGroupHeaderOffset(), + chunkOffset - chunkLength, + layout.firstChunkOfGroup() && i == 0, + chunk, + chunkOffset); } } @@ -548,6 +562,14 @@ private void writeDeletion(DataRegion dataRegion, DeletionData deletionData) } } + private int getChunkHeaderSerializedSize(final ChunkHeader chunkHeader) { + try (java.io.ByteArrayOutputStream output = new java.io.ByteArrayOutputStream()) { + return chunkHeader.serializeTo(output); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + private void flush() throws IOException { for (TsFilePrecalculatedChunkWriter writer : dataPartition2Writer.values()) { writer.getOutput().flush(); @@ -571,6 +593,11 @@ private void prepare( for (final Map.Entry entry : dataPartition2Writer.entrySet()) { final TsFilePrecalculatedChunkWriter writer = entry.getValue(); + final LoadTsFileProgress progress = dataPartition2Progress.get(entry.getKey()); + if (progress != null && progress.exists() && !progress.isReady(writer.getFile().length())) { + throw new LoadFileException( + "Staged LOAD TsFile is not contiguous yet: " + writer.getFile().getAbsolutePath()); + } writer.close(); final TsFileResource tsFileResource = dataPartition2Resource.get(entry.getKey()); @@ -745,6 +772,22 @@ private void close() { } } } + if (dataPartition2Progress != null) { + for (final LoadTsFileProgress progress : dataPartition2Progress.values()) { + try { + final Path progressPath = progress.getProgressFile().toPath(); + if (Files.exists(progressPath)) { + RetryUtils.retryOnException( + () -> { + Files.delete(progressPath); + return null; + }); + } + } catch (IOException e) { + LOGGER.warn(MESSAGE_DELETE_FAIL, progress.getProgressFile().getAbsolutePath(), e); + } + } + } try { RetryUtils.retryOnException( () -> { @@ -758,6 +801,7 @@ private void close() { } dataPartition2Writer = null; dataPartition2Resource = null; + dataPartition2Progress = null; dataPartition2ModificationFile = null; isClosed = true; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java new file mode 100644 index 0000000000000..6ef3e042897f9 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java @@ -0,0 +1,404 @@ +/* + * 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.storageengine.load; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.apache.tsfile.file.metadata.statistics.Statistics; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.utils.ReadWriteIOUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.File; +import java.io.FileInputStream; +import java.io.FileOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; + +/** + * Persistent physical-range tracker for a staged TsFile whose data zone may be written with holes + * during IoTConsensus V1 leader handoff. + * + *

The file is append-only. Every successfully placed physical chunk is recorded with enough + * information to rebuild the {@link org.apache.tsfile.write.TsFilePrecalculatedChunkWriter} + * metadata later: device, aligned flag, chunk-group-header offset, chunk offset, first-chunk + * marker, serialized chunk header, serialized chunk statistics and the exact data range in the + * staged file. + */ +public class LoadTsFileProgress { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadTsFileProgress.class); + + public static final String PROGRESS_SUFFIX = ".progress"; + private static final String MAGIC = "LTPROG1"; + private static final byte PROGRESS_ENTRY_VERSION = 1; + + /** TsFile header size written by {@code TsFilePrecalculatedChunkWriter}. */ + private static final long TS_FILE_HEADER_SIZE = 7L; + + private final File progressFile; + private long totalLength = -1L; + private final List records = new ArrayList<>(); + + public LoadTsFileProgress(final File tsFile) { + this.progressFile = new File(tsFile.getAbsolutePath() + PROGRESS_SUFFIX); + } + + public static File progressFileFor(final File tsFile) { + return new File(tsFile.getAbsolutePath() + PROGRESS_SUFFIX); + } + + public boolean exists() { + return progressFile.isFile(); + } + + public void activate(final long totalLength) throws IOException { + this.totalLength = totalLength; + if (!progressFile.exists()) { + try (final DataOutputStream output = + new DataOutputStream(new FileOutputStream(progressFile, false))) { + ReadWriteIOUtils.write(MAGIC, output); + ReadWriteIOUtils.write(totalLength, output); + } + } + } + + public void recordChunk( + final String device, + final boolean aligned, + final long chunkGroupHeaderOffset, + final long chunkOffset, + final boolean firstChunkOfGroup, + final Chunk chunk, + final long dataEnd) + throws IOException { + activate(dataEnd); + + final byte[] chunkHeaderBytes = serializeChunkHeader(chunk); + final byte[] statisticsBytes = serializeStatistics(chunk); + final long physicalStart = firstChunkOfGroup ? chunkGroupHeaderOffset : chunkOffset; + + final ChunkRangeRecord record = + new ChunkRangeRecord( + device, + aligned, + chunkGroupHeaderOffset, + chunkOffset, + firstChunkOfGroup, + chunk.getHeader().getDataType(), + chunk.getHeader().getChunkType(), + chunkHeaderBytes, + statisticsBytes, + physicalStart, + dataEnd); + records.add(record); + + try (final DataOutputStream output = + new DataOutputStream(new FileOutputStream(progressFile, true))) { + writeEntry(record, output); + } + } + + public boolean isReady(final long currentFileLength) throws IOException { + if (currentFileLength <= TS_FILE_HEADER_SIZE) { + return false; + } + + final List sorted = readAllRecords(); + if (sorted.isEmpty()) { + return false; + } + sorted.sort(Comparator.comparingLong(ChunkRangeRecord::physicalStart)); + + long coveredStart = sorted.get(0).physicalStart(); + long coveredEnd = sorted.get(0).physicalEnd(); + if (coveredStart != TS_FILE_HEADER_SIZE) { + return false; + } + for (int i = 1; i < sorted.size(); i++) { + final ChunkRangeRecord record = sorted.get(i); + if (record.physicalStart() > coveredEnd) { + return false; + } + coveredEnd = Math.max(coveredEnd, record.physicalEnd()); + } + return coveredEnd == currentFileLength; + } + + public List readAllRecords() throws IOException { + if (!progressFile.isFile()) { + return new ArrayList<>(); + } + + final List result = new ArrayList<>(); + try (final DataInputStream input = new DataInputStream(new FileInputStream(progressFile))) { + final String magic = ReadWriteIOUtils.readString(input); + if (!MAGIC.equals(magic)) { + throw new IOException("Invalid load progress file: " + progressFile); + } + final long persistedTotalLength = ReadWriteIOUtils.readLong(input); + this.totalLength = persistedTotalLength; + while (input.available() > 0) { + final int entryLength = ReadWriteIOUtils.readInt(input); + if (entryLength <= 0) { + throw new IOException("Invalid progress entry length: " + entryLength); + } + final byte version = input.readByte(); + final byte[] payload = new byte[entryLength - 1]; + input.readFully(payload); + if (version == PROGRESS_ENTRY_VERSION) { + try (final DataInputStream payloadInput = + new DataInputStream(new ByteArrayInputStream(payload))) { + result.add(ChunkRangeRecord.deserializePayload(payloadInput)); + } + } else { + LOGGER.warn( + "Skipping unsupported load progress entry version {} in {}", version, progressFile); + } + } + } + records.clear(); + records.addAll(result); + return result; + } + + public File getProgressFile() { + return progressFile; + } + + public long getTotalLength() { + return totalLength; + } + + private byte[] serializeChunkHeader(final Chunk chunk) throws IOException { + try (final ByteArrayOutputStream output = new ByteArrayOutputStream()) { + chunk.getHeader().serializeTo(output); + return output.toByteArray(); + } + } + + private byte[] serializeStatistics(final Chunk chunk) throws IOException { + try (final ByteArrayOutputStream output = new ByteArrayOutputStream()) { + chunk.getChunkStatistic().serialize(output); + return output.toByteArray(); + } + } + + private static void writeEntry(final ChunkRangeRecord record, final DataOutputStream output) + throws IOException { + final ByteArrayOutputStream payloadOutput = new ByteArrayOutputStream(); + try (final DataOutputStream payloadStream = new DataOutputStream(payloadOutput)) { + record.serializePayload(payloadStream); + } + final byte[] payload = payloadOutput.toByteArray(); + ReadWriteIOUtils.write(payload.length + 1, output); + output.writeByte(PROGRESS_ENTRY_VERSION); + output.write(payload); + } + + public static ChunkHeader deserializeChunkHeader( + final byte chunkType, final byte[] chunkHeaderBytes) throws IOException { + try (final ByteArrayInputStream input = new ByteArrayInputStream(chunkHeaderBytes)) { + return ChunkHeader.deserializeFrom(input, chunkType); + } + } + + public static Statistics deserializeStatistics( + final TSDataType dataType, final byte[] statisticsBytes) throws IOException { + try (final ByteArrayInputStream input = new ByteArrayInputStream(statisticsBytes)) { + return Statistics.deserialize(input, dataType); + } + } + + public static Chunk deserializeChunk(final ChunkRangeRecord record, final byte[] chunkDataBytes) + throws IOException { + final ChunkHeader chunkHeader = + deserializeChunkHeader(record.chunkType(), record.chunkHeaderBytes()); + final Statistics statistics = + deserializeStatistics(record.dataType(), record.statisticsBytes()); + return new Chunk(chunkHeader, ByteBuffer.wrap(chunkDataBytes), null, statistics); + } + + /** A physical interval plus the metadata required to recover the writer entry. */ + public static final class ChunkRangeRecord { + private final String device; + private final boolean aligned; + private final long chunkGroupHeaderOffset; + private final long chunkOffset; + private final boolean firstChunkOfGroup; + private final TSDataType dataType; + private final byte chunkType; + private final byte[] chunkHeaderBytes; + private final byte[] statisticsBytes; + private final long physicalStart; + private final long physicalEnd; + + public ChunkRangeRecord( + final String device, + final boolean aligned, + final long chunkGroupHeaderOffset, + final long chunkOffset, + final boolean firstChunkOfGroup, + final TSDataType dataType, + final byte chunkType, + final byte[] chunkHeaderBytes, + final byte[] statisticsBytes, + final long physicalStart, + final long physicalEnd) { + this.device = device; + this.aligned = aligned; + this.chunkGroupHeaderOffset = chunkGroupHeaderOffset; + this.chunkOffset = chunkOffset; + this.firstChunkOfGroup = firstChunkOfGroup; + this.dataType = dataType; + this.chunkType = chunkType; + this.chunkHeaderBytes = chunkHeaderBytes; + this.statisticsBytes = statisticsBytes; + this.physicalStart = physicalStart; + this.physicalEnd = physicalEnd; + } + + public String device() { + return device; + } + + public StringArrayDeviceID deviceId() { + return new StringArrayDeviceID(device); + } + + public boolean aligned() { + return aligned; + } + + public long chunkGroupHeaderOffset() { + return chunkGroupHeaderOffset; + } + + public long chunkOffset() { + return chunkOffset; + } + + public boolean firstChunkOfGroup() { + return firstChunkOfGroup; + } + + public TSDataType dataType() { + return dataType; + } + + public byte chunkType() { + return chunkType; + } + + public byte[] chunkHeaderBytes() { + return chunkHeaderBytes; + } + + public byte[] statisticsBytes() { + return statisticsBytes; + } + + public long physicalStart() { + return physicalStart; + } + + public long physicalEnd() { + return physicalEnd; + } + + public long dataLength() { + return physicalEnd - chunkOffset - deserializeHeaderSizeSafe(); + } + + private int deserializeHeaderSizeSafe() { + try { + return deserializeChunkHeader(chunkType, chunkHeaderBytes).getSerializedSize(); + } catch (final IOException e) { + throw new IllegalStateException(e); + } + } + + public void serializePayload(final DataOutputStream output) throws IOException { + ReadWriteIOUtils.write(device, output); + ReadWriteIOUtils.write(aligned, output); + ReadWriteIOUtils.write(chunkGroupHeaderOffset, output); + ReadWriteIOUtils.write(chunkOffset, output); + ReadWriteIOUtils.write(firstChunkOfGroup, output); + ReadWriteIOUtils.write(dataType.ordinal(), output); + output.writeByte(chunkType); + writeBytes(chunkHeaderBytes, output); + writeBytes(statisticsBytes, output); + ReadWriteIOUtils.write(physicalStart, output); + ReadWriteIOUtils.write(physicalEnd, output); + } + + public static ChunkRangeRecord deserializePayload(final DataInputStream input) + throws IOException { + final String device = ReadWriteIOUtils.readString(input); + final boolean aligned = ReadWriteIOUtils.readBool(input); + final long chunkGroupHeaderOffset = ReadWriteIOUtils.readLong(input); + final long chunkOffset = ReadWriteIOUtils.readLong(input); + final boolean firstChunkOfGroup = ReadWriteIOUtils.readBool(input); + final TSDataType dataType = TSDataType.values()[ReadWriteIOUtils.readInt(input)]; + final byte chunkType = input.readByte(); + final byte[] chunkHeaderBytes = readBytes(input); + final byte[] statisticsBytes = readBytes(input); + final long physicalStart = ReadWriteIOUtils.readLong(input); + final long physicalEnd = ReadWriteIOUtils.readLong(input); + return new ChunkRangeRecord( + device, + aligned, + chunkGroupHeaderOffset, + chunkOffset, + firstChunkOfGroup, + dataType, + chunkType, + chunkHeaderBytes, + statisticsBytes, + physicalStart, + physicalEnd); + } + + private static void writeBytes(final byte[] bytes, final DataOutputStream output) + throws IOException { + ReadWriteIOUtils.write(bytes.length, output); + output.write(bytes); + } + + private static byte[] readBytes(final DataInputStream input) throws IOException { + final int length = ReadWriteIOUtils.readInt(input); + if (length < 0) { + throw new IOException("Negative byte array length in progress file"); + } + final byte[] bytes = new byte[length]; + input.readFully(bytes); + return bytes; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java index 49eb75cf43289..62d593ae1f2fc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java @@ -113,7 +113,7 @@ public void writeChunk( chunkLength); if (isFirstChunkOfGroup) { - checkOffsetMismatch("chunk group header", chunkGroupHeaderOffset, currentPos, device, chunk); + chunkGroupHeaderOffset = alignToOffset(chunkGroupHeaderOffset, device, chunk); new ChunkGroupHeader(device).serializeTo(stream); currentPos = out.getPosition(); LOGGER.info( @@ -126,7 +126,7 @@ public void writeChunk( currentPos - chunkGroupHeaderOffset); } - checkOffsetMismatch("chunk", chunkOffset, currentPos, device, chunk); + chunkOffset = alignToOffset(chunkOffset, device, chunk); chunk.getHeader().serializeTo(stream); out.write(chunk.getData().duplicate()); @@ -160,33 +160,40 @@ public void writeChunk( .add(metadata); } - private void checkOffsetMismatch( - final String positionType, - final long expectedOffset, - final long actualOffset, - final IDeviceID device, - final Chunk chunk) + private long alignToOffset(final long expectedOffset, final IDeviceID device, final Chunk chunk) throws IOException { - if (expectedOffset != actualOffset) { - LOGGER.error( - "Precalculated offset mismatch: file={}, positionType={}, device={}, measurement={}, " - + "expectedOffset={}, actualOffset={}, delta={}", + final long currentOffset = out.getPosition(); + if (currentOffset < expectedOffset) { + long remaining = expectedOffset - currentOffset; + while (remaining > 0) { + final int step = (int) Math.min(Integer.MAX_VALUE, remaining); + out.write(new byte[step]); + remaining -= step; + } + LOGGER.warn( + "Filled physical hole before writing chunk: file={}, device={}, measurement={}, " + + "expectedOffset={}, actualOffset={}, fillBytes={}", + getFileForLog(), + device, + chunk.getHeader().getMeasurementID(), + expectedOffset, + currentOffset, + expectedOffset - currentOffset); + return expectedOffset; + } + if (currentOffset > expectedOffset) { + LOGGER.warn( + "Precalculated offset is behind actual file position; using actual position: file={}, " + + "device={}, measurement={}, expectedOffset={}, actualOffset={}, delta={}", getFileForLog(), - positionType, device, chunk.getHeader().getMeasurementID(), expectedOffset, - actualOffset, - actualOffset - expectedOffset); - throw new IOException( - String.format( - "Precalculated %s offset mismatch for %s/%s: expected=%d, actual=%d", - positionType, - device, - chunk.getHeader().getMeasurementID(), - expectedOffset, - actualOffset)); + currentOffset, + currentOffset - expectedOffset); + return currentOffset; } + return expectedOffset; } @Override From 6b53124d817f47ad84a4ffdee144380e9c00e12b Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Wed, 16 Sep 2026 11:31:41 +0800 Subject: [PATCH 6/8] Fix LOAD progress length and count assertions --- .../iotdb/db/it/IoTDBLoadTsFileClusterIT.java | 16 ++-- .../storageengine/load/LoadTsFileManager.java | 8 +- .../load/LoadTsFileProgress.java | 9 ++- .../load/LoadTsFileProgressTest.java | 80 +++++++++++++++++++ .../splitter/ChunkDataDirectWriteTest.java | 69 +++++++++++++++- 5 files changed, 172 insertions(+), 10 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgressTest.java diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java index e2fdf1069c660..3cc62889de815 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java @@ -127,12 +127,18 @@ public void testLoadTsFileInCluster() throws Exception { } statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); - try (final ResultSet resultSet = - statement.executeQuery("select count(*) from " + DATABASE + ".**")) { - Assert.assertTrue(resultSet.next()); - final long count = resultSet.getLong(1); - Assert.assertEquals(writtenPointCount, count); + long actualPointCount = 0; + for (final String device : DEVICES) { + try (final ResultSet resultSet = + statement.executeQuery( + "select count(s1), count(s2), count(s3), count(s4) from " + device)) { + Assert.assertTrue(resultSet.next()); + for (int columnIndex = 1; columnIndex <= MEASUREMENTS.size(); columnIndex++) { + actualPointCount += resultSet.getLong(columnIndex); + } + } } + Assert.assertEquals(writtenPointCount, actualPointCount); } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java index 40671567fb7c0..ddea08082046e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java @@ -594,9 +594,13 @@ private void prepare( dataPartition2Writer.entrySet()) { final TsFilePrecalculatedChunkWriter writer = entry.getValue(); final LoadTsFileProgress progress = dataPartition2Progress.get(entry.getKey()); - if (progress != null && progress.exists() && !progress.isReady(writer.getFile().length())) { + if (progress != null + && progress.exists() + && progress.getTotalLength() > 0 + && writer.getFile().length() < progress.getTotalLength()) { throw new LoadFileException( - "Staged LOAD TsFile is not contiguous yet: " + writer.getFile().getAbsolutePath()); + "Staged LOAD TsFile length does not reach prepare expectation: " + + writer.getFile().getAbsolutePath()); } writer.close(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java index 6ef3e042897f9..7c2ed3d3b5be7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java @@ -79,12 +79,12 @@ public boolean exists() { } public void activate(final long totalLength) throws IOException { - this.totalLength = totalLength; + this.totalLength = Math.max(this.totalLength, totalLength); if (!progressFile.exists()) { try (final DataOutputStream output = new DataOutputStream(new FileOutputStream(progressFile, false))) { ReadWriteIOUtils.write(MAGIC, output); - ReadWriteIOUtils.write(totalLength, output); + ReadWriteIOUtils.write(this.totalLength, output); } } } @@ -183,6 +183,11 @@ public List readAllRecords() throws IOException { } } } + long maxPhysicalEnd = this.totalLength; + for (final ChunkRangeRecord record : result) { + maxPhysicalEnd = Math.max(maxPhysicalEnd, record.physicalEnd()); + } + this.totalLength = maxPhysicalEnd; records.clear(); records.addAll(result); return result; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgressTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgressTest.java new file mode 100644 index 0000000000000..e35cb9e90202d --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgressTest.java @@ -0,0 +1,80 @@ +/* + * 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.storageengine.load; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.file.metadata.statistics.Statistics; +import org.apache.tsfile.read.common.Chunk; +import org.junit.Assert; +import org.junit.Test; + +import java.io.File; +import java.nio.ByteBuffer; +import java.nio.file.Files; + +public class LoadTsFileProgressTest { + + @Test + public void testTotalLengthKeepsMaximumWhenFillingEarlierHole() throws Exception { + final File tsFile = Files.createTempFile("load-progress", ".tsfile").toFile(); + final LoadTsFileProgress progress = new LoadTsFileProgress(tsFile); + + progress.activate(100L); + progress.activate(50L); + + Assert.assertEquals(100L, progress.getTotalLength()); + + Files.deleteIfExists(tsFile.toPath()); + Files.deleteIfExists(LoadTsFileProgress.progressFileFor(tsFile).toPath()); + } + + @Test + public void testRecordAndReadChunkRange() throws Exception { + final File tsFile = Files.createTempFile("load-progress-record", ".tsfile").toFile(); + final ChunkHeader chunkHeader = + new ChunkHeader( + "s1", 4, TSDataType.INT32, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 1); + final Statistics statistics = Statistics.getStatsByType(TSDataType.INT32); + statistics.update(1L, 1); + statistics.update(2L, 2); + final Chunk chunk = + new Chunk(chunkHeader, ByteBuffer.wrap(new byte[] {0, 0, 0, 1}), null, statistics); + + final LoadTsFileProgress progress = new LoadTsFileProgress(tsFile); + progress.recordChunk("root.sg.d1", false, 7L, 30L, true, chunk, 40L); + + Assert.assertEquals(40L, progress.getTotalLength()); + Assert.assertTrue(progress.isReady(40L)); + + final LoadTsFileProgress recovered = new LoadTsFileProgress(tsFile); + final LoadTsFileProgress.ChunkRangeRecord record = recovered.readAllRecords().get(0); + Assert.assertEquals("root.sg.d1", record.device()); + Assert.assertEquals(7L, record.physicalStart()); + Assert.assertEquals(40L, record.physicalEnd()); + Assert.assertEquals(30L, record.chunkOffset()); + Assert.assertTrue(record.firstChunkOfGroup()); + + Files.deleteIfExists(tsFile.toPath()); + Files.deleteIfExists(LoadTsFileProgress.progressFileFor(tsFile).toPath()); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java index 4e85ef36e1878..54f212b8dbdc2 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java @@ -21,7 +21,10 @@ import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.utils.TimePartitionUtils; +import org.apache.iotdb.db.queryengine.plan.scheduler.load.ChunkOffsetCalculator; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData.ChunkLayout; +import org.apache.tsfile.common.constant.TsFileConstant; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.header.ChunkHeader; import org.apache.tsfile.file.metadata.IChunkMetadata; @@ -32,11 +35,16 @@ import org.apache.tsfile.file.metadata.statistics.Statistics; import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.utils.TsPrimitiveType; +import org.apache.tsfile.write.TsFilePrecalculatedChunkWriter; import org.apache.tsfile.write.writer.TsFileIOWriter; import org.junit.Test; import org.mockito.Mockito; +import java.io.ByteArrayOutputStream; +import java.io.File; import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.util.ArrayList; import java.util.List; import static org.apache.tsfile.common.constant.TsFileConstant.TIME_COLUMN_MASK; @@ -146,6 +154,55 @@ public void testAlignedChunkSelectsPointsFromCurrentTimePartition() throws Excep assertEquals(partitionStart + 1, statistics.getStartTime()); } + @Test + public void testChunkOffsetCalculatorMatchesPrecalculatedWriter() throws Exception { + final File tsFile = + Files.createTempFile("chunk-offset-calculator", TsFileConstant.TSFILE_SUFFIX).toFile(); + Files.deleteIfExists(tsFile.toPath()); + + final ChunkOffsetCalculator calculator = new ChunkOffsetCalculator(); + final List chunkDataList = new ArrayList<>(); + + for (int deviceIndex = 0; deviceIndex < 5; deviceIndex++) { + final StringArrayDeviceID device = new StringArrayDeviceID("root", "sg", "d" + deviceIndex); + for (int measurementIndex = 0; measurementIndex < 4; measurementIndex++) { + final String measurement = "s" + measurementIndex; + final NonAlignedChunkData chunkData = + (NonAlignedChunkData) + ChunkData.createChunkData( + false, device, createChunkHeader(measurement), new TTimePartitionSlot(0L)); + chunkData.writeDecodePage( + new long[] {1L, 2L, 3L}, + new Object[] {measurementIndex, measurementIndex + 1, measurementIndex + 2}, + 3); + chunkData.endChunk(); + calculator.assign(chunkData); + chunkDataList.add(chunkData); + } + } + + try (final TsFilePrecalculatedChunkWriter writer = new TsFilePrecalculatedChunkWriter(tsFile)) { + for (final NonAlignedChunkData chunkData : chunkDataList) { + final ChunkLayout layout = chunkData.getChunkLayout(); + final Chunk chunk = chunkData.getChunks().get(0); + final long chunkLength = + serializeChunkHeaderSize(chunk.getHeader()) + chunk.getData().remaining(); + + writer.writeChunk( + chunkData.getDevice(), + chunkData.isAligned(), + layout.chunkGroupHeaderOffset(), + layout.firstChunkOfGroup(), + chunk, + layout.offset()); + + assertEquals(layout.offset() + chunkLength, tsFile.length()); + } + } finally { + Files.deleteIfExists(tsFile.toPath()); + } + } + private static Statistics createInt32Statistics() { final Statistics statistics = Statistics.getStatsByType(TSDataType.INT32); statistics.update(1L, 1); @@ -190,7 +247,17 @@ private static AlignedChunkData createAlignedTimeChunkData( } private static ChunkHeader createChunkHeader() { + return createChunkHeader("temperature"); + } + + private static ChunkHeader createChunkHeader(final String measurement) { return new ChunkHeader( - "temperature", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0); + measurement, 0, TSDataType.INT32, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0); + } + + private static int serializeChunkHeaderSize(final ChunkHeader chunkHeader) throws Exception { + try (final ByteArrayOutputStream output = new ByteArrayOutputStream()) { + return chunkHeader.serializeTo(output); + } } } From 2c9c9b09ef377918334b20e310ce0a79518da455 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Wed, 16 Sep 2026 11:44:28 +0800 Subject: [PATCH 7/8] Fix chunk offset writer test flushing --- .../db/storageengine/load/splitter/ChunkDataDirectWriteTest.java | 1 + 1 file changed, 1 insertion(+) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java index 54f212b8dbdc2..6d0cee8d3efbf 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java @@ -196,6 +196,7 @@ public void testChunkOffsetCalculatorMatchesPrecalculatedWriter() throws Excepti chunk, layout.offset()); + writer.getOutput().flush(); assertEquals(layout.offset() + chunkLength, tsFile.length()); } } finally { From 564e30058339c1802d1cf5a2ad72da35edbb8918 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Thu, 17 Sep 2026 15:00:19 +0800 Subject: [PATCH 8/8] Apply LOAD patch and remove temporary debug logs --- .../iotdb/db/it/IoTDBLoadTsFileClusterIT.java | 98 +++++++- .../iotdb/db/i18n/StorageEngineMessages.java | 4 + .../iotdb/db/i18n/StorageEngineMessages.java | 4 + .../org/apache/iotdb/db/conf/IoTDBConfig.java | 5 + .../dataregion/DataExecutionVisitor.java | 11 + .../load/TwoPhaseConsensusLoadStrategy.java | 15 ++ .../storageengine/load/LoadTsFileManager.java | 120 +++++++-- .../load/LoadTsFileProgress.java | 45 ++-- .../write/TsFilePrecalculatedChunkWriter.java | 74 +----- .../load/LoadTsFileManagerTest.java | 236 +++++++++++++++++ .../splitter/ChunkDataDirectWriteTest.java | 237 ++++++++++++++++++ 11 files changed, 747 insertions(+), 102 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java index 3cc62889de815..3fd4baf88da50 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java @@ -24,6 +24,7 @@ import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.it.utils.TsFileGenerator; import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.jdbc.Config; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.enums.TSEncoding; @@ -39,6 +40,7 @@ import java.io.File; import java.nio.file.Files; import java.sql.Connection; +import java.sql.DriverManager; import java.sql.ResultSet; import java.sql.Statement; import java.util.Arrays; @@ -53,8 +55,10 @@ public class IoTDBLoadTsFileClusterIT { private static final List DEVICES = Arrays.asList( DATABASE + ".d1", DATABASE + ".d2", DATABASE + ".d3", DATABASE + ".d4", DATABASE + ".d5"); + private static final List ALIGNED_DEVICES = + Arrays.asList(DATABASE + ".a1", DATABASE + ".a2"); private static final List MEASUREMENTS = Arrays.asList("s1", "s2", "s3", "s4"); - private static final int POINT_COUNT_PER_DEVICE = 20_000; + private static final int POINT_COUNT_PER_DEVICE = 10_000; private static final int REPLICATION_FACTOR = 3; private static final int CONFIG_NODE_NUM = 3; private static final int DATA_NODE_NUM = 3; @@ -68,6 +72,8 @@ public void setUp() throws Exception { EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false); EnvFactory.getEnv().getConfig().getCommonConfig().setPipeMemoryManagementEnabled(false); EnvFactory.getEnv().getConfig().getCommonConfig().setDatanodeMemoryProportion("1:10:1:1:1:0"); + EnvFactory.getEnv().getConfig().getCommonConfig().setTargetChunkPointNum(1_000); + EnvFactory.getEnv().getConfig().getCommonConfig().setMaxNumberOfPointsInPage(500); EnvFactory.getEnv() .getConfig() .getCommonConfig() @@ -113,6 +119,10 @@ public void testLoadTsFileInCluster() throws Exception { generator.registerTimeseries(device, schemas); generator.generateData(device, POINT_COUNT_PER_DEVICE, 1L, false); } + for (final String device : ALIGNED_DEVICES) { + generator.registerAlignedTimeseries(device, schemas); + generator.generateData(device, POINT_COUNT_PER_DEVICE, 1L, true, 1_000_000L); + } writtenPointCount = generator.getTotalNumber(); } @@ -125,6 +135,10 @@ public void testLoadTsFileInCluster() throws Exception { "create timeseries " + device + "." + measurement + " " + TSDataType.INT64.name()); } } + for (final String device : ALIGNED_DEVICES) { + statement.execute( + "create aligned timeseries " + device + "(s1 INT64, s2 INT64, s3 INT64, s4 INT64)"); + } statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); long actualPointCount = 0; @@ -141,4 +155,86 @@ public void testLoadTsFileInCluster() throws Exception { Assert.assertEquals(writtenPointCount, actualPointCount); } } + + @Test + public void testLoadAfterOneDataNodeDown() throws Exception { + final long firstPointCount = generateTsFile("load-first.tsfile", 5_000, 0L); + + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute("create database " + DATABASE); + for (final String device : DEVICES) { + for (final String measurement : MEASUREMENTS) { + statement.execute( + "create timeseries " + device + "." + measurement + " " + TSDataType.INT64.name()); + } + } + for (final String device : ALIGNED_DEVICES) { + statement.execute( + "create aligned timeseries " + device + "(s1 INT64, s2 INT64, s3 INT64, s4 INT64)"); + } + + statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); + } + + EnvFactory.getEnv().getDataNodeWrapper(0).stop(); + Thread.sleep(1_000L); + + final long secondPointCount = generateTsFile("load-second.tsfile", 3_000, 100_000L); + try (final Connection connection = + DriverManager.getConnection( + Config.IOTDB_URL_PREFIX + + EnvFactory.getEnv().getDataNodeWrapper(1).getIpAndPortString(), + "root", + "root"); + final Statement statement = connection.createStatement()) { + statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); + + final long expectedPointCount = firstPointCount + secondPointCount; + long actualPointCount = queryPointCount(statement); + for (int retry = 0; actualPointCount < expectedPointCount && retry < 30; retry++) { + actualPointCount = queryPointCount(statement); + } + Assert.assertEquals(expectedPointCount, actualPointCount); + } + } + + private long generateTsFile( + final String fileName, final int pointCountPerDevice, final long startTimestamp) + throws Exception { + final List schemas = + Arrays.asList( + new MeasurementSchema(MEASUREMENTS.get(0), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(1), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(2), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(3), TSDataType.INT64, TSEncoding.PLAIN)); + try (final TsFileGenerator generator = new TsFileGenerator(new File(tmpDir, fileName))) { + for (final String device : DEVICES) { + generator.registerTimeseries(device, schemas); + generator.generateData(device, pointCountPerDevice, 1L, false, startTimestamp); + } + final long alignedStartTimestamp = startTimestamp + 1_000_000L; + for (final String device : ALIGNED_DEVICES) { + generator.registerAlignedTimeseries(device, schemas); + generator.generateData(device, pointCountPerDevice, 1L, true, alignedStartTimestamp); + } + return generator.getTotalNumber(); + } + } + + private long queryPointCount(final Statement statement) throws Exception { + long actualPointCount = 0; + final List allDevices = new java.util.ArrayList<>(DEVICES); + allDevices.addAll(ALIGNED_DEVICES); + for (final String device : allDevices) { + for (final String measurement : MEASUREMENTS) { + try (final ResultSet resultSet = + statement.executeQuery("select count(" + measurement + ") from " + device)) { + Assert.assertTrue(resultSet.next()); + actualPointCount += resultSet.getLong(1); + } + } + } + return actualPointCount; + } } diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 2e3bba6cd3e57..39745b301e9f3 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -530,6 +530,10 @@ private StorageEngineMessages() {} "Preparing LOAD TsFile {}: sealing staged resources."; public static final String LOG_COMMITTING_LOAD_TSFILE_ARG_LOADING_PREPARED_RESOURCES_INTO_DATAREGION_EA1D6335 = "Committing LOAD TsFile {}: loading prepared resources into DataRegion."; + public static final String LOG_RECEIVE_LOAD_TSFILE_NODE_ARG_C36E832B = + "Receive LOAD TsFile node: {}."; + public static final String LOG_RECEIVE_LOAD_TSFILE_NODE_SUCCESS_ARG_27F8ECD6 = + "Receive LOAD TsFile node success: {}."; public static final String EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93 = "Table '%s.%s' does not exist when applying LOAD chunk data. It may have been dropped after the LOAD was analyzed."; public static final String TASK_DIR_NOT_EMPTY_SKIP_DELETE = "Task dir {} is not empty, skip deleting."; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 44c3126c6d446..883cc39f4a71f 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -528,6 +528,10 @@ private StorageEngineMessages() {} "正在准备 LOAD TsFile {}：封存暂存资源。"; public static final String LOG_COMMITTING_LOAD_TSFILE_ARG_LOADING_PREPARED_RESOURCES_INTO_DATAREGION_EA1D6335 = "正在提交 LOAD TsFile {}：将已准备资源加载到 DataRegion。"; + public static final String LOG_RECEIVE_LOAD_TSFILE_NODE_ARG_C36E832B = + "接收 LOAD TsFile 节点：{}。"; + public static final String LOG_RECEIVE_LOAD_TSFILE_NODE_SUCCESS_ARG_27F8ECD6 = + "接收 LOAD TsFile 节点成功：{}。"; public static final String EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93 = "应用 LOAD chunk 数据时表 '%s.%s' 不存在，可能在 LOAD 分析之后被删除了。"; public static final String TASK_DIR_NOT_EMPTY_SKIP_DELETE = "任务目录 {} 非空，跳过删除。"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java index ef6f9b08ca688..4f75f5eecaf1b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java @@ -1704,6 +1704,11 @@ public String[] getLoadTsFileDirs() { return this.loadTsFileDirs; } + public void setLoadTsFileDirs(final String[] loadTsFileDirs) { + this.loadTsFileDirs = loadTsFileDirs; + this.loadTsFileDirCanonicalPaths = canonicalPaths(loadTsFileDirs); + } + public String[] getLoadTsFileAllowedDirs() { return this.loadTsFileAllowedDirs.length == 0 ? getLoadTsFileDirs() diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java index c143d048baa10..fe7132f0afe9e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java @@ -36,6 +36,7 @@ import org.apache.iotdb.db.exception.runtime.TableLostRuntimeException; import org.apache.iotdb.db.i18n.DataNodeMiscMessages; import org.apache.iotdb.db.i18n.DataNodePipeMessages; +import org.apache.iotdb.db.i18n.StorageEngineMessages; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedDeleteDataNode; @@ -75,12 +76,16 @@ public TSStatus visitPlan(PlanNode node, DataRegion context) { @Override public TSStatus visitLoadTsFileConsensus( final LoadTsFileConsensusNode node, final DataRegion dataRegion) { + LOGGER.info(StorageEngineMessages.LOG_RECEIVE_LOAD_TSFILE_NODE_ARG_C36E832B, node); switch (node.getOp()) { case BEGIN: + LOGGER.info(StorageEngineMessages.LOG_RECEIVE_LOAD_TSFILE_NODE_SUCCESS_ARG_27F8ECD6, node); return StatusUtils.OK; case PIECE: try { dataRegion.writeLoadTsFilePiece(node); + LOGGER.info( + StorageEngineMessages.LOG_RECEIVE_LOAD_TSFILE_NODE_SUCCESS_ARG_27F8ECD6, node); return StatusUtils.OK; } catch (final IOException | PageException e) { LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); @@ -95,6 +100,8 @@ public TSStatus visitLoadTsFileConsensus( progressIndexes.put( slot, ProgressIndexType.deserializeFrom(ByteBuffer.wrap(bytes)))); final boolean prepared = dataRegion.writeLoadTsFilePrepare(node, progressIndexes); + LOGGER.info( + StorageEngineMessages.LOG_RECEIVE_LOAD_TSFILE_NODE_SUCCESS_ARG_27F8ECD6, node); return prepared ? StatusUtils.OK : RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR); } catch (final Exception e) { LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); @@ -109,6 +116,8 @@ public TSStatus visitLoadTsFileConsensus( progressIndexes.put( slot, ProgressIndexType.deserializeFrom(ByteBuffer.wrap(bytes)))); final boolean committed = dataRegion.writeLoadTsFileCommit(node, progressIndexes); + LOGGER.info( + StorageEngineMessages.LOG_RECEIVE_LOAD_TSFILE_NODE_SUCCESS_ARG_27F8ECD6, node); return committed ? StatusUtils.OK : RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR); } catch (final Exception e) { LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); @@ -117,6 +126,8 @@ public TSStatus visitLoadTsFileConsensus( case ABORT: try { dataRegion.writeLoadTsFileAbort(node); + LOGGER.info( + StorageEngineMessages.LOG_RECEIVE_LOAD_TSFILE_NODE_SUCCESS_ARG_27F8ECD6, node); return StatusUtils.OK; } catch (final Exception e) { LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java index fd12732957377..63106c985a2c2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java @@ -243,6 +243,15 @@ private boolean dispatchConsensusPiece( regionLoadIds.computeIfAbsent(regionId, ignored -> UUID.randomUUID().toString()); final long pieceIndex = pieceNode.getPieceIndex(); + LOGGER.info( + "Dispatch LOAD piece: loadId={}, regionId={}, pieceIndex={}, tsFile={}, dataSize={}, " + + "replicaSet={}", + loadId, + regionId, + pieceIndex, + pieceNode.getTsFile() == null ? null : pieceNode.getTsFile().getName(), + pieceNode.getDataSize(), + replicaSet); final LoadTsFileConsensusNode piece = LoadTsFileConsensusNode.piece( new PlanNodeId("load-piece-" + loadId + "-" + pieceIndex), @@ -262,6 +271,12 @@ private boolean dispatchConsensusPiece( pieceNode); return false; } + LOGGER.info( + "Dispatch LOAD piece success: loadId={}, regionId={}, pieceIndex={}, dataSize={}", + loadId, + regionId, + pieceIndex, + piece.getDataSize()); regionPieceCounts.merge(regionId, 1L, Long::sum); regionTotalBytes.merge(regionId, piece.getDataSize(), Long::sum); return true; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java index ddea08082046e..bf8c0d665167e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java @@ -119,6 +119,7 @@ public class LoadTsFileManager { public LoadTsFileManager(final DataRegion dataRegion) { this.dataRegion = Objects.requireNonNull(dataRegion); + recover(); } public void start() {} @@ -127,6 +128,35 @@ public void stop() { new HashSet<>(uuid2WriterManager.keySet()).forEach(this::forceCloseWriterManager); } + private void recover() { + final File[] baseDirs = + Arrays.stream(CONFIG.getLoadTsFileDirs()) + .map(File::new) + .map(this::getDataRegionLoadDir) + .toArray(File[]::new); + for (final File baseDir : baseDirs) { + final File[] uuidDirs = baseDir.listFiles(); + if (uuidDirs == null) { + continue; + } + for (final File uuidDir : uuidDirs) { + if (!uuidDir.isDirectory()) { + continue; + } + final String uuid = uuidDir.getName(); + try { + final TsFileWriterManager recovered = new TsFileWriterManager(uuidDir).recoverFromDisk(); + if (recovered.hasPendingTsFiles()) { + uuid2WriterManager.put(uuid, recovered); + LOGGER.info("Recovered LOAD writer manager for uuid {}", uuid); + } + } catch (final Exception e) { + LOGGER.warn("Failed to recover LOAD writer manager for uuid {}", uuid, e); + } + } + } + } + public void writeToDataRegion(LoadTsFilePieceNode pieceNode, String uuid) throws IOException, PageException { writePiece(uuid, pieceNode.getAllTsFileData()); @@ -224,6 +254,9 @@ private static List createPieceRefs( if (!file.isFile()) { continue; } + if (file.getName().endsWith(LoadTsFileProgress.PROGRESS_SUFFIX)) { + continue; + } final long previousLength = previousLengths.getOrDefault(file, 0L); final long size = file.length() - previousLength; if (size > 0) { @@ -445,7 +478,7 @@ public static void cleanTsFile(final File tsFile) { } } - private static class TsFileWriterManager { + private class TsFileWriterManager { private final File taskDir; private Map dataPartition2Writer; @@ -473,6 +506,53 @@ private void ensureDir(File dir) { } } + private TsFileWriterManager recoverFromDisk() throws IOException { + final File[] files = taskDir.listFiles(); + if (files == null) { + return this; + } + for (final File progressFile : files) { + if (!progressFile.getName().endsWith(LoadTsFileProgress.PROGRESS_SUFFIX)) { + continue; + } + final String progressName = progressFile.getName(); + final File tsFile = + new File( + progressFile.getParentFile(), + progressName.substring( + 0, progressName.length() - LoadTsFileProgress.PROGRESS_SUFFIX.length())); + if (!tsFile.isFile()) { + continue; + } + final LoadTsFileProgress progress = new LoadTsFileProgress(tsFile); + final List records = progress.readAllRecords(); + if (records.isEmpty()) { + continue; + } + + final TTimePartitionSlot timePartitionSlot = + new TTimePartitionSlot(parseTimePartitionStart(tsFile.getName())); + final DataPartitionInfo partitionInfo = + new DataPartitionInfo(dataRegion, timePartitionSlot); + dataPartition2Resource.put(partitionInfo, new TsFileResource(tsFile)); + dataPartition2Progress.put(partitionInfo, new LoadTsFileProgress(tsFile)); + if (tsFile.length() >= progress.getTotalLength()) { + isPrepared = true; + } + } + return this; + } + + private boolean hasPendingTsFiles() { + return !dataPartition2Resource.isEmpty(); + } + + private long parseTimePartitionStart(final String tsFileName) { + final String baseName = tsFileName.substring(0, tsFileName.lastIndexOf('.')); + final int lastDash = baseName.lastIndexOf(IoTDBConstant.FILE_NAME_SEPARATOR); + return Long.parseLong(baseName.substring(lastDash + 1)); + } + /** * It should be noted that all AlignedChunkData of the same partition split from a source file * should be guaranteed to be written to the same new file. Otherwise, for detached @@ -511,24 +591,31 @@ private void write(DataPartitionInfo partitionInfo, ChunkData chunkData) final Chunk chunk = chunks.get(i); final long chunkLength = getChunkHeaderSerializedSize(chunk.getHeader()) + (long) chunk.getData().remaining(); - writer.writeChunk( - chunkData.getDevice(), - chunkData.isAligned(), - layout.chunkGroupHeaderOffset(), - layout.firstChunkOfGroup() && i == 0, - chunk, - chunkOffset); + final TsFilePrecalculatedChunkWriter.ChunkWriteResult writeResult = + writer.writeChunk( + chunkData.getDevice(), + chunkData.isAligned(), + layout.chunkGroupHeaderOffset(), + layout.firstChunkOfGroup() && i == 0, + chunk, + chunkOffset); chunkOffset += chunkLength; + final boolean firstChunkOfGroup = layout.firstChunkOfGroup() && i == 0; + final long actualPhysicalStart = + firstChunkOfGroup + ? writeResult.actualChunkGroupHeaderOffset() + : writeResult.actualChunkOffset(); dataPartition2Progress .get(partitionInfo) .recordChunk( chunkData.getDevice().toString(), chunkData.isAligned(), - layout.chunkGroupHeaderOffset(), - chunkOffset - chunkLength, - layout.firstChunkOfGroup() && i == 0, + writeResult.actualChunkGroupHeaderOffset(), + writeResult.actualChunkOffset(), + firstChunkOfGroup, chunk, - chunkOffset); + writeResult.actualChunkEndOffset(), + actualPhysicalStart); } } @@ -623,15 +710,16 @@ private void loadAll(boolean isGeneratedByPipe) throws LoadFileException { dataRegion.loadNewTsFile(tsFileResource, true, isGeneratedByPipe, false, Optional.empty()); // Metrics + final TsFilePrecalculatedChunkWriter writer = dataPartition2Writer.get(entry.getKey()); + if (writer == null) { + continue; + } dataRegion .getNonSystemDatabaseName() .ifPresent( databaseName -> updateWritePointCountMetrics( - dataRegion, - databaseName, - getTsFileWritePointCount(dataPartition2Writer.get(entry.getKey())), - false)); + dataRegion, databaseName, getTsFileWritePointCount(writer), false)); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java index 7c2ed3d3b5be7..2049bea36f408 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileProgress.java @@ -38,7 +38,6 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Comparator; import java.util.List; /** @@ -98,11 +97,32 @@ public void recordChunk( final Chunk chunk, final long dataEnd) throws IOException { + final long physicalStart = firstChunkOfGroup ? chunkGroupHeaderOffset : chunkOffset; + recordChunk( + device, + aligned, + chunkGroupHeaderOffset, + chunkOffset, + firstChunkOfGroup, + chunk, + dataEnd, + physicalStart); + } + + public void recordChunk( + final String device, + final boolean aligned, + final long chunkGroupHeaderOffset, + final long chunkOffset, + final boolean firstChunkOfGroup, + final Chunk chunk, + final long dataEnd, + final long actualPhysicalStart) + throws IOException { activate(dataEnd); final byte[] chunkHeaderBytes = serializeChunkHeader(chunk); final byte[] statisticsBytes = serializeStatistics(chunk); - final long physicalStart = firstChunkOfGroup ? chunkGroupHeaderOffset : chunkOffset; final ChunkRangeRecord record = new ChunkRangeRecord( @@ -115,7 +135,7 @@ public void recordChunk( chunk.getHeader().getChunkType(), chunkHeaderBytes, statisticsBytes, - physicalStart, + actualPhysicalStart, dataEnd); records.add(record); @@ -130,25 +150,10 @@ public boolean isReady(final long currentFileLength) throws IOException { return false; } - final List sorted = readAllRecords(); - if (sorted.isEmpty()) { - return false; - } - sorted.sort(Comparator.comparingLong(ChunkRangeRecord::physicalStart)); - - long coveredStart = sorted.get(0).physicalStart(); - long coveredEnd = sorted.get(0).physicalEnd(); - if (coveredStart != TS_FILE_HEADER_SIZE) { + if (readAllRecords().isEmpty()) { return false; } - for (int i = 1; i < sorted.size(); i++) { - final ChunkRangeRecord record = sorted.get(i); - if (record.physicalStart() > coveredEnd) { - return false; - } - coveredEnd = Math.max(coveredEnd, record.physicalEnd()); - } - return coveredEnd == currentFileLength; + return currentFileLength >= totalLength; } public List readAllRecords() throws IOException { diff --git a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java index 62d593ae1f2fc..6a4b591eef6a2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java @@ -72,13 +72,9 @@ public TsFilePrecalculatedChunkWriter(TsFileOutput out) throws IOException { private void startFile() throws IOException { out.write(BytesUtils.stringToBytes(TSFileConfig.MAGIC_STRING)); out.write(new byte[] {TSFileConfig.VERSION_NUMBER}); - LOGGER.info( - "Started precalculated TsFile writer: file={}, headerEndOffset={}", - getFileForLog(), - out.getPosition()); } - public void writeChunk( + public ChunkWriteResult writeChunk( IDeviceID device, boolean isAligned, long chunkGroupHeaderOffset, @@ -91,56 +87,18 @@ public void writeChunk( final int chunkHeaderSize = chunk.getHeader().getSerializedSize(); final int chunkDataSize = chunk.getData().remaining(); final long chunkLength = chunkHeaderSize + (long) chunkDataSize; - - LOGGER.info( - "Writing precalculated Chunk: file={}, device={}, aligned={}, measurement={}, dataType={}, " - + "chunkType={}, pageCount={}, firstChunkOfGroup={}, expectedGroupHeaderOffset={}, " - + "expectedChunkOffset={}, actualOffsetBeforeWrite={}, chunkHeaderSize={}, " - + "chunkDataSize={}, chunkLength={}", - getFileForLog(), - device, - isAligned, - chunk.getHeader().getMeasurementID(), - chunk.getHeader().getDataType(), - chunk.getHeader().getChunkType(), - chunk.getHeader().getNumOfPages(), - isFirstChunkOfGroup, - chunkGroupHeaderOffset, - chunkOffset, - currentPos, - chunkHeaderSize, - chunkDataSize, - chunkLength); + long actualChunkGroupHeaderOffset = -1L; if (isFirstChunkOfGroup) { - chunkGroupHeaderOffset = alignToOffset(chunkGroupHeaderOffset, device, chunk); + actualChunkGroupHeaderOffset = alignToOffset(chunkGroupHeaderOffset, device, chunk); new ChunkGroupHeader(device).serializeTo(stream); currentPos = out.getPosition(); - LOGGER.info( - "Wrote ChunkGroupHeader: file={}, device={}, headerOffset={}, headerEndOffset={}, " - + "headerLength={}", - getFileForLog(), - device, - chunkGroupHeaderOffset, - currentPos, - currentPos - chunkGroupHeaderOffset); } chunkOffset = alignToOffset(chunkOffset, device, chunk); chunk.getHeader().serializeTo(stream); out.write(chunk.getData().duplicate()); - - LOGGER.info( - "Wrote precalculated Chunk: file={}, device={}, measurement={}, chunkOffset={}, " - + "chunkEndOffset={}, expectedChunkEndOffset={}, metadataMask={}", - getFileForLog(), - device, - chunk.getHeader().getMeasurementID(), - chunkOffset, - out.getPosition(), - chunkOffset + chunkLength, - chunk.getHeader().getChunkType() - & (TsFileConstant.TIME_COLUMN_MASK | TsFileConstant.VALUE_COLUMN_MASK)); + final long actualChunkEndOffset = out.getPosition(); ChunkMetadata metadata = new ChunkMetadata( @@ -158,6 +116,8 @@ public void writeChunk( .computeIfAbsent(device, k -> new TreeMap<>()) .computeIfAbsent(chunk.getHeader().getMeasurementID(), k -> new ArrayList<>()) .add(metadata); + + return new ChunkWriteResult(actualChunkGroupHeaderOffset, chunkOffset, actualChunkEndOffset); } private long alignToOffset(final long expectedOffset, final IDeviceID device, final Chunk chunk) @@ -199,17 +159,6 @@ private long alignToOffset(final long expectedOffset, final IDeviceID device, fi @Override public void close() throws IOException { final long metaOffset = out.getPosition(); - LOGGER.info( - "Closing precalculated TsFile writer: file={}, dataEndOffset={}, deviceCount={}, " - + "seriesCount={}, chunkCount={}", - getFileForLog(), - metaOffset, - device2MetadataMap.size(), - device2MetadataMap.values().stream().mapToInt(Map::size).sum(), - device2MetadataMap.values().stream() - .flatMap(measurements -> measurements.values().stream()) - .mapToInt(List::size) - .sum()); final OutputStream stream = out.wrapAsStream(); ReadWriteIOUtils.write(MetaMarker.SEPARATOR, stream); final Map deviceMetadataIndexMap = new TreeMap<>(); @@ -261,14 +210,6 @@ public void close() throws IOException { ReadWriteIOUtils.write(tsFileMetadataSize, stream); out.write(BytesUtils.stringToBytes(TSFileConfig.MAGIC_STRING)); - LOGGER.info( - "Finished precalculated TsFile writer: file={}, metaOffset={}, tsFileMetadataOffset={}, " - + "tsFileMetadataSize={}, fileEndOffset={}", - getFileForLog(), - metaOffset, - tsFileMetadataOffset, - tsFileMetadataSize, - out.getPosition()); out.close(); } @@ -298,4 +239,7 @@ public Map> getChunkMetadataListMap() { .collect(java.util.stream.Collectors.toList()))); return result; } + + public record ChunkWriteResult( + long actualChunkGroupHeaderOffset, long actualChunkOffset, long actualChunkEndOffset) {} } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java new file mode 100644 index 0000000000000..0332da6b7c35a --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java @@ -0,0 +1,236 @@ +/* + * 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.storageengine.load; + +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.db.conf.IoTDBConfig; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; +import org.apache.iotdb.db.queryengine.plan.scheduler.load.ChunkOffsetCalculator; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.NonAlignedChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.TsFileReader; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.common.RowRecord; +import org.apache.tsfile.read.expression.QueryExpression; +import org.apache.tsfile.read.query.dataset.QueryDataSet; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; + +import java.io.File; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class LoadTsFileManagerTest { + + private File tempDir; + private String[] originalLoadBaseDirs; + private DataRegion dataRegion; + private IoTDBConfig config; + + @Before + public void setUp() throws Exception { + tempDir = Files.createTempDirectory("load-tsfile-manager-test").toFile(); + config = IoTDBDescriptor.getInstance().getConfig(); + originalLoadBaseDirs = config.getLoadTsFileDirs(); + config.setLoadTsFileDirs(new String[] {tempDir.getAbsolutePath()}); + + dataRegion = Mockito.mock(DataRegion.class); + Mockito.when(dataRegion.getDatabaseName()).thenReturn("root.load_manager_test"); + Mockito.when(dataRegion.getDataRegionIdString()).thenReturn("0"); + Mockito.when(dataRegion.getNonSystemDatabaseName()) + .thenReturn(Optional.of("root.load_manager_test")); + } + + @After + public void tearDown() throws Exception { + config.setLoadTsFileDirs(originalLoadBaseDirs); + deleteRecursively(tempDir); + } + + @Test + public void testOutOfOrderPieceWriteKeepsProgressAndDataReadable() throws Exception { + final String uuid = "test-uuid"; + final int pointCount = 100; + final int deviceCount = 5; + final int measurementCount = 4; + final ChunkOffsetCalculator calculator = new ChunkOffsetCalculator(); + final List earlierChunks = new ArrayList<>(); + final List laterChunks = new ArrayList<>(); + final List expectations = new ArrayList<>(); + int valueBase = 0; + + for (int deviceIndex = 0; deviceIndex < deviceCount; deviceIndex++) { + final StringArrayDeviceID device = + new StringArrayDeviceID("root", "load_manager_test", "d" + deviceIndex); + for (int measurementIndex = 0; measurementIndex < measurementCount; measurementIndex++) { + final String measurement = "s" + measurementIndex; + final NonAlignedChunkData chunkData = + createNonAlignedChunkData(device, measurement, valueBase, pointCount); + calculator.assign(chunkData); + if (deviceIndex < 2) { + earlierChunks.add(chunkData); + } else { + laterChunks.add(chunkData); + } + expectations.add( + new SeriesExpectation(device, measurement, pointCount, valueBase, valueBase + 99)); + valueBase += pointCount; + } + } + + final LoadTsFileManager manager = new LoadTsFileManager(dataRegion); + File tsFile = null; + try { + // Normally the earlier piece would be sent first. Simulate the opposite: the later piece + // arrives first, then the earlier piece is filled in afterwards. + final List laterRefs = + manager.writePiece(uuid, toTsFileDataList(laterChunks)); + assertEquals(1, laterRefs.size()); + tsFile = new File(laterRefs.get(0).getRelativePath()); + + manager.writePiece(uuid, toTsFileDataList(earlierChunks)); + assertTrue(manager.prepare(uuid, 2, 1L, false, Collections.emptyMap())); + + final LoadTsFileProgress progress = new LoadTsFileProgress(tsFile); + final List records = progress.readAllRecords(); + assertEquals(deviceCount * measurementCount, records.size()); + assertTrue(progress.isReady(tsFile.length())); + assertAllSeriesReadable(tsFile, expectations); + } finally { + manager.deleteAll(uuid); + } + } + + private static NonAlignedChunkData createNonAlignedChunkData( + final IDeviceID device, final String measurement, final int valueBase, final int pointCount) { + final NonAlignedChunkData chunkData = + (NonAlignedChunkData) + ChunkData.createChunkData( + false, + device, + new ChunkHeader( + measurement, + 0, + TSDataType.INT32, + CompressionType.UNCOMPRESSED, + TSEncoding.PLAIN, + 0), + new TTimePartitionSlot(0L)); + final long[] times = new long[pointCount]; + final Object[] values = new Object[pointCount]; + for (int i = 0; i < pointCount; i++) { + times[i] = i + 1L; + values[i] = valueBase + i; + } + chunkData.writeDecodePage(times, values, pointCount); + chunkData.endChunk(); + return chunkData; + } + + private static List toTsFileDataList(final List chunkDataList) { + final List result = new ArrayList<>(); + for (final ChunkData chunkData : chunkDataList) { + result.add(chunkData); + } + return result; + } + + private static void assertAllSeriesReadable( + final File tsFile, final List expectations) throws Exception { + try (final TsFileReader reader = new TsFileReader(tsFile)) { + for (final SeriesExpectation expectation : expectations) { + final QueryExpression queryExpression = + QueryExpression.create( + Collections.singletonList( + new Path(expectation.device, expectation.measurement, false)), + null); + final QueryDataSet dataSet = reader.query(queryExpression); + int count = 0; + while (dataSet.hasNext()) { + final RowRecord row = dataSet.next(); + final org.apache.tsfile.read.common.Field field = row.getField(0); + final int value = field.getIntV(); + if (count == 0) { + assertEquals(expectation.firstValue, value); + } + if (count == expectation.pointCount - 1) { + assertEquals(expectation.lastValue, value); + } + count++; + } + assertEquals(expectation.pointCount, count); + } + } + } + + private static void deleteRecursively(final File file) { + if (file == null || !file.exists()) { + return; + } + if (file.isDirectory()) { + final File[] files = file.listFiles(); + if (files != null) { + for (final File child : files) { + deleteRecursively(child); + } + } + } + file.delete(); + } + + private static final class SeriesExpectation { + private final IDeviceID device; + private final String measurement; + private final int pointCount; + private final int firstValue; + private final int lastValue; + + private SeriesExpectation( + final IDeviceID device, + final String measurement, + final int pointCount, + final int firstValue, + final int lastValue) { + this.device = device; + this.measurement = measurement; + this.pointCount = pointCount; + this.firstValue = firstValue; + this.lastValue = lastValue; + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java index 6d0cee8d3efbf..45be593f242ad 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkDataDirectWriteTest.java @@ -26,6 +26,7 @@ import org.apache.tsfile.common.constant.TsFileConstant; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkGroupHeader; import org.apache.tsfile.file.header.ChunkHeader; import org.apache.tsfile.file.metadata.IChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID; @@ -33,7 +34,14 @@ import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; import org.apache.tsfile.file.metadata.statistics.Statistics; +import org.apache.tsfile.read.TsFileReader; +import org.apache.tsfile.read.TsFileSequenceReader; import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.read.common.Field; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.common.RowRecord; +import org.apache.tsfile.read.expression.QueryExpression; +import org.apache.tsfile.read.query.dataset.QueryDataSet; import org.apache.tsfile.utils.TsPrimitiveType; import org.apache.tsfile.write.TsFilePrecalculatedChunkWriter; import org.apache.tsfile.write.writer.TsFileIOWriter; @@ -45,6 +53,7 @@ import java.nio.ByteBuffer; import java.nio.file.Files; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import static org.apache.tsfile.common.constant.TsFileConstant.TIME_COLUMN_MASK; @@ -204,6 +213,201 @@ public void testChunkOffsetCalculatorMatchesPrecalculatedWriter() throws Excepti } } + @Test + public void testRichChunkOffsetCalculatorMatchesWriterForAlignedAndNonAligned() throws Exception { + final File tsFile = + Files.createTempFile("rich-chunk-offset-calculator", TsFileConstant.TSFILE_SUFFIX).toFile(); + Files.deleteIfExists(tsFile.toPath()); + + final ChunkOffsetCalculator calculator = new ChunkOffsetCalculator(); + final List chunkDataList = new ArrayList<>(); + final int pointCount = 300_000; + final long[] times = new long[pointCount]; + final Object[] values = new Object[pointCount]; + for (int i = 0; i < pointCount; i++) { + times[i] = i + 1L; + values[i] = i; + } + + for (int deviceIndex = 0; deviceIndex < 5; deviceIndex++) { + final StringArrayDeviceID device = new StringArrayDeviceID("root", "rich", "d" + deviceIndex); + for (int measurementIndex = 0; measurementIndex < 4; measurementIndex++) { + final NonAlignedChunkData chunkData = + createNonAlignedChunkData(device, "s" + measurementIndex); + chunkData.writeDecodePage(times, values, pointCount); + chunkData.endChunk(); + calculator.assign(chunkData); + chunkDataList.add(chunkData); + } + } + + final AlignedChunkData alignedChunkData = createAlignedTimeChunkData(); + alignedChunkData.writeDecodePage(times, new Object[pointCount], pointCount); + alignedChunkData.endChunk(); + alignedChunkData.addValueChunk(createChunkHeader("s1")); + final TsPrimitiveType[] alignedValues = new TsPrimitiveType[pointCount]; + for (int i = 0; i < pointCount; i++) { + alignedValues[i] = TsPrimitiveType.getByType(TSDataType.INT32, i); + } + alignedChunkData.writeDecodeValuePage(times, alignedValues, TSDataType.INT32); + alignedChunkData.endChunk(); + calculator.assign(alignedChunkData); + chunkDataList.add(alignedChunkData); + + try (final TsFilePrecalculatedChunkWriter writer = new TsFilePrecalculatedChunkWriter(tsFile)) { + for (final ChunkData chunkData : chunkDataList) { + final ChunkData.ChunkLayout layout = chunkData.getChunkLayout(); + final List chunks = chunkData.getChunks(); + long chunkOffset = layout.offset(); + for (int i = 0; i < chunks.size(); i++) { + final Chunk chunk = chunks.get(i); + final long chunkLength = + serializeChunkHeaderSize(chunk.getHeader()) + chunk.getData().remaining(); + writer.writeChunk( + chunkData.getDevice(), + chunkData.isAligned(), + layout.chunkGroupHeaderOffset(), + layout.firstChunkOfGroup() && i == 0, + chunk, + chunkOffset); + chunkOffset += chunkLength; + } + writer.getOutput().flush(); + assertEquals( + chunkData.getChunkLayout().offset() + chunkData.getChunkLayout().length(), + tsFile.length()); + } + } finally { + try (final TsFileSequenceReader reader = new TsFileSequenceReader(tsFile.getAbsolutePath())) { + assertEquals(6, reader.getAllDevices().size()); + } + assertSeriesReadable( + tsFile, + new StringArrayDeviceID("root", "rich", "d0"), + "s0", + false, + pointCount, + 0, + pointCount - 1, + TSDataType.INT32); + assertChunkMetadataCount( + tsFile, new StringArrayDeviceID("root", "sg", "d1"), "s1", true, pointCount); + Files.deleteIfExists(tsFile.toPath()); + } + } + + @Test + public void testPrecalculatedWriterHandlesOutOfOrderChunkDispatch() throws Exception { + final File tsFile = + Files.createTempFile("out-of-order-chunk-write", TsFileConstant.TSFILE_SUFFIX).toFile(); + Files.deleteIfExists(tsFile.toPath()); + + final ChunkOffsetCalculator calculator = new ChunkOffsetCalculator(); + final StringArrayDeviceID firstDevice = new StringArrayDeviceID("root", "order", "d0"); + final StringArrayDeviceID secondDevice = new StringArrayDeviceID("root", "order", "d1"); + final NonAlignedChunkData firstChunkData = createNonAlignedChunkData(firstDevice, "s0", 0, 4); + final NonAlignedChunkData secondChunkData = + createNonAlignedChunkData(secondDevice, "s0", 10, 4); + calculator.assign(firstChunkData); + calculator.assign(secondChunkData); + + final ChunkLayout secondLayout = secondChunkData.getChunkLayout(); + final ChunkLayout firstLayout = firstChunkData.getChunkLayout(); + final Chunk secondChunk = secondChunkData.getChunks().get(0); + final Chunk firstChunk = firstChunkData.getChunks().get(0); + final long secondChunkLength = + serializeChunkHeaderSize(secondChunk.getHeader()) + secondChunk.getData().remaining(); + final long firstChunkLength = + serializeChunkHeaderSize(firstChunk.getHeader()) + firstChunk.getData().remaining(); + final long expectedLengthAfterSecondChunk = secondLayout.offset() + secondChunkLength; + final long expectedFinalLength = + expectedLengthAfterSecondChunk + + serializeChunkGroupHeaderSize(firstDevice) + + firstChunkLength; + + try (final TsFilePrecalculatedChunkWriter writer = new TsFilePrecalculatedChunkWriter(tsFile)) { + writer.writeChunk( + secondChunkData.getDevice(), + secondChunkData.isAligned(), + secondLayout.chunkGroupHeaderOffset(), + secondLayout.firstChunkOfGroup(), + secondChunk, + secondLayout.offset()); + writer.getOutput().flush(); + assertEquals(expectedLengthAfterSecondChunk, tsFile.length()); + + writer.writeChunk( + firstChunkData.getDevice(), + firstChunkData.isAligned(), + firstLayout.chunkGroupHeaderOffset(), + firstLayout.firstChunkOfGroup(), + firstChunk, + firstLayout.offset()); + writer.getOutput().flush(); + assertEquals(expectedFinalLength, tsFile.length()); + } finally { + assertSeriesReadable(tsFile, firstDevice, "s0", false, 4, 0, 3, TSDataType.INT32); + assertSeriesReadable(tsFile, secondDevice, "s0", false, 4, 10, 13, TSDataType.INT32); + Files.deleteIfExists(tsFile.toPath()); + } + } + + private static void assertSeriesReadable( + final File tsFile, + final IDeviceID device, + final String measurement, + final boolean isAligned, + final int expectedPointCount, + final int expectedFirstValue, + final int expectedLastValue, + final TSDataType dataType) + throws Exception { + final QueryExpression queryExpression = + QueryExpression.create( + Collections.singletonList(new Path(device, measurement, isAligned)), null); + try (final TsFileReader reader = new TsFileReader(tsFile)) { + final QueryDataSet dataSet = reader.query(queryExpression); + int count = 0; + while (dataSet.hasNext()) { + final RowRecord row = dataSet.next(); + final Field field = row.getField(0); + final int value; + switch (dataType) { + case INT32: + value = field.getIntV(); + break; + case INT64: + value = Math.toIntExact(field.getLongV()); + break; + default: + throw new IllegalArgumentException("Unsupported type in assertion: " + dataType); + } + if (count == 0) { + assertEquals(expectedFirstValue, value); + } + if (count == expectedPointCount - 1) { + assertEquals(expectedLastValue, value); + } + count++; + } + assertEquals(expectedPointCount, count); + } + } + + private static void assertChunkMetadataCount( + final File tsFile, + final IDeviceID device, + final String measurement, + final boolean isAligned, + final int expectedPointCount) + throws Exception { + try (final TsFileSequenceReader reader = new TsFileSequenceReader(tsFile.getAbsolutePath())) { + final org.apache.tsfile.file.metadata.ChunkMetadata chunkMetadata = + reader.getChunkMetadataList(new Path(device, measurement, isAligned), false).get(0); + assertEquals(expectedPointCount, chunkMetadata.getStatistics().getCount()); + } + } + private static Statistics createInt32Statistics() { final Statistics statistics = Statistics.getStatsByType(TSDataType.INT32); statistics.update(1L, 1); @@ -221,6 +425,33 @@ private static NonAlignedChunkData createNonAlignedChunkData( ChunkData.createChunkData(false, device, createChunkHeader(), timePartitionSlot); } + private static NonAlignedChunkData createNonAlignedChunkData( + final IDeviceID device, final String measurement) { + return (NonAlignedChunkData) + ChunkData.createChunkData( + false, device, createChunkHeader(measurement), new TTimePartitionSlot(0L)); + } + + private static NonAlignedChunkData createNonAlignedChunkData( + final IDeviceID device, + final String measurement, + final int startValue, + final int pointCount) { + final NonAlignedChunkData chunkData = + (NonAlignedChunkData) + ChunkData.createChunkData( + false, device, createChunkHeader(measurement), new TTimePartitionSlot(0L)); + final long[] times = new long[pointCount]; + final Object[] values = new Object[pointCount]; + for (int i = 0; i < pointCount; i++) { + times[i] = i + 1L; + values[i] = startValue + i; + } + chunkData.writeDecodePage(times, values, pointCount); + chunkData.endChunk(); + return chunkData; + } + private static AlignedChunkData createAlignedChunkData() { final IDeviceID device = new StringArrayDeviceID("root", "sg", "d1"); return (AlignedChunkData) @@ -261,4 +492,10 @@ private static int serializeChunkHeaderSize(final ChunkHeader chunkHeader) throw return chunkHeader.serializeTo(output); } } + + private static int serializeChunkGroupHeaderSize(final IDeviceID device) throws Exception { + try (final ByteArrayOutputStream output = new ByteArrayOutputStream()) { + return new ChunkGroupHeader(device).serializeTo(output); + } + } }