diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/autocreate/IoTDBPipeConnectorCompressionIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/autocreate/IoTDBPipeConnectorCompressionIT.java index be4aa458d9fb0..f7c0c63b842e8 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/autocreate/IoTDBPipeConnectorCompressionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/autocreate/IoTDBPipeConnectorCompressionIT.java @@ -22,6 +22,8 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient; import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq; +import org.apache.iotdb.confignode.rpc.thrift.TShowPipeInfo; +import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq; import org.apache.iotdb.consensus.ConsensusFactory; import org.apache.iotdb.db.it.utils.TestUtils; import org.apache.iotdb.it.env.MultiEnvFactory; @@ -36,11 +38,17 @@ import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; +import java.sql.Connection; +import java.sql.SQLException; +import java.sql.Statement; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; +import static org.junit.Assert.fail; + @RunWith(IoTDBTestRunner.class) @Category({MultiClusterIT2AutoCreateSchema.class}) public class IoTDBPipeConnectorCompressionIT extends AbstractPipeDualAutoIT { @@ -179,4 +187,122 @@ private void doTest( Collections.singleton("8,")); } } + + @Test + public void testZstdCompressorLevel() throws Exception { + final DataNodeWrapper receiverDataNode = receiverEnv.getDataNodeWrapper(0); + + final String receiverIp = receiverDataNode.getIp(); + final int receiverPort = receiverDataNode.getPort(); + + try (final SyncConfigNodeIServiceClient client = + (SyncConfigNodeIServiceClient) senderEnv.getLeaderConfigNodeConnection()) { + if (!TestUtils.tryExecuteNonQueriesWithRetry( + senderEnv, + Arrays.asList( + "insert into root.db.d1(time, s1) values (1, 1)", + "insert into root.db.d1(time, s2) values (1, 1)", + "insert into root.db.d1(time, s3) values (1, 1)", + "insert into root.db.d1(time, s4) values (1, 1)", + "insert into root.db.d1(time, s5) values (1, 1)", + "flush"))) { + return; + } + + // Create 5 pipes with different zstd compression levels, p4 and p5 should fail. + + try (final Connection connection = senderEnv.getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute( + String.format( + "create pipe p1" + + " with extractor ('extractor.pattern'='root.db.d1.s1')" + + " with connector (" + + "'connector.ip'='%s'," + + "'connector.port'='%s'," + + "'connector.compressor'='zstd, zstd'," + + "'connector.compressor.zstd.level'='3')", + receiverIp, receiverPort)); + } catch (SQLException e) { + e.printStackTrace(); + fail(e.getMessage()); + } + + try (final Connection connection = senderEnv.getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute( + String.format( + "create pipe p2" + + " with extractor ('extractor.pattern'='root.db.d1.s2')" + + " with connector (" + + "'connector.ip'='%s'," + + "'connector.port'='%s'," + + "'connector.compressor'='zstd, zstd'," + + "'connector.compressor.zstd.level'='22')", + receiverIp, receiverPort)); + } catch (SQLException e) { + e.printStackTrace(); + fail(e.getMessage()); + } + + try (final Connection connection = senderEnv.getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute( + String.format( + "create pipe p3" + + " with extractor ('extractor.pattern'='root.db.d1.s3')" + + " with connector (" + + "'connector.ip'='%s'," + + "'connector.port'='%s'," + + "'connector.compressor'='zstd, zstd'," + + "'connector.compressor.zstd.level'='-131072')", + receiverIp, receiverPort)); + } catch (SQLException e) { + e.printStackTrace(); + fail(e.getMessage()); + } + + try (final Connection connection = senderEnv.getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute( + String.format( + "create pipe p4" + + " with extractor ('extractor.pattern'='root.db.d1.s4')" + + " with connector (" + + "'connector.ip'='%s'," + + "'connector.port'='%s'," + + "'connector.compressor'='zstd, zstd'," + + "'connector.compressor.zstd.level'='-131073')", + receiverIp, receiverPort)); + fail(); + } catch (SQLException e) { + // Make sure the error message in IoTDBConnector.java is returned + Assert.assertTrue(e.getMessage().contains("Zstd compression level should be in the range")); + } + + try (final Connection connection = senderEnv.getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute( + String.format( + "create pipe p5" + + " with extractor ('extractor.pattern'='root.db.d1.s5')" + + " with connector (" + + "'connector.ip'='%s'," + + "'connector.port'='%s'," + + "'connector.compressor'='zstd, zstd'," + + "'connector.compressor.zstd.level'='23')", + receiverIp, receiverPort)); + fail(); + } catch (SQLException e) { + // Make sure the error message in IoTDBConnector.java is returned + Assert.assertTrue(e.getMessage().contains("Zstd compression level should be in the range")); + } + + final List showPipeResult = client.showPipe(new TShowPipeReq()).pipeInfoList; + Assert.assertEquals(3, showPipeResult.size()); + + TestUtils.assertDataEventuallyOnEnv( + receiverEnv, "count timeseries", "count(timeseries),", Collections.singleton("3,")); + } + } } diff --git a/iotdb-core/node-commons/pom.xml b/iotdb-core/node-commons/pom.xml index 37371509d0cf4..a45bfb5f1e651 100644 --- a/iotdb-core/node-commons/pom.xml +++ b/iotdb-core/node-commons/pom.xml @@ -155,6 +155,10 @@ org.apache.commons commons-jexl3 + + com.github.luben + zstd-jni + org.mockito mockito-core diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeConnectorConstant.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeConnectorConstant.java index 3d0bd4cab7468..8d6879ae8c23b 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeConnectorConstant.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeConnectorConstant.java @@ -22,6 +22,8 @@ import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.pipe.config.PipeConfig; +import com.github.luben.zstd.Zstd; + import java.io.File; import java.util.Arrays; import java.util.Collections; @@ -177,6 +179,14 @@ public class PipeConnectorConstant { CONNECTOR_COMPRESSOR_ZSTD, CONNECTOR_COMPRESSOR_LZMA2))); + public static final String CONNECTOR_COMPRESSOR_ZSTD_LEVEL_KEY = + "connector.compressor.zstd.level"; + public static final String SINK_COMPRESSOR_ZSTD_LEVEL_KEY = "sink.compressor.zstd.level"; + public static final int CONNECTOR_COMPRESSOR_ZSTD_LEVEL_DEFAULT_VALUE = + Zstd.defaultCompressionLevel(); + public static final int CONNECTOR_COMPRESSOR_ZSTD_LEVEL_MIN_VALUE = Zstd.minCompressionLevel(); + public static final int CONNECTOR_COMPRESSOR_ZSTD_LEVEL_MAX_VALUE = Zstd.maxCompressionLevel(); + public static final String CONNECTOR_RATE_LIMIT_KEY = "connector.rate-limit-bytes-per-second"; public static final String SINK_RATE_LIMIT_KEY = "sink.rate-limit-bytes-per-second"; public static final double CONNECTOR_RATE_LIMIT_DEFAULT_VALUE = -1; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressorConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressorConfig.java new file mode 100644 index 0000000000000..c13028c4d31f0 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressorConfig.java @@ -0,0 +1,39 @@ +/* + * 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.commons.pipe.connector.compressor; + +public class PipeCompressorConfig { + + private final String name; + private final int zstdCompressionLevel; + + public PipeCompressorConfig(String name, int zstdCompressionLevel) { + this.name = name; + this.zstdCompressionLevel = zstdCompressionLevel; + } + + public String getName() { + return name; + } + + public int getZstdCompressionLevel() { + return zstdCompressionLevel; + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressorFactory.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressorFactory.java index 14ae972fc4f04..bcee9b85aebe5 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressorFactory.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressorFactory.java @@ -19,50 +19,86 @@ package org.apache.iotdb.commons.pipe.connector.compressor; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import java.util.Collections; import java.util.HashMap; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_GZIP; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_LZ4; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_LZMA2; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_SNAPPY; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_ZSTD; +import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_ZSTD_LEVEL_DEFAULT_VALUE; public class PipeCompressorFactory { - private static Map COMPRESSOR_NAME_TO_INSTANCE = new HashMap<>(); + private static final Logger LOGGER = LoggerFactory.getLogger(PipeCompressorFactory.class); + + private static final Map COMPRESSOR_NAME_TO_INSTANCE = + new ConcurrentHashMap<>(); static { COMPRESSOR_NAME_TO_INSTANCE.put(CONNECTOR_COMPRESSOR_SNAPPY, new PipeSnappyCompressor()); COMPRESSOR_NAME_TO_INSTANCE.put(CONNECTOR_COMPRESSOR_GZIP, new PipeGZIPCompressor()); COMPRESSOR_NAME_TO_INSTANCE.put(CONNECTOR_COMPRESSOR_LZ4, new PipeLZ4Compressor()); - COMPRESSOR_NAME_TO_INSTANCE.put(CONNECTOR_COMPRESSOR_ZSTD, new PipeZSTDCompressor()); + COMPRESSOR_NAME_TO_INSTANCE.put( + CONNECTOR_COMPRESSOR_ZSTD, + new PipeZSTDCompressor(CONNECTOR_COMPRESSOR_ZSTD_LEVEL_DEFAULT_VALUE)); COMPRESSOR_NAME_TO_INSTANCE.put(CONNECTOR_COMPRESSOR_LZMA2, new PipeLZMA2Compressor()); - COMPRESSOR_NAME_TO_INSTANCE = Collections.unmodifiableMap(COMPRESSOR_NAME_TO_INSTANCE); } - public static PipeCompressor getCompressor(String name) { - final PipeCompressor compressor = COMPRESSOR_NAME_TO_INSTANCE.get(name); - if (compressor == null) { - throw new UnsupportedOperationException("PipeCompressor not found for name: " + name); + public static PipeCompressor getCompressor(PipeCompressorConfig config) { + if (config == null) { + throw new IllegalArgumentException("PipeCompressorConfig is null"); } - return compressor; + if (config.getName() == null) { + throw new IllegalArgumentException("PipeCompressorConfig.getName() is null"); + } + + final String compressorName = config.getName(); + + // For ZSTD compressor, we need to consider the compression level + if (compressorName.equals(CONNECTOR_COMPRESSOR_ZSTD)) { + final int zstdCompressionLevel = config.getZstdCompressionLevel(); + return COMPRESSOR_NAME_TO_INSTANCE.computeIfAbsent( + CONNECTOR_COMPRESSOR_ZSTD + "_" + zstdCompressionLevel, + key -> { + LOGGER.info("Create new PipeZSTDCompressor with level: {}", zstdCompressionLevel); + return new PipeZSTDCompressor(zstdCompressionLevel); + }); + } + + // For other compressors, we can directly get the instance by name + final PipeCompressor compressor = COMPRESSOR_NAME_TO_INSTANCE.get(compressorName); + if (compressor != null) { + return compressor; + } + + throw new UnsupportedOperationException("PipeCompressor not found for name: " + compressorName); } private static Map COMPRESSOR_INDEX_TO_INSTANCE = new HashMap<>(); static { COMPRESSOR_INDEX_TO_INSTANCE.put( - PipeCompressor.PipeCompressionType.SNAPPY.getIndex(), new PipeSnappyCompressor()); + PipeCompressor.PipeCompressionType.SNAPPY.getIndex(), + COMPRESSOR_NAME_TO_INSTANCE.get(CONNECTOR_COMPRESSOR_SNAPPY)); COMPRESSOR_INDEX_TO_INSTANCE.put( - PipeCompressor.PipeCompressionType.GZIP.getIndex(), new PipeGZIPCompressor()); + PipeCompressor.PipeCompressionType.GZIP.getIndex(), + COMPRESSOR_NAME_TO_INSTANCE.get(CONNECTOR_COMPRESSOR_GZIP)); COMPRESSOR_INDEX_TO_INSTANCE.put( - PipeCompressor.PipeCompressionType.LZ4.getIndex(), new PipeLZ4Compressor()); + PipeCompressor.PipeCompressionType.LZ4.getIndex(), + COMPRESSOR_NAME_TO_INSTANCE.get(CONNECTOR_COMPRESSOR_LZ4)); COMPRESSOR_INDEX_TO_INSTANCE.put( - PipeCompressor.PipeCompressionType.ZSTD.getIndex(), new PipeZSTDCompressor()); + PipeCompressor.PipeCompressionType.ZSTD.getIndex(), + COMPRESSOR_NAME_TO_INSTANCE.get(CONNECTOR_COMPRESSOR_ZSTD)); COMPRESSOR_INDEX_TO_INSTANCE.put( - PipeCompressor.PipeCompressionType.LZMA2.getIndex(), new PipeLZMA2Compressor()); + PipeCompressor.PipeCompressionType.LZMA2.getIndex(), + COMPRESSOR_NAME_TO_INSTANCE.get(CONNECTOR_COMPRESSOR_LZMA2)); COMPRESSOR_INDEX_TO_INSTANCE = Collections.unmodifiableMap(COMPRESSOR_INDEX_TO_INSTANCE); } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeZSTDCompressor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeZSTDCompressor.java index 72782353d689f..50e2e1f845cf7 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeZSTDCompressor.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeZSTDCompressor.java @@ -19,29 +19,26 @@ package org.apache.iotdb.commons.pipe.connector.compressor; -import org.apache.tsfile.compress.ICompressor; -import org.apache.tsfile.compress.IUnCompressor; -import org.apache.tsfile.file.metadata.enums.CompressionType; +import com.github.luben.zstd.Zstd; import java.io.IOException; public class PipeZSTDCompressor extends PipeCompressor { - private static final ICompressor COMPRESSOR = ICompressor.getCompressor(CompressionType.ZSTD); - private static final IUnCompressor DECOMPRESSOR = - IUnCompressor.getUnCompressor(CompressionType.ZSTD); + private final int compressionLevel; - public PipeZSTDCompressor() { + public PipeZSTDCompressor(int compressionLevel) { super(PipeCompressionType.ZSTD); + this.compressionLevel = compressionLevel; } @Override public byte[] compress(byte[] data) throws IOException { - return COMPRESSOR.compress(data); + return Zstd.compress(data, compressionLevel); } @Override - public byte[] decompress(byte[] byteArray) throws IOException { - return DECOMPRESSOR.uncompress(byteArray); + public byte[] decompress(byte[] byteArray) { + return Zstd.decompress(byteArray, (int) Zstd.decompressedSize(byteArray, 0, byteArray.length)); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/protocol/IoTDBConnector.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/protocol/IoTDBConnector.java index f65daf86d4ff8..ad8e328e45e03 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/protocol/IoTDBConnector.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/protocol/IoTDBConnector.java @@ -22,6 +22,7 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.commons.pipe.connector.PipeReceiverStatusHandler; import org.apache.iotdb.commons.pipe.connector.compressor.PipeCompressor; +import org.apache.iotdb.commons.pipe.connector.compressor.PipeCompressorConfig; import org.apache.iotdb.commons.pipe.connector.compressor.PipeCompressorFactory; import org.apache.iotdb.commons.pipe.connector.limiter.GlobalRateLimiter; import org.apache.iotdb.commons.pipe.connector.limiter.PipeEndPointRateLimiter; @@ -50,6 +51,10 @@ import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_DEFAULT_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_SET; +import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_ZSTD_LEVEL_DEFAULT_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_ZSTD_LEVEL_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_ZSTD_LEVEL_MAX_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_COMPRESSOR_ZSTD_LEVEL_MIN_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_EXCEPTION_CONFLICT_RECORD_IGNORED_DATA_DEFAULT_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_EXCEPTION_CONFLICT_RECORD_IGNORED_DATA_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_EXCEPTION_CONFLICT_RESOLVE_STRATEGY_DEFAULT_VALUE; @@ -72,6 +77,7 @@ import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_RATE_LIMIT_DEFAULT_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_RATE_LIMIT_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.SINK_COMPRESSOR_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.SINK_COMPRESSOR_ZSTD_LEVEL_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.SINK_EXCEPTION_CONFLICT_RECORD_IGNORED_DATA_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.SINK_EXCEPTION_CONFLICT_RESOLVE_STRATEGY_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.SINK_EXCEPTION_CONFLICT_RETRY_MAX_TIME_SECONDS_KEY; @@ -151,6 +157,21 @@ public void validate(PipeParameterValidator validator) throws Exception { CONNECTOR_LOAD_BALANCE_STRATEGY_SET, loadBalanceStrategy), loadBalanceStrategy); + final int zstdCompressionLevel = + parameters.getIntOrDefault( + Arrays.asList(CONNECTOR_COMPRESSOR_ZSTD_LEVEL_KEY, SINK_COMPRESSOR_ZSTD_LEVEL_KEY), + CONNECTOR_COMPRESSOR_ZSTD_LEVEL_DEFAULT_VALUE); + validator.validate( + arg -> + (int) arg >= CONNECTOR_COMPRESSOR_ZSTD_LEVEL_MIN_VALUE + && (int) arg <= CONNECTOR_COMPRESSOR_ZSTD_LEVEL_MAX_VALUE, + String.format( + "Zstd compression level should be in the range [%d, %d], but got %d.", + CONNECTOR_COMPRESSOR_ZSTD_LEVEL_MIN_VALUE, + CONNECTOR_COMPRESSOR_ZSTD_LEVEL_MAX_VALUE, + zstdCompressionLevel), + zstdCompressionLevel); + final String compressionTypes = parameters .getStringOrDefault( @@ -170,7 +191,9 @@ public void validate(PipeParameterValidator validator) throws Exception { "Compressor should be one of %s, but got %s.", CONNECTOR_COMPRESSOR_SET, trimmedCompressionType), trimmedCompressionType); - compressors.add(PipeCompressorFactory.getCompressor(trimmedCompressionType)); + compressors.add( + PipeCompressorFactory.getCompressor( + new PipeCompressorConfig(trimmedCompressionType, zstdCompressionLevel))); } } validator.validate(