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 a3d39c6fcd65..0cb700627ef2 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 @@ -19,8 +19,10 @@ package org.apache.iotdb.db.it; +import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; import org.apache.iotdb.consensus.ConsensusFactory; import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.it.utils.TsFileGenerator; import org.apache.iotdb.itbase.category.ClusterIT; @@ -49,6 +51,7 @@ import java.sql.Statement; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -96,6 +99,14 @@ public static void setUpClass() { .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS) .setSchemaReplicationFactor(REPLICATION_FACTOR) .setDataReplicationFactor(REPLICATION_FACTOR); + // The restart case compares the on-disk size of the replicas, so keep compaction from changing + // those sizes while the test waits for the restarted DataNode to catch up. + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setEnableSeqSpaceCompaction(false) + .setEnableUnseqSpaceCompaction(false) + .setEnableCrossSpaceCompaction(false); EnvFactory.getEnv().initClusterEnvironment(CONFIG_NODE_NUM, DATA_NODE_NUM); } @@ -260,6 +271,161 @@ public void testLoadAfterOneDataNodeDown() throws Exception { } } + /** + * A LOAD that runs while one DataNode is down is replicated by consensus to the surviving + * replicas. When the stopped DataNode restarts, its replicas replay the LOAD commands they + * missed, so the restarted DataNode ends up serving the very same file. + */ + @Test + public void testRestartedDataNodeIsSyncedWithTheLoadedFile() throws Exception { + final long firstPointCount = generateTsFile(tmpDir, "sync-first.tsfile", 5_000, 1L); + + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + createDatabaseAndSchema(statement); + statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); + awaitPointCount(statement, firstPointCount); + } + + // LOAD the second file while DataNode 0 is down. With three replicas it holds a copy of every + // region, so it is the DataNode that has to catch up with the LOAD when it comes back. + final DataNodeWrapper restartedDataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + restartedDataNode.stop(); + Assert.assertFalse(restartedDataNode.isAlive()); + + final long secondPointCount = generateTsFile(tmpDir, "sync-second.tsfile", 3_000, 100_000L); + final long expectedPointCount = firstPointCount + secondPointCount; + 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() + "\""); + awaitPointCount(statement, expectedPointCount); + } + + restartedDataNode.start(); + Assert.assertTrue(restartedDataNode.isAlive()); + + // A pinned connection may still read from another replica, so check the restarted node's own + // replicas before letting it answer the query. + awaitRestartedLocalReplicasHaveTheLoadedFile(restartedDataNode); + awaitPointCount(restartedDataNode, expectedPointCount); + } + + /** + * Waits until every data region replica on the given DataNode has the same on-disk size as the + * other replicas of that region. A node that did not replay the LOAD keeps its pre-LOAD size. + */ + private static void awaitRestartedLocalReplicasHaveTheLoadedFile( + final DataNodeWrapper restartedDataNode) { + Awaitility.await() + .atMost(5, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted( + () -> { + final Map> regionSizes = dataRegionSizes(); + Assert.assertFalse( + "the loaded database must have data regions", regionSizes.isEmpty()); + + int restartedDataNodeId; + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + restartedDataNodeId = dataNodeId(statement, restartedDataNode); + } + + for (final Map.Entry> region : regionSizes.entrySet()) { + final Map replicaSizes = region.getValue(); + // A region without data on any replica has nothing to catch up with. + if (replicaSizes.values().stream().noneMatch(IoTDBLoadTsFileClusterIT::hasData)) { + continue; + } + final String restartedSize = replicaSizes.get(restartedDataNodeId); + Assert.assertTrue( + "region " + region.getKey() + " has no data on the restarted DataNode", + hasData(restartedSize)); + replicaSizes.forEach( + (dataNodeId, size) -> { + if (dataNodeId != restartedDataNodeId) { + Assert.assertEquals( + "region " + region.getKey() + " differs on DataNode " + dataNodeId, + size, + restartedSize); + } + }); + } + }); + } + + /** Whether a `show regions` size means the replica holds data. */ + private static boolean hasData(final String size) { + return size != null && !"Unknown".equals(size) && !"0 B".equals(size); + } + + /** The DataRegion rows of `show regions`, mapped to each replica's on-disk size. */ + private static Map> dataRegionSizes() throws Exception { + final Map> regionSizes = new HashMap<>(); + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery("show regions")) { + while (resultSet.next()) { + if (!"DataRegion".equals(resultSet.getString(ColumnHeaderConstant.TYPE)) + || !DATABASE.equals(resultSet.getString(ColumnHeaderConstant.DATABASE))) { + continue; + } + regionSizes + .computeIfAbsent( + resultSet.getInt(ColumnHeaderConstant.REGION_ID), id -> new HashMap<>()) + .put( + resultSet.getInt(ColumnHeaderConstant.DATA_NODE_ID), + resultSet.getString(ColumnHeaderConstant.TSFILE_SIZE)); + } + } + return regionSizes; + } + + /** The cluster id of the given DataNode, matched by its client RPC endpoint. */ + private static int dataNodeId(final Statement statement, final DataNodeWrapper dataNode) + throws SQLException { + try (final ResultSet resultSet = statement.executeQuery("show datanodes")) { + while (resultSet.next()) { + if (dataNode.getIp().equals(resultSet.getString(ColumnHeaderConstant.RPC_ADDRESS)) + && dataNode.getPort() == resultSet.getInt(ColumnHeaderConstant.RPC_PORT)) { + return resultSet.getInt(ColumnHeaderConstant.NODE_ID); + } + } + } + throw new AssertionError("cannot find " + dataNode.getIpAndPortString() + " in show datanodes"); + } + + /** Waits until the given statement sees the expected number of imported points. */ + private static void awaitPointCount(final Statement statement, final long expectedPointCount) { + Awaitility.await() + .atMost(2, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted(() -> Assert.assertEquals(expectedPointCount, queryPointCount(statement))); + } + + /** Waits until the points of every loaded file are queryable through the given DataNode. */ + private static void awaitPointCount( + final DataNodeWrapper dataNode, final long expectedPointCount) { + Awaitility.await() + .atMost(2, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(dataNode); + final Statement statement = connection.createStatement()) { + Assert.assertEquals(expectedPointCount, queryPointCount(statement)); + } + }); + } + /** With two of three DataNodes down the LOAD must fail with a message. */ @Test public void testLoadWithTwoDataNodesDownFailsWithAMeaningfulMessage() throws Exception { @@ -365,7 +531,7 @@ private long generateTsFile( } } - private long queryPointCount(final Statement statement) throws Exception { + private static long queryPointCount(final Statement statement) throws Exception { long actualPointCount = 0; final List allDevices = new ArrayList<>(DEVICES); allDevices.addAll(ALIGNED_DEVICES); diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileRegionMigrationIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileRegionMigrationIT.java new file mode 100644 index 000000000000..9de078017e9a --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileRegionMigrationIT.java @@ -0,0 +1,403 @@ +/* + * 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.commons.cluster.RegionStatus; +import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.config.MppDataNodeConfig; +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.awaitility.Awaitility; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.Before; +import org.junit.BeforeClass; +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.SQLException; +import java.sql.Statement; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +/** + * LOAD TsFile has to keep working when the topology changes while a file is being imported: the + * route of a region can be re-resolved after a piece was rejected, and a LOAD that notices a + * migration falls back to inserting the file through the ordinary write path. Either way every + * point of the file has to end up in the cluster and the migration has to finish. + * + *

The data replication factor is 2 with three DataNodes, so a region always has exactly one + * DataNode that can be its migration destination. + */ +@RunWith(IoTDBTestRunner.class) +@Category({ClusterIT.class}) +public class IoTDBLoadTsFileRegionMigrationIT { + + private static final long PARTITION_INTERVAL = 100_000_000L; + private static final String DATABASE = "root.load_migrate"; + 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 DATA_REPLICATION_FACTOR = 2; + private static final int SCHEMA_REPLICATION_FACTOR = 3; + private static final int CONFIG_NODE_NUM = 1; + private static final int DATA_NODE_NUM = 3; + + /** + * Keep the LOAD blocked after it dispatches a piece until the migration becomes active, so the + * two operations cannot silently run one after the other on a fast machine. + */ + private static final double THROTTLED_LOAD_WRITE_THROUGHPUT_BYTES_PER_SECOND = 1; + + /** Sets the LOAD write throughput in the DataNode properties before the cluster starts. */ + private static void setLoadWriteThroughput(final double bytesPerSecond) { + final MppDataNodeConfig dataNodeConfig = + (MppDataNodeConfig) EnvFactory.getEnv().getConfig().getDataNodeConfig(); + final Properties properties = new Properties(); + properties.setProperty( + "load_write_throughput_bytes_per_second", String.valueOf(bytesPerSecond)); + dataNodeConfig.updateProperties(properties); + } + + private File tmpDir; + + @BeforeClass + public static void setUpClass() { + 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().setTargetChunkPointNum(1_000); + EnvFactory.getEnv().getConfig().getCommonConfig().setMaxNumberOfPointsInPage(500); + setLoadWriteThroughput(THROTTLED_LOAD_WRITE_THROUGHPUT_BYTES_PER_SECOND); + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setConfigNodeConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS) + .setSchemaReplicationFactor(SCHEMA_REPLICATION_FACTOR) + .setDataReplicationFactor(DATA_REPLICATION_FACTOR); + EnvFactory.getEnv().initClusterEnvironment(CONFIG_NODE_NUM, DATA_NODE_NUM); + } + + @AfterClass + public static void tearDownClass() { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Before + public void setUp() throws Exception { + tmpDir = new File(Files.createTempDirectory("load-migrate-it").toUri()); + } + + @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) { + } + + final File[] files = tmpDir.listFiles(); + if (files != null) { + for (final File file : files) { + Assert.assertTrue(deleteRecursively(file)); + } + } + Assert.assertTrue(tmpDir.delete()); + } + + /** + * Starts a LOAD of a file large enough to still be dispatching pieces when a region of its + * database is migrated away from a DataNode it was routed to, then asserts that both the LOAD and + * the migration finish with every point of the file imported. + */ + @Test + public void testLoadWhileARegionMigrates() throws Exception { + final long writtenPointCount = generateTsFile(new File(tmpDir, "load-migrating.tsfile")); + + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + createDatabaseAndSchema(statement); + } + + final ExecutorService loader = Executors.newSingleThreadExecutor(); + try { + final Future load = + loader.submit( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); + return null; + } catch (final Exception e) { + return e; + } + }); + + // Wait until the LOAD allocated the data region of its points, then migrate that region while + // the LOAD is still splitting and dispatching the pieces of the file. + final int[] migration = awaitMigratableRegion(); + Assert.assertFalse("the LOAD finished before the migration could overlap it", load.isDone()); + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute( + "migrate region " + migration[0] + " from " + migration[1] + " to " + migration[2]); + } + + // The migration procedure is submitted asynchronously. Keep the LOAD throttled until the + // procedure has actually started changing the region replicas, then resume it so the two + // operations overlap instead of running one after the other. + awaitRegionMigrationActive(migration[0]); + Assert.assertFalse("the LOAD finished while the migration was starting", load.isDone()); + restoreLoadWriteThroughput(); + + // The LOAD may follow the migration with a refreshed route or fall back to the tablet path, + // but it must not report a failure to the client. + final Exception loadFailure = load.get(15, TimeUnit.MINUTES); + Assert.assertNull("the LOAD must survive a region migration: " + loadFailure, loadFailure); + + awaitRegionMigrated(migration[0], migration[1], migration[2]); + + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + awaitPointCount(statement, writtenPointCount); + } + } finally { + try { + restoreLoadWriteThroughput(); + } catch (final Exception ignored) { + // Preserve the original failure; the cluster is discarded by tearDownClass. + } + loader.shutdownNow(); + } + } + + private static void restoreLoadWriteThroughput() throws SQLException { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute("set configuration \"load_write_throughput_bytes_per_second\"=\"-1\""); + } + } + + /** + * Waits until the LOAD has allocated a data region of the database, and returns the region, one + * of the DataNodes that replicates it and a DataNode that does not, in that order. + */ + private static int[] awaitMigratableRegion() { + final AtomicReference migration = new AtomicReference<>(); + Awaitility.await() + .atMost(5, TimeUnit.MINUTES) + .pollInterval(200, TimeUnit.MILLISECONDS) + .ignoreExceptions() + .untilAsserted( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + final Set allDataNodes = allDataNodeIds(statement); + final Map> regionMap = dataRegionMap(statement); + Assert.assertFalse("the LOAD must allocate a data region", regionMap.isEmpty()); + + for (final Map.Entry> entry : regionMap.entrySet()) { + final Set replicas = entry.getValue(); + final int destination = + allDataNodes.stream() + .filter(dataNodeId -> !replicas.contains(dataNodeId)) + .findAny() + .orElse(-1); + if (destination < 0) { + continue; + } + migration.set( + new int[] {entry.getKey(), replicas.iterator().next(), destination}); + return; + } + Assert.fail("no data region has a free migration destination yet"); + } + }); + return migration.get(); + } + + /** Waits until the migration procedure has started changing the region's replicas. */ + private static void awaitRegionMigrationActive(final int regionId) { + Awaitility.await() + .atMost(2, TimeUnit.MINUTES) + .pollInterval(200, TimeUnit.MILLISECONDS) + .ignoreExceptions() + .untilAsserted( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery("show regions")) { + while (resultSet.next()) { + if (!"DataRegion".equals(resultSet.getString(ColumnHeaderConstant.TYPE)) + || regionId != resultSet.getInt(ColumnHeaderConstant.REGION_ID)) { + continue; + } + final String status = resultSet.getString(ColumnHeaderConstant.STATUS); + if (RegionStatus.Adding.getStatus().equals(status) + || RegionStatus.Removing.getStatus().equals(status)) { + return; + } + } + } + Assert.fail("region " + regionId + " has not started migrating"); + }); + } + + /** Waits until the migrated region left the source DataNode and reached the destination. */ + private static void awaitRegionMigrated( + final int regionId, final int sourceDataNode, final int destinationDataNode) { + Awaitility.await() + .atMost(5, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + final Set replicas = dataRegionMap(statement).get(regionId); + Assert.assertNotNull("the migrated region must still exist", replicas); + Assert.assertFalse( + "the migrated region must leave DataNode " + sourceDataNode, + replicas.contains(sourceDataNode)); + Assert.assertTrue( + "the migrated region must reach DataNode " + destinationDataNode, + replicas.contains(destinationDataNode)); + } + }); + } + + /** Waits until the given statement sees the expected number of imported points. */ + private static void awaitPointCount(final Statement statement, final long expectedPointCount) { + Awaitility.await() + .atMost(5, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted(() -> Assert.assertEquals(expectedPointCount, queryPointCount(statement))); + } + + private void createDatabaseAndSchema(final Statement statement) throws SQLException { + statement.execute("create database " + DATABASE); + for (final String device : DEVICES) { + for (final String measurement : MEASUREMENTS) { + statement.execute( + "create timeseries " + device + "." + measurement + " " + TSDataType.INT64.name()); + } + } + } + + /** The DataRegion rows of `show regions`, mapped to the DataNodes that replicate each region. */ + private static Map> dataRegionMap(final Statement statement) + throws SQLException { + final Map> regionMap = new HashMap<>(); + try (final ResultSet resultSet = statement.executeQuery("show regions")) { + while (resultSet.next()) { + if (!"DataRegion".equals(resultSet.getString(ColumnHeaderConstant.TYPE)) + || !DATABASE.equals(resultSet.getString(ColumnHeaderConstant.DATABASE))) { + continue; + } + final int regionId = resultSet.getInt(ColumnHeaderConstant.REGION_ID); + final int dataNodeId = resultSet.getInt(ColumnHeaderConstant.DATA_NODE_ID); + regionMap.computeIfAbsent(regionId, id -> new HashSet<>()).add(dataNodeId); + } + } + return regionMap; + } + + private static Set allDataNodeIds(final Statement statement) throws SQLException { + final Set dataNodeIds = new HashSet<>(); + try (final ResultSet resultSet = statement.executeQuery("show datanodes")) { + while (resultSet.next()) { + dataNodeIds.add(resultSet.getInt(ColumnHeaderConstant.NODE_ID)); + } + } + return dataNodeIds; + } + + private long generateTsFile(final File tsFile) 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(tsFile)) { + for (final String device : DEVICES) { + generator.registerTimeseries(device, schemas); + generator.generateData(device, POINT_COUNT_PER_DEVICE, 1L, false, 1L); + } + return generator.getTotalNumber(); + } + } + + private static long queryPointCount(final Statement statement) throws Exception { + long actualPointCount = 0; + for (final String device : DEVICES) { + 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; + } + + private static boolean deleteRecursively(final File file) { + final File[] children = file.listFiles(); + if (children != null) { + for (final File child : children) { + if (!deleteRecursively(child)) { + return false; + } + } + } + return file.delete(); + } +} diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index f5ecb62b048f..dfb2fb7d441b 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java @@ -1605,4 +1605,10 @@ private DataNodeMiscMessages() {} public static final String LOG_FAILED_TO_COLLECT_THE_DEVICES_USING_THE_EXTENDED_TEMPLATE_ARG_AB32D554 = "Failed to collect the devices using the extended template {}"; + public static final String + MISC_EXCEPTION_LOAD_TSFILE_WRITE_BUFFER_SIZE_SHOULD_BE_GREATER_THAN_0_AB1BA936 = + "loadTsFileWriteBufferSize should be greater than 0"; + public static final String EXCEPTION_BUFFER_SIZE_SHOULD_BE_GREATER_THAN_0_BB68B5C4 = + "bufferSize should be greater than 0"; + } 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 294af1dbb28d..25b73c301527 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 @@ -84,8 +84,10 @@ private StorageEngineMessages() {} "The LOAD snapshot progress file {} cannot be read back: {}"; 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 + EXCEPTION_UNEXPECTED_END_OF_FILE_WHEN_READING_THE_STAGED_PIECE_ARG_ARG_BYTES_WERE_REQUESTED_FROM_OFFSET_ARG_BUT_THE_FILE_ENDED_AT_POSITION_ARG_3F9CA32B = + "Unexpected end of file when reading the staged piece %s: %d bytes were requested from " + + "offset %d, but the file ended at position %d."; public static final String EXCEPTION_STAGED_PIECE_FILE_ARG_IS_MISSING_AND_CANNOT_BE_READ_BACK_4F62F9C6 = "Staged piece file %s is missing and cannot be read back."; public static final String EXCEPTION_LOAD_CONSENSUS_INVALID_CHUNK_PAYLOAD_REF_PATH_ARG_OFFSET_ARG_SIZE_ARG_6B70A61B = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index 26e6a16a5558..d91103c16bab 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java @@ -1587,4 +1587,10 @@ private DataNodeMiscMessages() {} public static final String LOG_FAILED_TO_COLLECT_THE_DEVICES_USING_THE_EXTENDED_TEMPLATE_ARG_AB32D554 = "收集使用扩展模板 {} 的设备失败"; + public static final String + MISC_EXCEPTION_LOAD_TSFILE_WRITE_BUFFER_SIZE_SHOULD_BE_GREATER_THAN_0_AB1BA936 = + "loadTsFileWriteBufferSize 应大于 0"; + public static final String EXCEPTION_BUFFER_SIZE_SHOULD_BE_GREATER_THAN_0_BB68B5C4 = + "bufferSize 应大于 0"; + } 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 49f0e760e0a3..93f9b29b6f8d 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 @@ -83,8 +83,9 @@ private StorageEngineMessages() {} "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 + EXCEPTION_UNEXPECTED_END_OF_FILE_WHEN_READING_THE_STAGED_PIECE_ARG_ARG_BYTES_WERE_REQUESTED_FROM_OFFSET_ARG_BUT_THE_FILE_ENDED_AT_POSITION_ARG_3F9CA32B = + "读取已暂存的分片 %s 时意外到达文件末尾：期望读取 %d 字节（起始偏移 %d），但文件在位置 %d 处结束。"; public static final String EXCEPTION_STAGED_PIECE_FILE_ARG_IS_MISSING_AND_CANNOT_BE_READ_BACK_4F62F9C6 = "已暂存的分片文件 %s 不存在，无法回读。"; public static final String EXCEPTION_LOAD_CONSENSUS_INVALID_CHUNK_PAYLOAD_REF_PATH_ARG_OFFSET_ARG_SIZE_ARG_6B70A61B = 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 f8006e8db8f4..919351473b52 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 @@ -1180,6 +1180,8 @@ public class IoTDBConfig implements DefaultEncodingProvider { private double loadWriteThroughputBytesPerSecond = -1; // Bytes/s + private int loadTsFileWriteBufferSize = 1024 * 1024; // 1 MiB + private long loadTabletConversionThresholdBytes = -1; private boolean loadActiveListeningEnable = true; @@ -4326,6 +4328,19 @@ public void setLoadWriteThroughputBytesPerSecond(double loadWriteThroughputBytes this.loadWriteThroughputBytesPerSecond = loadWriteThroughputBytesPerSecond; } + public int getLoadTsFileWriteBufferSize() { + return loadTsFileWriteBufferSize; + } + + public void setLoadTsFileWriteBufferSize(final int loadTsFileWriteBufferSize) { + if (loadTsFileWriteBufferSize <= 0) { + throw new IllegalArgumentException( + DataNodeMiscMessages + .MISC_EXCEPTION_LOAD_TSFILE_WRITE_BUFFER_SIZE_SHOULD_BE_GREATER_THAN_0_AB1BA936); + } + this.loadTsFileWriteBufferSize = loadTsFileWriteBufferSize; + } + public long getLoadTabletConversionThresholdBytes() { return loadTabletConversionThresholdBytes; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java index 71d67f0f1b88..b22ad2dc9692 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java @@ -2621,6 +2621,12 @@ private void loadLoadTsFileProps(TrimProperties properties) { "load_write_throughput_bytes_per_second", String.valueOf(conf.getLoadWriteThroughputBytesPerSecond())))); + conf.setLoadTsFileWriteBufferSize( + Integer.parseInt( + properties.getProperty( + "load_tsfile_write_buffer_size_in_byte", + String.valueOf(conf.getLoadTsFileWriteBufferSize())))); + conf.setLoadTsFileAllowedDirs( Arrays.stream(properties.getProperty("load_tsfile_allowed_dirs", "").trim().split(",")) .filter(dir -> !dir.isEmpty()) @@ -2770,6 +2776,12 @@ private void loadLoadTsFileHotModifiedProp(TrimProperties properties) throws IOE ConfigurationFileUtils.getConfigurationDefaultValue( "load_write_throughput_bytes_per_second")))); + conf.setLoadTsFileWriteBufferSize( + Integer.parseInt( + properties.getProperty( + "load_tsfile_write_buffer_size_in_byte", + String.valueOf(conf.getLoadTsFileWriteBufferSize())))); + conf.setLoadTsFileAllowedDirs( Arrays.stream( properties 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 402c7afc8ae5..745643678bae 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 @@ -27,6 +27,7 @@ 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.WritePlanNode; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; import org.apache.tsfile.exception.NotImplementedException; @@ -220,7 +221,8 @@ public String toString() { + '}'; } - public static class ByteBufferInputStream extends InputStream { + public static class ByteBufferInputStream extends InputStream + implements ChunkData.ByteBufferSliceSource { private final ByteBuffer buffer; public ByteBufferInputStream(ByteBuffer buffer) { @@ -245,7 +247,8 @@ public int read(byte[] b, int off, int len) { return toRead; } - public ByteBuffer read(int length) { + @Override + public ByteBuffer readSlice(int length) { if (length < 0 || length > buffer.remaining()) { throw new IllegalArgumentException( DataNodeQueryMessages.INVALID_LENGTH_FOR_SLICING + length); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionResolver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionResolver.java index f20e74636ff2..3a9dbadf12f0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionResolver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionResolver.java @@ -43,7 +43,8 @@ class DataPartitionResolver { private final String userName; /** The replica set resolved for every slot so far. */ - private final Map resolvedPartitions = new HashMap<>(); + private final Map, TRegionReplicaSet> resolvedPartitions = + new HashMap<>(); DataPartitionResolver(final DataPartitionBatchFetcher partitionFetcher, final String userName) { this.partitionFetcher = @@ -63,7 +64,7 @@ List resolve(final List> final List indexesToFetch = new ArrayList<>(); for (int i = 0; i < slotList.size(); i++) { final Pair slot = slotList.get(i); - final TRegionReplicaSet resolved = resolvedPartitions.get(new SlotKey(slot.left, slot.right)); + final TRegionReplicaSet resolved = resolvedPartitions.get(slot); replicaSets.add(resolved); if (resolved == null) { slotsToFetch.add(slot); @@ -88,13 +89,10 @@ List resolve(final List> final TRegionReplicaSet replicaSet = fetched.get(i); replicaSets.set(indexesToFetch.get(i), replicaSet); final Pair slot = slotsToFetch.get(i); - resolvedPartitions.put(new SlotKey(slot.left, slot.right), replicaSet); + resolvedPartitions.put(slot, replicaSet); } } return replicaSets; } - - /** One (device, time partition) pair, the key the partitions are resolved by. */ - private record SlotKey(IDeviceID device, TTimePartitionSlot slot) {} } 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 index 0d7478b307fa..f9caf454f0b7 100644 --- 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 @@ -67,8 +67,9 @@ List route(final List chunkDataList) { final int totalChunks = chunkDataList.size(); final int[] chunkPartitionIndexes = new int[totalChunks]; - final List distinctSlotList = new ArrayList<>(); - final Map slotToIndexMap = new HashMap<>(totalChunks); + final List> distinctSlotList = new ArrayList<>(); + final Map, Integer> slotToIndexMap = + new HashMap<>(totalChunks); // 1. Deduplicate (device, time-partition) pairs and record mapping index per chunk for (int i = 0; i < totalChunks; i++) { @@ -88,7 +89,7 @@ List route(final List chunkDataList) { chunkData.getTimePartitionSlot(), DataNodeQueryMessages.EXCEPTION_CHUNK_TIME_PARTITION_SLOT_CANNOT_BE_NULL_E5B04C6F); - final SlotKey slotKey = new SlotKey(device, slot); + final Pair slotKey = new Pair<>(device, slot); final Integer existingIndex = slotToIndexMap.get(slotKey); if (existingIndex == null) { @@ -102,8 +103,7 @@ List route(final List chunkDataList) { } // 2. Resolve the replica sets of the distinct slots, reusing the ones the judgement resolved. - final List replicaSets = - partitionResolver.resolve(toQuerySlots(distinctSlotList)); + final List replicaSets = partitionResolver.resolve(distinctSlotList); // 3. Map resolved replica sets back to original input order final List routedReplicaSets = new ArrayList<>(totalChunks); @@ -111,34 +111,17 @@ List route(final List chunkDataList) { final int partitionIndex = chunkPartitionIndexes[i]; final TRegionReplicaSet replicaSet = replicaSets.get(partitionIndex); if (replicaSet == null) { - final SlotKey slotKey = distinctSlotList.get(partitionIndex); + final Pair slotKey = distinctSlotList.get(partitionIndex); throw new IllegalStateException( String.format( DataNodeQueryMessages .EXCEPTION_NULL_REPLICA_SET_RESOLVED_FOR_DEVICE_ARG_AT_PARTITION_ARG_19CA94B2, - slotKey.device(), - slotKey.slot())); + slotKey.getLeft(), + slotKey.getRight())); } routedReplicaSets.add(replicaSet); } return routedReplicaSets; } - - // ------------------------------------------------------------------------- - // Internal Model - // ------------------------------------------------------------------------- - - /** The distinct slots in the pair form the resolver takes. */ - private static List> toQuerySlots( - final List distinctSlots) { - final List> slots = new ArrayList<>(distinctSlots.size()); - for (final SlotKey slot : distinctSlots) { - slots.add(new Pair<>(slot.device(), slot.slot())); - } - return slots; - } - - /** Composite key for slot deduplication; the only slot representation this class holds. */ - private record SlotKey(IDeviceID device, TTimePartitionSlot slot) {} } 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 370e6854c5f4..7354a9d21b07 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 @@ -28,6 +28,7 @@ import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.PageException; +import org.apache.tsfile.file.MetaMarker; import org.apache.tsfile.file.header.ChunkHeader; import org.apache.tsfile.file.header.PageHeader; import org.apache.tsfile.file.metadata.IChunkMetadata; @@ -50,11 +51,14 @@ import java.io.DataOutputStream; import java.io.IOException; import java.io.InputStream; +import java.io.OutputStream; +import java.nio.BufferOverflowException; import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.LinkedHashSet; +import java.util.LinkedHashMap; import java.util.LinkedList; import java.util.List; +import java.util.Map; import java.util.Queue; import java.util.Set; @@ -154,6 +158,13 @@ public AlignedChunkData( addAttrDataSize(); } + /** + * Detaches a consumed chunk data from the splitter: the copy is stored back into the map the + * splitter keeps draining, so the pages that follow land in a fresh object instead of the object + * that was just handed to the consumer. Without the copy the same object would be consumed again + * with every later flush and would re-emit the chunks it already emitted, which leaves the split + * TsFile with more time chunks than each value column has chunks. + */ protected AlignedChunkData(final AlignedChunkData alignedChunkData) { this( alignedChunkData.device, @@ -325,18 +336,45 @@ public void endChunk() { } protected void encodeAndBuildChunks() throws IOException, PageException { - final AlignedChunkWriterImpl writer = createAlignedChunkWriter(); + // Group the buffered pages by chunk once; scanning the whole page list once per chunk turns + // wide aligned tables into O(columns * pages). + final Map> pagesByChunkIndex = new LinkedHashMap<>(); + for (final PageBuffer pageBuffer : pageBuffers) { + pagesByChunkIndex + .computeIfAbsent(pageBuffer.chunkIndex, ignored -> new ArrayList<>()) + .add(pageBuffer); + } - writePageBuffersToWriter(writer); + // Fast-path: assemble chunks that need no splitting directly, avoiding re-encoding and its + // buffer amplification. + final Map directlyBuiltChunks = new LinkedHashMap<>(); + for (final Map.Entry> entry : pagesByChunkIndex.entrySet()) { + final int chunkIndex = entry.getKey(); + final Chunk chunk = + buildChunkFromUndecodedPages( + chunkHeaderList.get(chunkIndex), + chunkIndex == 0 ? TIME_COLUMN_MASK : VALUE_COLUMN_MASK, + entry.getValue()); + if (chunk != null) { + directlyBuiltChunks.put(chunkIndex, chunk); + } + } - // 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); + final AlignedChunkWriterImpl writer = + directlyBuiltChunks.size() == pagesByChunkIndex.size() ? null : createAlignedChunkWriter(); + if (writer != null) { + writePageBuffersToWriter(writer, directlyBuiltChunks.keySet()); } - for (final int chunkIndex : bufferedChunkIndexes) { + + for (final int chunkIndex : pagesByChunkIndex.keySet()) { + final Chunk directChunk = directlyBuiltChunks.get(chunkIndex); + if (directChunk != null) { + entireChunks.add(directChunk); + dataSize += + directChunk.getData().remaining() + directChunk.getChunkStatistic().getSerializedSize(); + continue; + } + final ChunkHeader sourceHeader = chunkHeaderList.get(chunkIndex); final ByteBuffer encodedData; final Statistics statistics; @@ -362,6 +400,132 @@ protected void encodeAndBuildChunks() throws IOException, PageException { } } + /** + * Reassembles a chunk from pages that were stored without decoding. Returns {@code null} when the + * chunk has a decoded page, a non-empty page without statistics, or no page with data at all; the + * caller then falls back to {@link AlignedChunkWriterImpl}. + */ + protected Chunk buildChunkFromUndecodedPages( + final ChunkHeader sourceHeader, final int mask, final List chunkPages) + throws IOException { + // A time chunk header carries VECTOR, and its pages use TimeStatistics rather than the + // LongStatistics that Statistics.getStatsByType(INT64) would create. + final TSDataType statisticsType = sourceHeader.getDataType(); + Statistics statistics = null; + int totalSize = 0; + int pageCount = 0; + int emptyPageCount = 0; + for (final PageBuffer page : chunkPages) { + if (page.needDecode) { + // A decoded page is kept as points instead of bytes, so it cannot be copied as it is. + return null; + } + if (page.pageHeader == null) { + return null; + } + if (page.pageData.remaining() == 0) { + // An empty page (an aligned series stores an all-null column as an empty page) has no data + // and no statistics, but it still occupies one page slot. The aligned reader pairs the + // pages of a value chunk with the pages of its time chunk by position, so dropping it would + // shift every later page of the value chunk onto the wrong timestamps. It is kept as the + // single 0 varint it serializes to, unless the whole chunk is empty: an all-null chunk + // keeps the zero-page layout instead. + emptyPageCount++; + continue; + } + // A non-empty page without statistics belongs to a single-page chunk: the statistics of such + // a chunk live in the chunk metadata, not in the page header, and cannot be rebuilt from the + // page bytes. Such a chunk cannot be assembled here and has to go through the writer path. + if (page.pageHeader.getStatistics() == null) { + return null; + } + if (statistics == null) { + statistics = Statistics.getStatsByType(statisticsType); + } + statistics.mergeStatistics(page.pageHeader.getStatistics()); + try { + totalSize = Math.addExact(totalSize, page.pageHeader.getSerializedPageSize()); + } catch (final ArithmeticException e) { + // A chunk that does not fit in one buffer cannot be assembled here: re-encode it instead. + return null; + } + pageCount++; + } + if (statistics == null) { + return null; + } + pageCount += emptyPageCount; + try { + // Each kept empty page is exactly the one byte of the 0 varint that represents it. + totalSize = Math.addExact(totalSize, emptyPageCount); + } catch (final ArithmeticException e) { + return null; + } + + final ByteBuffer encodedData = ByteBuffer.allocate(totalSize); + final OutputStream headerStream = new ByteBufferOutputStream(encodedData); + try { + for (final PageBuffer page : chunkPages) { + // The pages are copied as they were read, so each page header has to serialize back to + // exactly the bytes getSerializedPageSize() accounted for. Verify that instead of trusting + // it: a header that serializes longer would overflow the buffer, and a shorter one would + // leave a chunk whose declared data size does not describe its pages. + final int expectedPageSize = page.pageHeader.getSerializedPageSize(); + final int pageDataSize = page.pageData.remaining(); + if (expectedPageSize < pageDataSize || expectedPageSize > encodedData.remaining()) { + return null; + } + final int pageStart = encodedData.position(); + page.pageHeader.serializeTo(headerStream); + if (encodedData.position() - pageStart != expectedPageSize - pageDataSize + || pageDataSize > encodedData.remaining()) { + return null; + } + encodedData.put(page.pageData.duplicate()); + } + } catch (final BufferOverflowException e) { + // Defensive only: the checks above already keep every write within the allocated size. + return null; + } + if (encodedData.position() != totalSize) { + // Never hand out a chunk whose header advertises more bytes than it carries. + return null; + } + encodedData.flip(); + // A chunk with one page is normally written with ONLY_ONE_PAGE_CHUNK_HEADER, whose page + // header does not carry statistics. The pages copied here keep their statistics, so use the + // multi-page marker to make them readable. + final ChunkHeader encodedHeader = + pageCount == 1 + ? new ChunkHeader( + (byte) (MetaMarker.CHUNK_HEADER | mask), + sourceHeader.getMeasurementID(), + totalSize, + sourceHeader.getDataType(), + sourceHeader.getCompressionType(), + sourceHeader.getEncodingType()) + : createEncodedChunkHeader(sourceHeader, totalSize, pageCount, mask); + return new Chunk(encodedHeader, encodedData, null, statistics); + } + + private static final class ByteBufferOutputStream extends OutputStream { + private final ByteBuffer buffer; + + private ByteBufferOutputStream(final ByteBuffer buffer) { + this.buffer = buffer; + } + + @Override + public void write(final int b) { + buffer.put((byte) b); + } + + @Override + public void write(final byte[] bytes, final int offset, final int length) { + buffer.put(bytes, offset, length); + } + } + protected static ChunkHeader createEncodedChunkHeader( final ChunkHeader sourceHeader, final int encodedDataSize, @@ -396,13 +560,23 @@ private AlignedChunkWriterImpl createAlignedChunkWriter() { return new AlignedChunkWriterImpl(timeSchema, valueSchemaList, encryptParameter); } - private void writePageBuffersToWriter(final AlignedChunkWriterImpl writer) + private void writePageBuffersToWriter( + final AlignedChunkWriterImpl writer, final Set directlyBuiltChunkIndexes) throws IOException, PageException { + // Cache partition bounds: resolving them per point performs BigInteger work and allocations. + final long partitionStart = timePartitionSlot.getStartTime(); + final long partitionEnd = TimePartitionUtils.getTimePartitionEndTime(partitionStart); for (final PageBuffer page : pageBuffers) { + if (directlyBuiltChunkIndexes.contains(page.chunkIndex)) { + continue; + } if (!page.needDecode) { - // Skip empty pages: an aligned series stores an all-null column as an empty page, and - // writing its page header alone yields a chunk that cannot be parsed back. - if (page.pageData.remaining() == 0) { + // An empty page (an aligned series stores an all-null column as an empty page) has to be + // written whenever its chunk also has a page with data: the aligned reader pairs the pages + // of a value chunk with the pages of its time chunk by position, so dropping the page would + // shift every later page of the value chunk onto the wrong timestamps. A chunk that has no + // page with data at all keeps the zero-page layout of an all-null column instead. + if (page.pageData.remaining() == 0 && !hasPageWithData(page.chunkIndex)) { continue; } if (page.chunkIndex == 0) { @@ -413,43 +587,73 @@ private void writePageBuffersToWriter(final AlignedChunkWriterImpl writer) } continue; } + // Timestamps within a page are strictly increasing: a chunk group writer rejects an + // out-of-order point (see Aligned/NonAlignedChunkGroupWriterImpl#checkIsHistoryData), so the + // points of this partition are always a contiguous slice and the upper bound can start at the + // lower bound. + final int from = lowerBound(page.timeBatch, 0, page.timeBatch.length, partitionStart); + final int to = upperBound(page.timeBatch, from, page.timeBatch.length, partitionEnd); + if (from >= to) { + continue; + } if (page.valueBatch == null) { - boolean hasPoint = false; - for (final long time : page.timeBatch) { - if (!isInCurrentTimePartition(time)) { - continue; - } - writer.writeTime(time); - hasPoint = true; - } - if (hasPoint) { - writer.sealCurrentTimePage(); + for (int i = from; i < to; i++) { + writer.writeTime(page.timeBatch[i]); } + writer.sealCurrentTimePage(); } else { final TSDataType dataType = chunkHeaderList.get(page.chunkIndex).getDataType(); - boolean hasPoint = false; - for (int i = 0; i < page.timeBatch.length; i++) { - if (!isInCurrentTimePartition(page.timeBatch[i])) { - continue; - } + for (int i = from; i < to; i++) { 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(); - // Use the overflow-safe end time. A naive "start + interval - 1" wraps around for the last - // representable partition (and equals the start when interval == 1), which would either drop - // the upper bound or exclude points that belong to the current partition. - final long partitionEnd = TimePartitionUtils.getTimePartitionEndTime(partitionStart); - return time >= partitionStart && time <= partitionEnd; + /** + * Whether the chunk written here carries at least one point. The empty pages of such a chunk are + * written as well, so that the page sequence of every value chunk stays aligned with the page + * sequence of the time chunk; a chunk without any point keeps the zero-page layout of an all-null + * column. Only empty pages ask this question, so the scan stays off the common path. + */ + private boolean hasPageWithData(final int chunkIndex) { + for (final PageBuffer page : pageBuffers) { + if (page.chunkIndex == chunkIndex + && (page.needDecode ? page.satisfiedLength > 0 : page.pageData.remaining() > 0)) { + return true; + } + } + return false; + } + + private static int lowerBound(final long[] times, final int low, final int high, final long key) { + int left = low; + int right = high; + while (left < right) { + final int mid = (left + right) >>> 1; + if (times[mid] < key) { + left = mid + 1; + } else { + right = mid; + } + } + return left; + } + + private static int upperBound(final long[] times, final int low, final int high, final long key) { + int left = low; + int right = high; + while (left < right) { + final int mid = (left + right) >>> 1; + if (times[mid] <= key) { + left = mid + 1; + } else { + right = mid; + } + } + return left; } private void writeValueToAlignedChunkWriter( 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 9db2c9590b2b..56bff063d705 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 @@ -35,6 +35,11 @@ import java.io.IOException; import java.io.Serializable; import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; import static org.apache.tsfile.common.constant.TsFileConstant.VALUE_COLUMN_MASK; @@ -69,9 +74,29 @@ protected void encodeAndBuildChunks() throws IOException, PageException { // partition and equals the start when interval == 1. final long endTime = TimePartitionUtils.getTimePartitionEndTime(startTime); + // Group the buffered pages by chunk once; scanning the whole page list once per chunk turns + // batched aligned tables into O(columns * pages). + final Map> pagesByChunkIndex = new LinkedHashMap<>(); + for (final PageBuffer pageBuffer : pageBuffers) { + pagesByChunkIndex + .computeIfAbsent(pageBuffer.chunkIndex, ignored -> new ArrayList<>()) + .add(pageBuffer); + } + // Process each standalone value chunk for (int i = 0; i < chunkHeaderList.size(); i++) { final ChunkHeader chunkHeader = chunkHeaderList.get(i); + final List chunkPages = + pagesByChunkIndex.getOrDefault(i, Collections.emptyList()); + final Chunk directlyBuiltChunk = + buildChunkFromUndecodedPages(chunkHeader, VALUE_COLUMN_MASK, chunkPages); + if (directlyBuiltChunk != null) { + entireChunks.add(directlyBuiltChunk); + dataSize += + directlyBuiltChunk.getData().remaining() + + directlyBuiltChunk.getChunkStatistic().getSerializedSize(); + continue; + } MeasurementSchema measurementSchema = new MeasurementSchema( chunkHeader.getMeasurementID(), @@ -93,11 +118,7 @@ protected void encodeAndBuildChunks() throws IOException, PageException { // untouched writer would produce a zero-point page whose data size does not match the bytes // actually written, and the resulting TsFile could not be parsed back. boolean hasBufferedData = false; - for (PageBuffer page : pageBuffers) { - if (page.chunkIndex != i) { - continue; - } - + for (PageBuffer page : chunkPages) { if (!page.needDecode) { // An undecoded page may carry no data at all (for aligned series an all-null column is // stored as an empty page). Writing its page header alone would produce a chunk whose 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 c35e41897146..7deb07b9bd79 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 @@ -44,7 +44,14 @@ public interface ChunkData extends TsFileData { ByteBuffer EMPTY_BUFFER = ByteBuffer.allocate(0); - byte[] EMPTY_BYTES = new byte[0]; + + /** + * An input stream over a {@link ByteBuffer} that can expose the next bytes as a slice instead of + * copying them into a new array. + */ + interface ByteBufferSliceSource { + ByteBuffer readSlice(int length); + } // ------------------------------------------------------------------------- // Metadata & Status @@ -168,9 +175,28 @@ static void serializeChunkPayload( // 2. Write inline payload bytes ReadWriteIOUtils.write(false, stream); // isRef = false - byte[] payload = extractPayloadBytes(data, ref); - ReadWriteIOUtils.write(payload.length, stream); - stream.write(payload); + final int payloadLength = (int) inlinePayloadSize(data, ref); + ReadWriteIOUtils.write(payloadLength, stream); + if (data != null && data.hasRemaining()) { + writeByteBuffer(stream, data); + } else if (ref != null) { + ref.writePayloadTo(stream); + } + } + + private static void writeByteBuffer(final DataOutputStream stream, final ByteBuffer data) + throws IOException { + if (data.hasArray() && !data.isReadOnly()) { + stream.write(data.array(), data.arrayOffset() + data.position(), data.remaining()); + return; + } + final ByteBuffer duplicate = data.duplicate(); + final byte[] buffer = new byte[Math.min(duplicate.remaining(), 8192)]; + while (duplicate.hasRemaining()) { + final int length = Math.min(duplicate.remaining(), buffer.length); + duplicate.get(buffer, 0, length); + stream.write(buffer, 0, length); + } } /** Deserializes payload, appending to {@code refs} if referenced externally. */ @@ -193,18 +219,13 @@ static ByteBuffer deserializeChunkPayload(InputStream stream, List file.length()) { - throw new ChunkPayloadUnavailableException( - String.format( - StorageEngineMessages - .EXCEPTION_THE_CHUNK_PAYLOAD_OF_ARG_BYTES_AT_OFFSET_ARG_IS_LONGER_THAN_THE_STAGED_FILE_ARG_14081256, - size, - offset, - file.getAbsolutePath())); + if (size == 0) { + // Nothing to read: an empty payload must not require the staged file to still exist. + return new byte[0]; } + final File file = locatePayloadFile(); try (final RandomAccessFile input = new RandomAccessFile(file, "r")) { - final byte[] payload = new byte[(int) size]; + final byte[] payload = new byte[Math.toIntExact(size)]; input.seek(offset); - input.readFully(payload); - return payload; - } catch (final EOFException e) { + try { + input.readFully(payload); + return payload; + } catch (final EOFException e) { + throw newStagedFileEofException(file, input.getFilePointer(), e); + } + } catch (final IOException | ArithmeticException e) { throw new ChunkPayloadUnavailableException( String.format( - StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_EOF_8743387D, + StorageEngineMessages + .EXCEPTION_FAILED_TO_READ_BACK_THE_STAGED_PIECE_FILE_ARG_ARG_3F54CB90, file.getAbsolutePath(), - offset), + e.getMessage()), e); - } catch (final IOException | ArithmeticException | NegativeArraySizeException e) { + } + } + + /** + * Streams the referenced bytes into {@code stream} without materializing them in a byte array. + * Used when a payload has to be inlined into a request but is still stored in the staged file. + */ + public void writePayloadTo(final DataOutputStream stream) throws IOException { + if (size == 0) { + // Nothing to copy: an empty payload must not require the staged file to still exist. + return; + } + final File file = locatePayloadFile(); + final byte[] buffer = new byte[(int) Math.min(size, 8192L)]; + try (final RandomAccessFile input = new RandomAccessFile(file, "r")) { + input.seek(offset); + long remaining = size; + long position = offset; + while (remaining > 0) { + final int read = input.read(buffer, 0, (int) Math.min(buffer.length, remaining)); + if (read < 0) { + // The file shrank after it was checked: report the position the file actually ends at. + throw newStagedFileEofException(file, position, null); + } + stream.write(buffer, 0, read); + remaining -= read; + position += read; + } + } catch (final IOException | ArithmeticException e) { throw new ChunkPayloadUnavailableException( String.format( StorageEngineMessages @@ -148,6 +173,37 @@ public byte[] readPayload() { } } + private ChunkPayloadUnavailableException newStagedFileEofException( + final File file, final long position, final Throwable cause) { + final String message = + String.format( + StorageEngineMessages + .EXCEPTION_UNEXPECTED_END_OF_FILE_WHEN_READING_THE_STAGED_PIECE_ARG_ARG_BYTES_WERE_REQUESTED_FROM_OFFSET_ARG_BUT_THE_FILE_ENDED_AT_POSITION_ARG_3F9CA32B, + file.getAbsolutePath(), + size, + offset, + position); + return cause == null + ? new ChunkPayloadUnavailableException(message) + : new ChunkPayloadUnavailableException(message, cause); + } + + private File locatePayloadFile() { + final File file = locateFile(); + // Check against the file before allocating: the subtraction form cannot overflow even for an + // offset or size close to Long.MAX_VALUE. + if (offset < 0 || size < 0 || size > file.length() - offset) { + throw new ChunkPayloadUnavailableException( + String.format( + StorageEngineMessages + .EXCEPTION_THE_CHUNK_PAYLOAD_OF_ARG_BYTES_AT_OFFSET_ARG_IS_LONGER_THAN_THE_STAGED_FILE_ARG_14081256, + size, + offset, + file.getAbsolutePath())); + } + return file; + } + /** * Resolves the staged file this reference points at. * 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 9b17746f942a..7a976838d705 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 @@ -472,30 +472,37 @@ private void consumeAllAlignedChunkData( return; } - // Insertion-ordered on purpose: the chunks of a series must keep the page order of the - // source TsFile, so that the emitted pieces (and the replicas applying them) write the time - // partitions in ascending order. - Map chunkDataMap = new LinkedHashMap<>(); - for (Map.Entry> entry : pageIndex2ChunkData.entrySet()) { - List alignedChunkDataList = entry.getValue(); + // Insertion-ordered and de-duplicating on purpose: the chunks of a series must keep the page + // order of the source TsFile, so that the emitted pieces (and the replicas applying them) write + // the time partitions in ascending order, and a chunk filled by more than one page index is + // only written once. + // + // Every entry is swapped for a fresh wrapper while the object that holds the pages is the one + // handed to the consumer: a consumed chunk data must not be drained again by a later flush, so + // the map keeps objects that start empty instead of the ones that were just emitted. + final Map chunkDataMap = new LinkedHashMap<>(); + for (final Map.Entry> entry : pageIndex2ChunkData.entrySet()) { + final List alignedChunkDataList = entry.getValue(); for (int i = 0; i < alignedChunkDataList.size(); i++) { - AlignedChunkData oldChunkData = alignedChunkDataList.get(i); - BatchedAlignedValueChunkData chunkData = + final AlignedChunkData oldChunkData = alignedChunkDataList.get(i); + final BatchedAlignedValueChunkData chunkData = chunkDataMap.computeIfAbsent(oldChunkData, BatchedAlignedValueChunkData::new); alignedChunkDataList.set(i, chunkData); } } - for (AlignedChunkData chunkData : chunkDataMap.keySet()) { - chunkData.endChunk(); + for (final AlignedChunkData chunkData : chunkDataMap.keySet()) { timePartitionSlots.add(chunkData.getTimePartitionSlot()); if (deletions.isEmpty() && timePartitionSlots.size() > CONFIG.getLoadTsFileSpiltPartitionMaxSize()) { + // Check before encoding: the chunk that crosses the limit does not have to be turned into + // standard chunks only to be rejected. throw new LoadFileException( String.format( StorageEngineMessages .STORAGE_EXCEPTION_TIME_PARTITION_SLOTS_SIZE_IS_GREATER_THAN_S_D076F78E, CONFIG.getLoadTsFileSpiltPartitionMaxSize())); } + chunkData.endChunk(); if (Boolean.FALSE.equals(consumer.apply(chunkData))) { throw new IllegalStateException( String.format( @@ -555,7 +562,8 @@ private Pair decodePage( if (isAligned) { TimePageReader timePageReader = new TimePageReader(pageHeader, pageData, timeDecoder); long[] times = timePageReader.getNextTimeBatch(); - return new Pair<>(times, new Object[times.length]); + // Aligned time pages carry no value array; return null to avoid a dummy Object[]. + return new Pair<>(times, null); } valueDecoder.reset(); @@ -596,19 +604,6 @@ private void handleEmptyValueChunk( } } - /** - * handle empty page in aligned chunk, if uncompressedSize and compressedSize are both 0, and the - * statistics is null, then the page is empty. - * - * @param pageHeader page header - * @return true if the page is empty - */ - private boolean isEmptyPage(PageHeader pageHeader) { - return pageHeader.getUncompressedSize() == 0 - && pageHeader.getCompressedSize() == 0 - && pageHeader.getStatistics() == null; - } - private TsPrimitiveType[] decodeValuePage( TsFileSequenceReader reader, ChunkHeader chunkHeader, @@ -616,7 +611,9 @@ private TsPrimitiveType[] decodeValuePage( long[] times, Decoder valueDecoder) throws IOException { - if (pageHeader.getSerializedPageSize() == 0) { + // An empty page stands for an all-null slice of the column. It has to expand to one null per + // time of the page, otherwise the decoded values would not line up with the decoded times. + if (pageHeader.getUncompressedSize() == 0) { return new TsPrimitiveType[times.length]; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/BufferedFileChannelOutput.java b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/BufferedFileChannelOutput.java new file mode 100644 index 000000000000..946b99e5c70d --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/BufferedFileChannelOutput.java @@ -0,0 +1,192 @@ +/* + * 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.tsfile.write; + +import org.apache.iotdb.db.i18n.DataNodeMiscMessages; +import org.apache.iotdb.db.i18n.StorageEngineMessages; + +import org.apache.tsfile.write.writer.TsFileOutput; + +import java.io.IOException; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.util.Objects; + +/** + * A {@link TsFileOutput} that buffers writes and hands them to the {@link FileChannel} once the + * buffer size the caller passes in is reached. The buffer is allocated at its full size when the + * output is created, and it is flushed by {@link #force()}, {@link #flush()}, {@link + * #truncate(long)}, a seek to another position, or {@link #close()}. + * + *

This class is not thread-safe; callers must serialize access to one instance. + */ +public final class BufferedFileChannelOutput extends OutputStream implements TsFileOutput { + + private final FileChannel channel; + private final int bufferSize; + private final byte[] buffer; + + /** Reused wrapper over {@link #buffer} so each flush does not allocate a new ByteBuffer. */ + private final ByteBuffer bufferView; + + private int bufferLength; + + /** + * Creates an output that buffers the given number of bytes. + * + * @param bufferSize the size of the buffer in bytes, must be positive + */ + public BufferedFileChannelOutput(final FileChannel channel, final int bufferSize) { + this.channel = + Objects.requireNonNull( + channel, StorageEngineMessages.EXCEPTION_CHANNEL_CANNOT_BE_NULL_0F79E0FB); + if (bufferSize <= 0) { + throw new IllegalArgumentException( + DataNodeMiscMessages.EXCEPTION_BUFFER_SIZE_SHOULD_BE_GREATER_THAN_0_BB68B5C4); + } + this.bufferSize = bufferSize; + this.buffer = new byte[bufferSize]; + this.bufferView = ByteBuffer.wrap(buffer); + } + + @Override + public void write(final int b) throws IOException { + write((byte) b); + } + + @Override + public void write(final byte b) throws IOException { + if (bufferLength == bufferSize) { + flushBuffer(); + } + buffer[bufferLength++] = b; + } + + @Override + public void write(final byte[] bytes) throws IOException { + write(bytes, 0, bytes.length); + } + + @Override + public void write(final byte[] bytes, final int offset, final int length) throws IOException { + Objects.checkFromIndexSize(offset, length, bytes.length); + if (length == 0) { + return; + } + if (length >= bufferSize) { + // A block at least as large as the buffer is not worth a copy. + flushBuffer(); + writeDirect(bytes, offset, length); + return; + } + if (length > bufferSize - bufferLength) { + flushBuffer(); + } + System.arraycopy(bytes, offset, buffer, bufferLength, length); + bufferLength += length; + } + + @Override + public void write(final ByteBuffer src) throws IOException { + final int length = src.remaining(); + if (length == 0) { + return; + } + if (length >= bufferSize) { + flushBuffer(); + writeFully(src); + return; + } + if (length > bufferSize - bufferLength) { + flushBuffer(); + } + src.get(buffer, bufferLength, length); + bufferLength += length; + } + + private void writeDirect(final byte[] bytes, final int offset, final int length) + throws IOException { + writeFully(ByteBuffer.wrap(bytes, offset, length)); + } + + /** Hands the buffered bytes to the channel, keeping the previous writes in order. */ + private void flushBuffer() throws IOException { + if (bufferLength == 0) { + return; + } + bufferView.position(0).limit(bufferLength); + writeFully(bufferView); + bufferLength = 0; + } + + private void writeFully(final ByteBuffer src) throws IOException { + while (src.hasRemaining()) { + channel.write(src); + } + } + + @Override + public long getPosition() throws IOException { + // The buffered bytes are already logically written, so a position query must not flush them. + return channel.position() + bufferLength; + } + + /** Moves the output to the given absolute position, flushing first when the position changes. */ + void position(final long offset) throws IOException { + if (offset == channel.position() + bufferLength) { + return; + } + flushBuffer(); + channel.position(offset); + } + + @Override + public OutputStream wrapAsStream() { + return this; + } + + @Override + public void flush() throws IOException { + flushBuffer(); + } + + @Override + public void force() throws IOException { + flushBuffer(); + // Pass false to fsync file content without blocking on OS metadata modifications + channel.force(false); + } + + @Override + public void truncate(final long size) throws IOException { + flushBuffer(); + channel.truncate(size); + } + + @Override + public void close() throws IOException { + try { + flushBuffer(); + } finally { + channel.close(); + } + } +} 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 7e375d9fe704..8ceee739d323 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 @@ -19,6 +19,7 @@ package org.apache.tsfile.write; +import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.StorageEngineMessages; import org.apache.tsfile.common.conf.TSFileConfig; @@ -87,16 +88,25 @@ public class TsFilePrecalculatedChunkWriter implements AutoCloseable { private long maxPhysicalFileSize; private boolean sealed; + /** + * The buffer size of the staged file output, from {@code load_tsfile_write_buffer_size_in_byte}. + * Only the caller knows the configuration; the output itself stays independent of it. + */ + private static int loadTsFileWriteBufferSize() { + return IoTDBDescriptor.getInstance().getConfig().getLoadTsFileWriteBufferSize(); + } + public TsFilePrecalculatedChunkWriter(final File file) throws IOException { this.file = Objects.requireNonNull(file, StorageEngineMessages.EXCEPTION_FILE_CANNOT_BE_NULL_29A83D70); this.out = - new WritableFileChannelOutput( + new BufferedFileChannelOutput( FileChannel.open( file.toPath(), StandardOpenOption.CREATE, StandardOpenOption.WRITE, - StandardOpenOption.TRUNCATE_EXISTING)); + StandardOpenOption.TRUNCATE_EXISTING), + loadTsFileWriteBufferSize()); this.sealed = false; startFile(); this.maxPhysicalFileSize = out.getPosition(); @@ -117,9 +127,10 @@ public TsFilePrecalculatedChunkWriter( final File file, final FileChannel channel, final boolean sealed) throws IOException { this.file = file; this.out = - new WritableFileChannelOutput( + new BufferedFileChannelOutput( Objects.requireNonNull( - channel, StorageEngineMessages.EXCEPTION_CHANNEL_CANNOT_BE_NULL_0F79E0FB)); + channel, StorageEngineMessages.EXCEPTION_CHANNEL_CANNOT_BE_NULL_0F79E0FB), + loadTsFileWriteBufferSize()); this.sealed = sealed; this.maxPhysicalFileSize = channel.size(); LOGGER.info( @@ -215,7 +226,7 @@ public ChunkWriteResult writeChunk( */ private long alignToOffset(final long expectedOffset, final IDeviceID device, final Chunk chunk) throws IOException { - if (out instanceof WritableFileChannelOutput channelOutput) { + if (out instanceof BufferedFileChannelOutput channelOutput) { // 1. Expand file with zeros if jumping past current physical maximum bound if (expectedOffset > maxPhysicalFileSize) { final long paddingStartOffset = maxPhysicalFileSize; @@ -349,7 +360,7 @@ public void close() throws IOException { try { // Seek past the highest written data zone byte before building indices - if (out instanceof WritableFileChannelOutput channelOutput) { + if (out instanceof BufferedFileChannelOutput channelOutput) { channelOutput.position(maxPhysicalFileSize); } @@ -472,88 +483,4 @@ private String getFileForLog() { public record ChunkWriteResult( long actualChunkGroupHeaderOffset, long actualChunkOffset, long actualChunkEndOffset) {} - - // ------------------------------------------------------------------------- - // Channel Wrapper - // ------------------------------------------------------------------------- - - private static final class WritableFileChannelOutput implements TsFileOutput { - - private final FileChannel channel; - private final ByteBuffer singleByteBuffer = ByteBuffer.allocateDirect(1); - - private WritableFileChannelOutput(final FileChannel channel) { - this.channel = channel; - } - - @Override - public void write(final byte[] bytes) throws IOException { - write(bytes, 0, bytes.length); - } - - public void write(final byte[] bytes, final int offset, final int length) throws IOException { - write(ByteBuffer.wrap(bytes, offset, length)); - } - - @Override - public synchronized void write(final byte b) throws IOException { - singleByteBuffer.clear(); - singleByteBuffer.put(b); - singleByteBuffer.flip(); - write(singleByteBuffer); - } - - @Override - public void write(final ByteBuffer buffer) throws IOException { - while (buffer.hasRemaining()) { - channel.write(buffer); - } - } - - @Override - public long getPosition() throws IOException { - return channel.position(); - } - - private void position(final long offset) throws IOException { - channel.position(offset); - } - - @Override - public OutputStream wrapAsStream() { - return new OutputStream() { - @Override - public void write(final int b) throws IOException { - WritableFileChannelOutput.this.write((byte) b); - } - - @Override - public void write(final byte[] bytes, final int offset, final int length) - throws IOException { - WritableFileChannelOutput.this.write(bytes, offset, length); - } - }; - } - - @Override - public void flush() throws IOException { - force(); - } - - @Override - public void force() throws IOException { - // Pass false to fsync file content without blocking on OS metadata modifications - channel.force(false); - } - - @Override - public void truncate(final long size) throws IOException { - channel.truncate(size); - } - - @Override - public void close() throws IOException { - channel.close(); - } - } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/PropertiesTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/PropertiesTest.java index b0d7aef7394d..0fe910f1d6b9 100755 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/PropertiesTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/PropertiesTest.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.conf; +import org.apache.iotdb.calc.exception.QueryProcessException; import org.apache.iotdb.commons.conf.CommonConfig; import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.conf.ConfigurationFileUtils; @@ -69,6 +70,39 @@ public void testWalFileListCacheConfiguration() throws Exception { } } + /** The LOAD write buffer defaults to 1 MiB, can be overridden, and must stay positive. */ + @Test + public void testLoadTsFileWriteBufferSizeDefaultsAndReloads() throws Exception { + final IoTDBDescriptor descriptor = IoTDBDescriptor.getInstance(); + final int originalValue = descriptor.getConfig().getLoadTsFileWriteBufferSize(); + + try { + Assert.assertEquals(1024 * 1024, new IoTDBConfig().getLoadTsFileWriteBufferSize()); + + final TrimProperties properties = new TrimProperties(); + properties.setProperty("load_tsfile_write_buffer_size_in_byte", "4096"); + descriptor.loadProperties(properties); + Assert.assertEquals(4096, descriptor.getConfig().getLoadTsFileWriteBufferSize()); + + properties.setProperty("load_tsfile_write_buffer_size_in_byte", "8192"); + descriptor.loadHotModifiedProps(properties); + Assert.assertEquals(8192, descriptor.getConfig().getLoadTsFileWriteBufferSize()); + + properties.setProperty("load_tsfile_write_buffer_size_in_byte", "0"); + Assert.assertThrows( + QueryProcessException.class, () -> descriptor.loadHotModifiedProps(properties)); + Assert.assertEquals( + "an invalid value keeps the current one", + 8192, + descriptor.getConfig().getLoadTsFileWriteBufferSize()); + } finally { + final TrimProperties properties = new TrimProperties(); + properties.setProperty( + "load_tsfile_write_buffer_size_in_byte", Integer.toString(originalValue)); + descriptor.loadHotModifiedProps(properties); + } + } + @Test public void testDeviceEntryBatchSizeIsCappedByThriftFrameSizeOnStartup() throws Exception { final IoTDBDescriptor descriptor = IoTDBDescriptor.getInstance(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategyMigrationTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategyMigrationTest.java new file mode 100644 index 000000000000..257e0fcf01d9 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategyMigrationTest.java @@ -0,0 +1,183 @@ +/* + * 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.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.queryengine.plan.planner.plan.node.PlanNodeId; +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.memory.LoadTsFileDataCacheMemoryBlock; +import org.apache.iotdb.rpc.RpcUtils; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Test; + +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.function.BiPredicate; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * The LOAD keeps dispatching pieces while a region of the file migrates, so the pieces dispatched + * after the migration must follow the route the region has now instead of the one pinned before it. + */ +public class TwoPhaseConsensusLoadStrategyMigrationTest { + + private static final TConsensusGroupId REGION = + new TConsensusGroupId(TConsensusGroupType.DataRegion, 1); + + private static final TRegionReplicaSet ROUTE_BEFORE_MIGRATION = + new TRegionReplicaSet(REGION, Collections.singletonList(location(11))); + + private static final TRegionReplicaSet ROUTE_AFTER_MIGRATION = + new TRegionReplicaSet(REGION, Collections.singletonList(location(12))); + + @Test + public void testAPieceDispatchedAfterTheMigrationFollowsTheNewRoute() throws Exception { + final List submitted = new ArrayList<>(); + final TwoPhaseConsensusLoadStrategy strategy = + strategy( + submitter( + submitted, + (pieceIndex, route) -> pieceIndex == 1 && route.equals(ROUTE_BEFORE_MIGRATION), + ROUTE_AFTER_MIGRATION)); + + assertTrue(dispatch(strategy, piece(0), ROUTE_BEFORE_MIGRATION)); + // The region migrated while the LOAD was between two pieces: the second piece still goes to the + // pinned route first, fails there, and is retried on the route the region has now. + assertTrue(dispatch(strategy, piece(1), ROUTE_AFTER_MIGRATION)); + + assertEquals(Arrays.asList("0@11", "1@11", "1@12"), submitted); + assertEquals(ROUTE_AFTER_MIGRATION, routesOf(strategy).get(REGION)); + assertEquals(Long.valueOf(2L), pieceCounts(strategy).get(REGION)); + } + + @Test + public void testAPieceWhoseRouteIsGoneAfterTheMigrationIsNotDispatched() throws Exception { + final List submitted = new ArrayList<>(); + final TwoPhaseConsensusLoadStrategy strategy = + strategy( + submitter( + submitted, + (pieceIndex, route) -> pieceIndex == 1 && route.equals(ROUTE_BEFORE_MIGRATION), + null)); + + assertTrue(dispatch(strategy, piece(0), ROUTE_BEFORE_MIGRATION)); + // The route the piece was pinned to is gone and cannot be looked up again: the piece fails + // instead of being sent to a route the partition table no longer knows. + assertFalse(dispatch(strategy, piece(1), ROUTE_AFTER_MIGRATION)); + + assertEquals(Arrays.asList("0@11", "1@11"), submitted); + assertEquals(Long.valueOf(1L), pieceCounts(strategy).get(REGION)); + } + + private static TwoPhaseConsensusLoadStrategy strategy(final LoadConsensusSubmitter submitter) { + return new TwoPhaseConsensusLoadStrategy( + mock(LoadTsFileDispatcherImpl.class), + new DataPartitionResolver(mock(DataPartitionBatchFetcher.class), "root"), + mock(LoadTsFileDataCacheMemoryBlock.class), + submitter, + false); + } + + /** + * Submits on the given route, fails when the predicate says the route is stale, and resolves the + * route to the given one afterwards. + */ + private static LoadConsensusSubmitter submitter( + final List submitted, + final BiPredicate staleRoute, + final TRegionReplicaSet resolvedRoute) { + final LoadConsensusSubmitter submitter = mock(LoadConsensusSubmitter.class); + when(submitter.submit(any(), any())) + .thenAnswer( + invocation -> { + final TRegionReplicaSet route = invocation.getArgument(0); + final LoadTsFileConsensusNode piece = invocation.getArgument(1); + submitted.add( + piece.getPieceIndex() + + "@" + + route.getDataNodeLocations().get(0).getDataNodeId()); + return staleRoute.test(piece.getPieceIndex(), route) + ? RpcUtils.getStatus(TSStatusCode.DISPATCH_ERROR) + : RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); + }); + when(submitter.resolveRoute(any(), any(Integer.class))).thenReturn(resolvedRoute); + return submitter; + } + + private static LoadTsFilePieceNode piece(final long pieceIndex) { + return new LoadTsFilePieceNode(new PlanNodeId("piece-" + pieceIndex), null, pieceIndex); + } + + private static boolean dispatch( + final TwoPhaseConsensusLoadStrategy strategy, + final LoadTsFilePieceNode piece, + final TRegionReplicaSet route) + throws Exception { + final Method method = + TwoPhaseConsensusLoadStrategy.class.getDeclaredMethod( + "dispatchOnePieceNode", LoadTsFilePieceNode.class, TRegionReplicaSet.class); + method.setAccessible(true); + return (boolean) method.invoke(strategy, piece, route); + } + + @SuppressWarnings("unchecked") + private static Map routesOf( + final TwoPhaseConsensusLoadStrategy strategy) throws Exception { + final Field field = TwoPhaseConsensusLoadStrategy.class.getDeclaredField("regionRoutes"); + field.setAccessible(true); + return (Map) field.get(strategy); + } + + @SuppressWarnings("unchecked") + private static Map pieceCounts( + final TwoPhaseConsensusLoadStrategy strategy) throws Exception { + final Field field = TwoPhaseConsensusLoadStrategy.class.getDeclaredField("regionPieceCounts"); + field.setAccessible(true); + return (Map) field.get(strategy); + } + + private static TDataNodeLocation location(final int dataNodeId) { + final TEndPoint endPoint = new TEndPoint("127.0.0." + dataNodeId, 9003); + return new TDataNodeLocation() + .setDataNodeId(dataNodeId) + .setClientRpcEndPoint(endPoint) + .setInternalEndPoint(endPoint) + .setMPPDataExchangeEndPoint(endPoint) + .setDataRegionConsensusEndPoint(endPoint) + .setSchemaRegionConsensusEndPoint(endPoint); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerSplitCoverageTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerSplitCoverageTest.java index a1c2cee8689b..7db4d160748c 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerSplitCoverageTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerSplitCoverageTest.java @@ -299,6 +299,56 @@ public void testAlignedMeasurementMatrix() throws Exception { 5_000L)); } + /** + * An aligned column whose values are all null in one page has to keep that page slot: the aligned + * reader pairs the pages of a value chunk with the pages of its time chunk by position, so a + * dropped page shifts every later value onto the wrong timestamp. Covers a page that is copied as + * it was read, a page that is split at a time partition boundary and a chunk that has to be + * re-encoded. + */ + @Test + public void testAllNullPageKeepsItsPageSlot() throws Exception { + final List copiedPages = new ArrayList<>(); + copiedPages.add(new long[][] {page(0, 9), page(10, 19), page(20, 29), page(100, 109)}); + final List splitPages = new ArrayList<>(); + splitPages.add(new long[][] {page(190, 209), page(210, 229)}); + // A page copied as it was read plus a page that has to be decoded: here the empty page is + // written + // by the re-encoding path, which has to keep it as well. + final List mixedPages = new ArrayList<>(); + mixedPages.add(new long[][] {page(0, 9), page(10, 19), page(90, 209)}); + + final List devices = new ArrayList<>(); + devices.add( + alignedDevice( + "null_page_copied", + ChunkOrder.CHUNK_MAJOR, + series("m0", TSDataType.INT64, NullMode.NONE, copiedPages), + series( + "m1", TSDataType.INT64, NullMode.NONE, copiedPages, new long[][] {page(10, 19)}))); + devices.add( + alignedDevice( + "null_page_split", + ChunkOrder.CHUNK_MAJOR, + series("m0", TSDataType.INT64, NullMode.NONE, splitPages), + series( + "m1", TSDataType.INT64, NullMode.NONE, splitPages, new long[][] {page(190, 209)}))); + devices.add( + alignedDevice( + "null_page_reencoded", + ChunkOrder.CHUNK_MAJOR, + series("m0", TSDataType.INT64, NullMode.NONE, mixedPages), + series( + "m1", TSDataType.INT64, NullMode.NONE, mixedPages, new long[][] {page(10, 19)}))); + runScenario( + new Scenario( + "all-null-page", + FINE_PARTITION_INTERVAL, + devices, + new DispatchPlan(2, false, false), + 400L)); + } + /** Non-aligned devices with 1..5 measurements, both chunk layouts and null-heavy series. */ @Test public void testNonAlignedMeasurementMatrix() throws Exception { @@ -1204,7 +1254,7 @@ private static List expectedPoints(final SeriesSpec series) { for (int chunkIndex = 0; chunkIndex < series.chunks.size(); chunkIndex++) { for (final long[] pageRange : series.chunks.get(chunkIndex)) { for (long time = pageRange[0]; time <= pageRange[1]; time++) { - if (series.nullMode.isNull(time, chunkIndex)) { + if (series.isNull(time, chunkIndex)) { continue; } expected.add(new Object[] {time, valueFor(series.dataType, time)}); @@ -1293,7 +1343,7 @@ private static void writeNonAlignedChunk( for (final long[] pageRange : series.chunks.get(chunkIndex)) { boolean wroteThisPage = false; for (long time = pageRange[0]; time <= pageRange[1]; time++) { - if (series.nullMode.isNull(time, chunkIndex)) { + if (series.isNull(time, chunkIndex)) { continue; } writeNonAlignedValue(chunkWriter, series.dataType, time, valueFor(series.dataType, time)); @@ -1335,7 +1385,7 @@ private static void writeAlignedChunkPages( series.dataType, time, valueFor(series.dataType, time), - series.nullMode.isNull(time, chunkIndex)); + series.isNull(time, chunkIndex)); } alignedWriter.write(time); } @@ -1639,10 +1689,15 @@ boolean isNull(final long time, final int chunkIndex) { } private static final class SeriesSpec { + static final long[][] NO_NULL_RANGES = new long[0][]; + final String measurement; final TSDataType dataType; final NullMode nullMode; + /** Inclusive time ranges that are null on top of {@link #nullMode}. */ + final long[][] nullRanges; + /** Outer list = chunks; inner list = pages; each page = [startTime, endTime] inclusive. */ final List chunks; @@ -1651,10 +1706,32 @@ private static final class SeriesSpec { final TSDataType dataType, final NullMode nullMode, final List chunks) { + this(measurement, dataType, nullMode, chunks, NO_NULL_RANGES); + } + + SeriesSpec( + final String measurement, + final TSDataType dataType, + final NullMode nullMode, + final List chunks, + final long[][] nullRanges) { this.measurement = measurement; this.dataType = dataType; this.nullMode = nullMode; this.chunks = chunks; + this.nullRanges = nullRanges; + } + + boolean isNull(final long time, final int chunkIndex) { + if (nullMode.isNull(time, chunkIndex)) { + return true; + } + for (final long[] range : nullRanges) { + if (time >= range[0] && time <= range[1]) { + return true; + } + } + return false; } } @@ -1745,6 +1822,16 @@ private static SeriesSpec series( return new SeriesSpec(measurement, dataType, nullMode, chunks); } + /** A series whose points inside {@code nullRanges} are written as nulls. */ + private static SeriesSpec series( + final String measurement, + final TSDataType dataType, + final NullMode nullMode, + final List chunks, + final long[][] nullRanges) { + return new SeriesSpec(measurement, dataType, nullMode, chunks, nullRanges); + } + private static DeviceSpec alignedDevice( final String name, final ChunkOrder chunkOrder, final SeriesSpec... series) { return new DeviceSpec(deviceId(name), true, chunkOrder, Arrays.asList(series)); 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 index a3d247331676..9c0060c5fa9f 100644 --- 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 @@ -417,6 +417,42 @@ public void testChunkPayloadRefRefusesToReadBeyondTheStagedFile() throws Excepti } } + @Test + public void testChunkPayloadRefStreamsPayloadsLargerThanItsScratchBuffer() throws Exception { + // Stage a piece first so the shared staging directory follows this test's tempDir, then append + // a payload larger than the 8 KiB scratch buffer to the staged file. + final LoadTsFileManager manager = new LoadTsFileManager(dataRegion); + final LoadTsFileConsensusNode piece = stagedPiece("large-payload-load", laidOutChunk()); + manager.writePiece(piece); + final File staged = new File(tempDir, piece.getPieceRefs().get(0).getRelativePath()); + final long offset = staged.length(); + final byte[] payload = new byte[20_000]; + for (int i = 0; i < payload.length; i++) { + payload[i] = (byte) (i * 31); + } + Files.write(staged.toPath(), payload, StandardOpenOption.APPEND); + + final ChunkPayloadRef ref = + new ChunkPayloadRef(staged.getAbsolutePath(), offset, payload.length); + final ByteArrayOutputStream output = new ByteArrayOutputStream(payload.length); + ref.writePayloadTo(new DataOutputStream(output)); + + assertArrayEquals(payload, output.toByteArray()); + } + + @Test + public void testEmptyChunkPayloadRefNeedsNoStagedFile() throws Exception { + // An empty payload owns no bytes, so it can be read back even after the staged file is gone: + // resolving a missing file first would fail a forward that has nothing to copy. + final File cleanedUp = new File(tempDir, "cleaned-up.tsfile"); + final ChunkPayloadRef empty = new ChunkPayloadRef(cleanedUp.getAbsolutePath(), 0L, 0L); + + assertArrayEquals(new byte[0], empty.readPayload()); + final ByteArrayOutputStream output = new ByteArrayOutputStream(); + empty.writePayloadTo(new DataOutputStream(output)); + assertEquals(0, output.size()); + } + /** * A PREPARE of a task that is still missing a piece must fail instead of sealing the file. * diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java index 4ccbc98aec63..d70cb40e8141 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java @@ -23,10 +23,16 @@ import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.path.MeasurementPath; import org.apache.iotdb.commons.utils.TimePartitionUtils; +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.storageengine.dataregion.modification.ModificationFile; import org.apache.iotdb.db.storageengine.dataregion.modification.TreeDeletionEntry; +import com.sun.management.ThreadMXBean; import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.apache.tsfile.encoding.encoder.PlainEncoder; import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.MetaMarker; @@ -40,17 +46,29 @@ import org.apache.tsfile.file.metadata.StringArrayDeviceID; import org.apache.tsfile.file.metadata.TableSchema; import org.apache.tsfile.file.metadata.TimeseriesMetadata; +import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.MetadataIndexNodeType; +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.ReadWriteIOUtils; import org.apache.tsfile.write.chunk.AlignedChunkWriterImpl; import org.apache.tsfile.write.chunk.ChunkWriterImpl; +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.schema.Schema; import org.apache.tsfile.write.writer.TsFileIOWriter; import org.apache.tsfile.write.writer.tsmiterator.TSMIterator; import org.junit.Assert; +import org.junit.Assume; import org.junit.Test; import java.io.ByteArrayInputStream; @@ -58,12 +76,17 @@ import java.io.DataOutputStream; import java.io.File; import java.io.IOException; +import java.lang.management.ManagementFactory; import java.nio.file.Files; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; public class TsFileSplitterTest { @@ -134,6 +157,121 @@ public void testSplitV3TsFile() throws Exception { } } + @Test + public void testSplitRejectsMoreTimePartitionsThanTheConfiguredLimit() throws Exception { + final File sourceTsFile = new File("split-partition-limit-source.tsfile"); + final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + final int originalPartitionSlotsLimit = config.getLoadTsFileSpiltPartitionMaxSize(); + try { + config.setLoadTsFileSpiltPartitionMaxSize(2); + writeTableTsFileWithOnePointPerTimePartition(sourceTsFile, 3); + + try { + new TsFileSplitter(sourceTsFile, tsFileData -> true).splitTsFileByDataPartition(); + Assert.fail("expected the splitter to reject more time partitions than the limit"); + } catch (final LoadFileException e) { + // expected: the third partition crosses the configured limit + } + } finally { + config.setLoadTsFileSpiltPartitionMaxSize(originalPartitionSlotsLimit); + if (sourceTsFile.exists()) { + Assert.assertTrue(sourceTsFile.delete()); + } + } + } + + @Test + public void testSplitAlignedChunkAcrossManyTimePartitionsDecodesEachPageOnce() throws Exception { + // One point per time partition in a single chunk group makes every value page span as many + // time partitions as the page holds points. Such a page must be decoded once and shared by the + // partitions that it covers: decoding it once per partition allocates one page-sized array per + // partition, i.e. O(points * partitions) objects, which exhausts the DataNode heap when a + // TsFile stores its points one time partition apart. + final int pointsPerPage = 1024; + final int pointCount = 2 * pointsPerPage; + final File sourceTsFile = new File("split-aligned-many-partitions-source.tsfile"); + final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + final int originalPartitionSlotsLimit = config.getLoadTsFileSpiltPartitionMaxSize(); + final int originalMaxPointsInPage = + TSFileDescriptor.getInstance().getConfig().getMaxNumberOfPointsInPage(); + try { + config.setLoadTsFileSpiltPartitionMaxSize(pointCount + 1); + TSFileDescriptor.getInstance().getConfig().setMaxNumberOfPointsInPage(pointsPerPage); + writeTableTsFileWithOnePointPerTimePartition(sourceTsFile, pointCount); + + final AtomicInteger emittedChunkDataCount = new AtomicInteger(); + final long allocatedBytes = + splitAndMeasureAllocatedBytes(sourceTsFile, emittedChunkDataCount); + + Assert.assertEquals(pointCount, emittedChunkDataCount.get()); + // Decoding the pages per partition too and recomputing the partition bounds per point both + // make this O(points * partitions) in allocations, which is hundreds of megabytes here. + Assert.assertTrue( + "splitting " + + pointCount + + " single-point time partitions allocated " + + allocatedBytes + + " bytes", + allocatedBytes < 128L * 1024 * 1024); + } finally { + config.setLoadTsFileSpiltPartitionMaxSize(originalPartitionSlotsLimit); + TSFileDescriptor.getInstance() + .getConfig() + .setMaxNumberOfPointsInPage(originalMaxPointsInPage); + if (sourceTsFile.exists()) { + Assert.assertTrue(sourceTsFile.delete()); + } + } + } + + private long splitAndMeasureAllocatedBytes( + final File sourceTsFile, final AtomicInteger emittedChunkDataCount) throws Exception { + final java.lang.management.ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean(); + Assume.assumeTrue( + "the JVM does not report the allocations of a thread", + threadMXBean instanceof ThreadMXBean); + final long allocatedBytesBefore = + ((ThreadMXBean) threadMXBean).getCurrentThreadAllocatedBytes(); + new TsFileSplitter( + sourceTsFile, + tsFileData -> { + if (tsFileData instanceof ChunkData) { + emittedChunkDataCount.incrementAndGet(); + } + return true; + }) + .splitTsFileByDataPartition(); + return ((ThreadMXBean) threadMXBean).getCurrentThreadAllocatedBytes() - allocatedBytesBefore; + } + + private void writeTableTsFileWithOnePointPerTimePartition(final File tsFile, final int pointCount) + throws Exception { + if (tsFile.exists()) { + Assert.assertTrue(tsFile.delete()); + } + final long baseTime = + TimePartitionUtils.getTimePartitionSlot(1_700_000_000_000L).getStartTime(); + final long timePartitionInterval = TimePartitionUtils.getTimePartitionInterval(); + + try (final TsFileIOWriter writer = new TsFileIOWriter(tsFile)) { + writer.setSchema(createSchema()); + writer.startChunkGroup(new StringArrayDeviceID("table1", "tagA")); + + final AlignedChunkWriterImpl chunkWriter = + new AlignedChunkWriterImpl( + Collections.singletonList(new MeasurementSchema("s1", TSDataType.INT64))); + for (int i = 0; i < pointCount; i++) { + final long time = baseTime + i * timePartitionInterval; + chunkWriter.getTimeChunkWriter().write(time); + chunkWriter.getValueChunkWriterByIndex(0).write(time, (long) i, false); + } + chunkWriter.writeToFileWriter(writer); + + writer.endChunkGroup(); + writer.endFile(); + } + } + private File constructV3TsFile() throws IOException { final File tsFile = Files.createTempFile("v3-tsfile-splitter", ".tsfile").toFile(); final IDeviceID deviceID = new PlainDeviceID("root.sg.d1"); @@ -182,6 +320,118 @@ private File constructV3TsFile() throws IOException { return tsFile; } + @Test + public void testSplitAlignedValueChunkWithUndecodedPagesPreservesValues() throws Exception { + assertAlignedValueChunkSplitPreservesValues(2); + } + + @Test + public void testSplitAlignedValueChunkSpanningTimePartitionsPreservesValues() throws Exception { + // A page that covers two time partitions has to be decoded and written once per partition. The + // partition ranges are located by binary search, so this also checks that no point is dropped + // or duplicated at the partition boundary. + assertAlignedValueChunkSplitPreservesValues(4); + } + + private void assertAlignedValueChunkSplitPreservesValues(final int maxNumberOfPointsInPage) + throws Exception { + final File sourceTsFile = + new File("split-table-aligned-values-" + maxNumberOfPointsInPage + "-source.tsfile"); + final File targetTsFile = + new File("split-table-aligned-values-" + maxNumberOfPointsInPage + "-target.tsfile"); + final IDeviceID deviceID = new StringArrayDeviceID("table1", "tagA"); + final int pointsPerPartition = 2; + final int partitionCount = 2; + final long baseTime = + TimePartitionUtils.getTimePartitionSlot(1_700_000_000_000L).getStartTime(); + final long timePartitionInterval = TimePartitionUtils.getTimePartitionInterval(); + final int pointCount = pointsPerPartition * partitionCount; + final long[] times = new long[pointCount]; + final long[] values = new long[pointCount]; + + final int previousMaxNumberOfPointsInPage = + TSFileDescriptor.getInstance().getConfig().getMaxNumberOfPointsInPage(); + try { + TSFileDescriptor.getInstance() + .getConfig() + .setMaxNumberOfPointsInPage(maxNumberOfPointsInPage); + if (sourceTsFile.exists()) { + Assert.assertTrue(sourceTsFile.delete()); + } + try (final TsFileIOWriter writer = new TsFileIOWriter(sourceTsFile)) { + writer.setSchema(createSchema()); + writer.startChunkGroup(deviceID); + final AlignedChunkWriterImpl chunkWriter = + new AlignedChunkWriterImpl( + Collections.singletonList(new MeasurementSchema("s1", TSDataType.INT64))); + for (int partition = 0; partition < partitionCount; partition++) { + for (int i = 0; i < pointsPerPartition; i++) { + final int index = partition * pointsPerPartition + i; + times[index] = baseTime + partition * timePartitionInterval + i; + values[index] = index; + chunkWriter.getTimeChunkWriter().write(times[index]); + chunkWriter.getValueChunkWriterByIndex(0).write(times[index], values[index], false); + } + } + chunkWriter.writeToFileWriter(writer); + writer.endChunkGroup(); + writer.endFile(); + } + + final List emittedChunkDataList = new ArrayList<>(); + new TsFileSplitter( + sourceTsFile, + tsFileData -> { + if (tsFileData instanceof ChunkData) { + emittedChunkDataList.add((ChunkData) tsFileData); + } + return true; + }) + .splitTsFileByDataPartition(); + + Assert.assertEquals(partitionCount, emittedChunkDataList.size()); + if (targetTsFile.exists()) { + Assert.assertTrue(targetTsFile.delete()); + } + try (final TsFileIOWriter writer = new TsFileIOWriter(targetTsFile)) { + writer.setSchema(createSchema()); + writer.startChunkGroup(deviceID); + for (final ChunkData chunkData : emittedChunkDataList) { + writeSerializedChunkDataToWriter(chunkData, writer); + } + writer.endChunkGroup(); + writer.endFile(); + } + + try (final TsFileReader reader = new TsFileReader(targetTsFile)) { + final QueryDataSet dataSet = + reader.query( + QueryExpression.create( + Collections.singletonList(new Path(deviceID, "s1", true)), null)); + int index = 0; + while (dataSet.hasNext()) { + final RowRecord row = dataSet.next(); + Assert.assertEquals(times[index], row.getTimestamp()); + final Field field = row.getField(0); + Assert.assertNotNull(field); + Assert.assertEquals(values[index], field.getLongV()); + index++; + } + Assert.assertEquals(pointCount, index); + } + } finally { + TSFileDescriptor.getInstance() + .getConfig() + .setMaxNumberOfPointsInPage(previousMaxNumberOfPointsInPage); + if (sourceTsFile.exists()) { + Assert.assertTrue(sourceTsFile.delete()); + } + if (targetTsFile.exists()) { + Assert.assertTrue(targetTsFile.delete()); + } + } + } + @Test public void testSplitTableTimeOnlyAlignedChunk() throws Exception { final File sourceTsFile = new File("split-table-time-only-source.tsfile"); @@ -394,6 +644,506 @@ public void testSplitTableTimeOnlyAlignedChunkAtExactLongMaxUpperBound() throws } } + private static final List BATCHED_MEASUREMENTS = + Arrays.asList("s1", "s2", "s3", "s4", "s5"); + + private static final List> BATCHED_MEASUREMENT_BATCHES = + Arrays.asList( + Arrays.asList("s1", "s2"), Arrays.asList("s3", "s4"), Collections.singletonList("s5")); + + /** + * A batched aligned TsFile holds one time chunk per time window, but the value chunks of the + * remaining series batches are appended behind all time chunks. That is what the compaction + * writer emits when {@code compactionMaxAlignedSeriesNumInOneBatch} is smaller than the series + * count, and it makes the splitter flush the chunk data of a time chunk several times. + */ + @Test + public void testSplitBatchedAlignedChunkAcrossTimePartitions() throws Exception { + final File sourceTsFile = new File("split-batched-aligned-source.tsfile"); + final IDeviceID deviceID = new StringArrayDeviceID("table1", "tagA"); + final int previousMaxNumberOfPointsInPage = + TSFileDescriptor.getInstance().getConfig().getMaxNumberOfPointsInPage(); + final long firstPartitionStartTime = + TimePartitionUtils.getTimePartitionSlot(1_700_000_000_000L).getStartTime(); + final long timePartitionInterval = TimePartitionUtils.getTimePartitionInterval(); + final List> timeWindows = + Arrays.asList( + Arrays.asList( + firstPartitionStartTime, firstPartitionStartTime + 1, firstPartitionStartTime + 2), + Arrays.asList( + firstPartitionStartTime + timePartitionInterval, + firstPartitionStartTime + timePartitionInterval + 1)); + final File[] targetTsFiles = { + new File("split-batched-aligned-target-0.tsfile"), + new File("split-batched-aligned-target-1.tsfile") + }; + try { + TSFileDescriptor.getInstance().getConfig().setMaxNumberOfPointsInPage(2); + writeBatchedAlignedTsFile( + sourceTsFile, + Collections.singletonList(deviceID), + timeWindows, + BATCHED_MEASUREMENT_BATCHES, + Collections.emptyMap(), + Collections.emptyMap()); + + final List emittedChunkDataList = splitSourceTsFile(sourceTsFile); + // The splitter flushes the chunk data at every series batch, so it is emitted several times. + Assert.assertTrue(emittedChunkDataList.size() > 2); + for (final ChunkData chunkData : emittedChunkDataList) { + assertNoDuplicatedChunk(chunkData); + } + + final Map> chunkDataByPartition = + groupChunkDataByPartition(emittedChunkDataList); + Assert.assertEquals(2, chunkDataByPartition.size()); + + final List firstPartitionChunkDataList = + chunkDataByPartition.get(new TTimePartitionSlot(firstPartitionStartTime)); + Assert.assertNotNull(firstPartitionChunkDataList); + Assert.assertTrue(firstPartitionChunkDataList.size() > 1); + assertBatchedAlignedTarget( + writeTargetTsFile(targetTsFiles[0], firstPartitionChunkDataList), + deviceID, + new long[] { + firstPartitionStartTime, firstPartitionStartTime + 1, firstPartitionStartTime + 2 + }, + 1, + Collections.emptySet()); + + final List secondPartitionChunkDataList = + chunkDataByPartition.get( + new TTimePartitionSlot(firstPartitionStartTime + timePartitionInterval)); + Assert.assertNotNull(secondPartitionChunkDataList); + Assert.assertTrue(secondPartitionChunkDataList.size() > 1); + assertBatchedAlignedTarget( + writeTargetTsFile(targetTsFiles[1], secondPartitionChunkDataList), + deviceID, + new long[] { + firstPartitionStartTime + timePartitionInterval, + firstPartitionStartTime + timePartitionInterval + 1 + }, + 1, + Collections.emptySet()); + } finally { + TSFileDescriptor.getInstance() + .getConfig() + .setMaxNumberOfPointsInPage(previousMaxNumberOfPointsInPage); + if (sourceTsFile.exists()) { + Assert.assertTrue(sourceTsFile.delete()); + } + for (final File targetTsFile : targetTsFiles) { + if (targetTsFile.exists()) { + Assert.assertTrue(targetTsFile.delete()); + } + } + } + } + + /** + * The same batched layout mixed with the remaining chunk shapes: the first time window stays + * inside one time partition and carries empty value chunks, while the second window crosses the + * partition boundary and carries the all-null (empty page) value chunks. + */ + @Test + public void testSplitBatchedAlignedChunkWithEmptyAndNullValueChunks() throws Exception { + final File sourceTsFile = new File("split-batched-aligned-special-source.tsfile"); + final IDeviceID deviceID = new StringArrayDeviceID("table1", "tagA"); + final int previousMaxNumberOfPointsInPage = + TSFileDescriptor.getInstance().getConfig().getMaxNumberOfPointsInPage(); + final long firstPartitionStartTime = + TimePartitionUtils.getTimePartitionSlot(1_700_100_000_000L).getStartTime(); + final long timePartitionInterval = TimePartitionUtils.getTimePartitionInterval(); + final long crossedTime = firstPartitionStartTime + timePartitionInterval - 1; + final long nextPartitionStartTime = firstPartitionStartTime + timePartitionInterval; + final List> timeWindows = + Arrays.asList( + Collections.singletonList(firstPartitionStartTime + 3), + Arrays.asList(crossedTime, nextPartitionStartTime)); + final Map> emptyMeasurementIdsByWindow = + Collections.singletonMap(0, new HashSet<>(Arrays.asList("s4", "s5"))); + final Map> nullMeasurementIdsByWindow = + Collections.singletonMap(1, new HashSet<>(Arrays.asList("s4", "s5"))); + final Set nullMeasurements = new HashSet<>(Arrays.asList("s4", "s5")); + final File[] targetTsFiles = { + new File("split-batched-aligned-special-target-0.tsfile"), + new File("split-batched-aligned-special-target-1.tsfile") + }; + try { + TSFileDescriptor.getInstance().getConfig().setMaxNumberOfPointsInPage(2); + writeBatchedAlignedTsFile( + sourceTsFile, + Collections.singletonList(deviceID), + timeWindows, + BATCHED_MEASUREMENT_BATCHES, + emptyMeasurementIdsByWindow, + nullMeasurementIdsByWindow); + + final List emittedChunkDataList = splitSourceTsFile(sourceTsFile); + Assert.assertTrue(emittedChunkDataList.size() > 2); + for (final ChunkData chunkData : emittedChunkDataList) { + assertNoDuplicatedChunk(chunkData); + } + + final Map> chunkDataByPartition = + groupChunkDataByPartition(emittedChunkDataList); + Assert.assertEquals(2, chunkDataByPartition.size()); + + final List firstPartitionChunkDataList = + chunkDataByPartition.get(new TTimePartitionSlot(firstPartitionStartTime)); + Assert.assertNotNull(firstPartitionChunkDataList); + assertBatchedAlignedTarget( + writeTargetTsFile(targetTsFiles[0], firstPartitionChunkDataList), + deviceID, + new long[] {firstPartitionStartTime + 3, crossedTime}, + 2, + nullMeasurements); + + final List secondPartitionChunkDataList = + chunkDataByPartition.get(new TTimePartitionSlot(nextPartitionStartTime)); + Assert.assertNotNull(secondPartitionChunkDataList); + assertBatchedAlignedTarget( + writeTargetTsFile(targetTsFiles[1], secondPartitionChunkDataList), + deviceID, + new long[] {nextPartitionStartTime}, + 1, + nullMeasurements); + } finally { + TSFileDescriptor.getInstance() + .getConfig() + .setMaxNumberOfPointsInPage(previousMaxNumberOfPointsInPage); + if (sourceTsFile.exists()) { + Assert.assertTrue(sourceTsFile.delete()); + } + for (final File targetTsFile : targetTsFiles) { + if (targetTsFile.exists()) { + Assert.assertTrue(targetTsFile.delete()); + } + } + } + } + + /** Every device of a batched TsFile keeps its own flush state. */ + @Test + public void testSplitBatchedAlignedChunkOfMultipleDevices() throws Exception { + final File sourceTsFile = new File("split-batched-aligned-multi-device-source.tsfile"); + final List deviceIDs = + Arrays.asList( + new StringArrayDeviceID("table1", "tagA"), new StringArrayDeviceID("table1", "tagB")); + final int previousMaxNumberOfPointsInPage = + TSFileDescriptor.getInstance().getConfig().getMaxNumberOfPointsInPage(); + final long firstPartitionStartTime = + TimePartitionUtils.getTimePartitionSlot(1_700_200_000_000L).getStartTime(); + final long timePartitionInterval = TimePartitionUtils.getTimePartitionInterval(); + final List> timeWindows = + Arrays.asList( + Arrays.asList(firstPartitionStartTime, firstPartitionStartTime + 1), + Arrays.asList( + firstPartitionStartTime + timePartitionInterval, + firstPartitionStartTime + timePartitionInterval + 1)); + final File[] targetTsFiles = { + new File("split-batched-aligned-multi-device-target-0.tsfile"), + new File("split-batched-aligned-multi-device-target-1.tsfile") + }; + try { + TSFileDescriptor.getInstance().getConfig().setMaxNumberOfPointsInPage(2); + writeBatchedAlignedTsFile( + sourceTsFile, + deviceIDs, + timeWindows, + BATCHED_MEASUREMENT_BATCHES, + Collections.emptyMap(), + Collections.emptyMap()); + + final List emittedChunkDataList = splitSourceTsFile(sourceTsFile); + for (final ChunkData chunkData : emittedChunkDataList) { + assertNoDuplicatedChunk(chunkData); + } + + final Map> chunkDataByPartition = + groupChunkDataByPartition(emittedChunkDataList); + Assert.assertEquals(2, chunkDataByPartition.size()); + + final List firstPartitionChunkDataList = + chunkDataByPartition.get(new TTimePartitionSlot(firstPartitionStartTime)); + Assert.assertNotNull(firstPartitionChunkDataList); + final File firstTargetTsFile = + writeTargetTsFile(targetTsFiles[0], firstPartitionChunkDataList); + final List secondPartitionChunkDataList = + chunkDataByPartition.get( + new TTimePartitionSlot(firstPartitionStartTime + timePartitionInterval)); + Assert.assertNotNull(secondPartitionChunkDataList); + final File secondTargetTsFile = + writeTargetTsFile(targetTsFiles[1], secondPartitionChunkDataList); + + for (final IDeviceID deviceID : deviceIDs) { + assertBatchedAlignedTarget( + firstTargetTsFile, + deviceID, + new long[] {firstPartitionStartTime, firstPartitionStartTime + 1}, + 1, + Collections.emptySet()); + assertBatchedAlignedTarget( + secondTargetTsFile, + deviceID, + new long[] { + firstPartitionStartTime + timePartitionInterval, + firstPartitionStartTime + timePartitionInterval + 1 + }, + 1, + Collections.emptySet()); + } + } finally { + TSFileDescriptor.getInstance() + .getConfig() + .setMaxNumberOfPointsInPage(previousMaxNumberOfPointsInPage); + if (sourceTsFile.exists()) { + Assert.assertTrue(sourceTsFile.delete()); + } + for (final File targetTsFile : targetTsFiles) { + if (targetTsFile.exists()) { + Assert.assertTrue(targetTsFile.delete()); + } + } + } + } + + private void writeBatchedAlignedTsFile( + final File tsFile, + final List deviceIDs, + final List> timeWindows, + final List> measurementBatches, + final Map> emptyMeasurementIdsByWindow, + final Map> nullMeasurementIdsByWindow) + throws Exception { + if (tsFile.exists()) { + Assert.assertTrue(tsFile.delete()); + } + try (final TsFileIOWriter writer = new TsFileIOWriter(tsFile)) { + writer.setSchema(createBatchedSchema()); + for (final IDeviceID deviceID : deviceIDs) { + writer.startChunkGroup(deviceID); + for (int batchIndex = 0; batchIndex < measurementBatches.size(); batchIndex++) { + final List measurementBatch = measurementBatches.get(batchIndex); + for (int windowIndex = 0; windowIndex < timeWindows.size(); windowIndex++) { + final List timeWindow = timeWindows.get(windowIndex); + final Set emptyMeasurements = + emptyMeasurementIdsByWindow.getOrDefault(windowIndex, Collections.emptySet()); + final Set nullMeasurements = + nullMeasurementIdsByWindow.getOrDefault(windowIndex, Collections.emptySet()); + if (batchIndex == 0) { + // The first batch carries the time chunk of the window, the following batches only + // append their value chunks behind it, exactly like the batched compaction writer. + final List measurementSchemas = new ArrayList<>(); + for (final String measurement : measurementBatch) { + if (!emptyMeasurements.contains(measurement)) { + measurementSchemas.add(new MeasurementSchema(measurement, TSDataType.INT64)); + } + } + final AlignedChunkWriterImpl alignedChunkWriter = + new AlignedChunkWriterImpl(measurementSchemas); + for (final long time : timeWindow) { + alignedChunkWriter.getTimeChunkWriter().write(time); + } + int valueIndex = 0; + for (final String measurement : measurementBatch) { + if (emptyMeasurements.contains(measurement)) { + continue; + } + final ValueChunkWriter valueChunkWriter = + alignedChunkWriter.getValueChunkWriterByIndex(valueIndex++); + for (final long time : timeWindow) { + valueChunkWriter.write( + time, + batchedValue(measurement, time), + nullMeasurements.contains(measurement)); + } + } + alignedChunkWriter.writeToFileWriter(writer); + } else { + for (final String measurement : measurementBatch) { + if (emptyMeasurements.contains(measurement)) { + continue; + } + final ValueChunkWriter valueChunkWriter = + new ValueChunkWriter( + measurement, + CompressionType.LZ4, + TSDataType.INT64, + TSEncoding.PLAIN, + new PlainEncoder(TSDataType.INT64, 8)); + for (final long time : timeWindow) { + valueChunkWriter.write( + time, + batchedValue(measurement, time), + nullMeasurements.contains(measurement)); + } + valueChunkWriter.writeToFileWriter(writer); + } + } + for (final String measurement : measurementBatch) { + if (emptyMeasurements.contains(measurement)) { + writer.writeEmptyValueChunk( + measurement, + CompressionType.LZ4, + TSDataType.INT64, + TSEncoding.PLAIN, + Statistics.getStatsByType(TSDataType.INT64)); + } + } + } + } + writer.endChunkGroup(); + } + writer.endFile(); + } + } + + private static long batchedValue(final String measurement, final long time) { + return (BATCHED_MEASUREMENTS.indexOf(measurement) + 1) * 1_000_000_000L + time; + } + + private List splitSourceTsFile(final File sourceTsFile) throws Exception { + final List emittedChunkDataList = new ArrayList<>(); + final TsFileSplitter splitter = + new TsFileSplitter( + sourceTsFile, + tsFileData -> { + if (tsFileData instanceof ChunkData) { + emittedChunkDataList.add((ChunkData) tsFileData); + } + return true; + }); + splitter.splitTsFileByDataPartition(); + return emittedChunkDataList; + } + + private Map> groupChunkDataByPartition( + final List emittedChunkDataList) { + final Map> chunkDataByPartition = new LinkedHashMap<>(); + for (final ChunkData chunkData : emittedChunkDataList) { + chunkDataByPartition + .computeIfAbsent(chunkData.getTimePartitionSlot(), slot -> new ArrayList<>()) + .add(chunkData); + } + return chunkDataByPartition; + } + + private void assertNoDuplicatedChunk(final ChunkData chunkData) { + final Set measurementIds = new HashSet<>(); + for (final Chunk chunk : chunkData.getChunks()) { + final String measurementId = chunk.getHeader().getMeasurementID(); + Assert.assertTrue( + "Duplicated chunk of measurement " + + measurementId + + " in one emitted ChunkData, the chunk data was flushed without being reset", + measurementIds.add(measurementId.isEmpty() ? "