diff --git a/CHANGES.txt b/CHANGES.txt index 95e18be773..e4ed0d362f 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -1,4 +1,5 @@ 5.1 + * Enable CQLSSTableWriter to create SSTables compressed with a dictionary (CASSANDRA-20938) * Support ZSTD dictionary compression (CASSANDRA-17021) * Fix ExceptionsTable when stacktrace has zero elements (CASSANDRA-20992) * Replace blocking wait with non-blocking delay in paxos repair (CASSANDRA-20983) diff --git a/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java b/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java index c12a76bc84..67d13c2327 100644 --- a/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java @@ -34,6 +34,7 @@ import org.apache.cassandra.config.DatabaseDescriptor; import org.apache.cassandra.db.DecoratedKey; import org.apache.cassandra.db.RegularAndStaticColumns; import org.apache.cassandra.db.SerializationHeader; +import org.apache.cassandra.db.compression.CompressionDictionary; import org.apache.cassandra.db.partitions.PartitionUpdate; import org.apache.cassandra.db.rows.EncodingStats; import org.apache.cassandra.index.Index; @@ -57,6 +58,8 @@ abstract class AbstractSSTableSimpleWriter implements Closeable protected final Collection indexGroups; protected Consumer> sstableProducedListener; protected boolean openSSTableOnProduced = false; + protected CompressionDictionary compressionDictionary; + protected SSTable.Owner owner; protected AbstractSSTableSimpleWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns) { @@ -81,6 +84,11 @@ abstract class AbstractSSTableSimpleWriter implements Closeable this.indexGroups.add(indexGroup); } + public void setCompressionDictionary(CompressionDictionary compressionDictionary) + { + this.compressionDictionary = compressionDictionary; + } + protected void setSSTableProducedListener(Consumer> listener) { this.sstableProducedListener = Objects.requireNonNull(listener, "sstableProducedListener cannot be null"); @@ -114,6 +122,20 @@ abstract class AbstractSSTableSimpleWriter implements Closeable if (makeRangeAware) return SSTableTxnWriter.createRangeAware(metadata, 0, ActiveRepairService.UNREPAIRED_SSTABLE, ActiveRepairService.NO_PENDING_REPAIR, false, format, header); + + SSTable.Owner effectiveOwner; + + if (this.owner != null && this.owner.compressionDictionaryManager() != null && compressionDictionary != null) + { + // already checks if it is cached or not + this.owner.compressionDictionaryManager().add(compressionDictionary); + effectiveOwner = this.owner; + } + else + { + effectiveOwner = owner; + } + return SSTableTxnWriter.create(metadata, createDescriptor(directory, metadata.keyspace, metadata.name, format), 0, @@ -122,7 +144,7 @@ abstract class AbstractSSTableSimpleWriter implements Closeable false, header, indexGroups, - owner); + effectiveOwner); } private static Descriptor createDescriptor(File directory, final String keyspace, final String columnFamily, final SSTableFormat fmt) throws IOException diff --git a/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java b/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java index 0676eb1f56..c1043a9fb7 100644 --- a/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java @@ -56,15 +56,18 @@ import org.apache.cassandra.db.Directories; import org.apache.cassandra.db.Keyspace; import org.apache.cassandra.db.Slice; import org.apache.cassandra.db.Slices; +import org.apache.cassandra.db.compression.CompressionDictionary; import org.apache.cassandra.db.marshal.AbstractType; import org.apache.cassandra.dht.IPartitioner; import org.apache.cassandra.dht.Murmur3Partitioner; import org.apache.cassandra.exceptions.InvalidRequestException; import org.apache.cassandra.exceptions.SyntaxException; import org.apache.cassandra.index.sai.StorageAttachedIndexGroup; +import org.apache.cassandra.io.compress.IDictionaryCompressor; import org.apache.cassandra.io.sstable.format.SSTableFormat; import org.apache.cassandra.io.sstable.format.SSTableReader; import org.apache.cassandra.io.util.File; +import org.apache.cassandra.schema.CompressionParams; import org.apache.cassandra.schema.KeyspaceMetadata; import org.apache.cassandra.schema.KeyspaceParams; import org.apache.cassandra.schema.Keyspaces; @@ -413,6 +416,7 @@ public class CQLSSTableWriter implements Closeable private boolean buildIndexes = true; private Consumer> sstableProducedListener; private boolean openSSTableOnProduced = false; + private CompressionDictionary compressionDictionary = null; protected Builder() { @@ -665,6 +669,18 @@ public class CQLSSTableWriter implements Closeable return this; } + /** + * Use specific compression dictionary upon writing the data. + * + * @param compressionDictionary compression dictionary to use + * @return this builder + */ + public Builder withCompressionDictionary(CompressionDictionary compressionDictionary) + { + this.compressionDictionary = compressionDictionary; + return this; + } + public CQLSSTableWriter build() { if (directory == null) @@ -729,8 +745,26 @@ public class CQLSSTableWriter implements Closeable Schema.instance.submit(SchemaTransformations.addTable(tableMetadata, true)); } + if (compressionDictionary != null) + { + CompressionParams compressionParams = tableMetadata.params.compression; + + if (!compressionParams.isDictionaryCompressionEnabled()) + { + throw new IllegalStateException("Table's compressor can not accept any dictionary: " + compressionParams.asMap()); + } + + IDictionaryCompressor compressor = (IDictionaryCompressor) compressionParams.getSstableCompressor(); + if (!compressor.canConsumeDictionary(compressionDictionary)) + { + throw new IllegalStateException("Provided dictionary can not be consumed by table's compressor. " + + "Provided dictionary type: " + compressionDictionary.kind() + + "; expected dictionary type by the compressor: " + compressor.acceptableDictionaryKind()); + } + } + ColumnFamilyStore cfs = null; - if (buildIndexes && !indexStatements.isEmpty()) + if ((buildIndexes && !indexStatements.isEmpty()) || compressionDictionary != null) { KeyspaceMetadata keyspaceMetadata = ClusterMetadata.current().schema.getKeyspaceMetadata(keyspaceName); Keyspace keyspace = Keyspace.mockKS(keyspaceMetadata); @@ -764,8 +798,8 @@ public class CQLSSTableWriter implements Closeable TableMetadataRef ref = tableMetadata.ref; AbstractSSTableSimpleWriter writer = sorted - ? new SSTableSimpleWriter(directory, ref, preparedModificationStatement.updatedColumns(), maxSSTableSizeInMiB) - : new SSTableSimpleUnsortedWriter(directory, ref, preparedModificationStatement.updatedColumns(), maxSSTableSizeInMiB); + ? new SSTableSimpleWriter(cfs, directory, ref, preparedModificationStatement.updatedColumns(), maxSSTableSizeInMiB) + : new SSTableSimpleUnsortedWriter(cfs, directory, ref, preparedModificationStatement.updatedColumns(), maxSSTableSizeInMiB); if (format != null) writer.setSSTableFormatType(format); @@ -777,6 +811,11 @@ public class CQLSSTableWriter implements Closeable writer.addIndexGroup(saiGroup); } + if (compressionDictionary != null && cfs != null) + { + writer.setCompressionDictionary(compressionDictionary); + } + if (sstableProducedListener != null) writer.setSSTableProducedListener(sstableProducedListener); diff --git a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java index cf406af4c4..d90cc898bc 100644 --- a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java @@ -69,12 +69,18 @@ class SSTableSimpleUnsortedWriter extends AbstractSSTableSimpleWriter private final DiskWriter diskWriter = new DiskWriter(); SSTableSimpleUnsortedWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns, long maxSSTableSizeInMiB) + { + this(null, directory, metadata, columns, maxSSTableSizeInMiB); + } + + SSTableSimpleUnsortedWriter(SSTable.Owner owner, File directory, TableMetadataRef metadata, RegularAndStaticColumns columns, long maxSSTableSizeInMiB) { super(directory, metadata, columns); this.maxSStableSizeInBytes = maxSSTableSizeInMiB * 1024L * 1024L; this.header = new SerializationHeader(true, metadata.get(), columns, EncodingStats.NO_STATS); this.helper = new SerializationHelper(this.header); diskWriter.start(); + this.owner = owner; } @Override diff --git a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java index bd0ba94cce..05120fa65d 100644 --- a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java @@ -65,9 +65,27 @@ class SSTableSimpleWriter extends AbstractSSTableSimpleWriter * Any non-positive value indicates the sstable size is unlimited. */ protected SSTableSimpleWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns, long maxSSTableSizeInMiB) + { + this(null, directory, metadata, columns, maxSSTableSizeInMiB); + } + + /** + * Create a SSTable writer for sorted input data. + * When a positive {@param maxSSTableSizeInMiB} is defined, the writer outputs a sequence of SSTables, + * whose sizes do not exceed the specified value. + * + * @param owner owner of sstables this writer creates + * @param directory directory to store the sstable files + * @param metadata table metadata + * @param columns columns to update + * @param maxSSTableSizeInMiB defines the max SSTable size if the value is positive. + * Any non-positive value indicates the sstable size is unlimited. + */ + protected SSTableSimpleWriter(SSTable.Owner owner, File directory, TableMetadataRef metadata, RegularAndStaticColumns columns, long maxSSTableSizeInMiB) { super(directory, metadata, columns); this.maxSSTableSizeInBytes = maxSSTableSizeInMiB * 1024L * 1024L; + this.owner = owner; } @Override diff --git a/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java b/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java index f2a131b716..19a9d4cae1 100644 --- a/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java +++ b/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java @@ -30,6 +30,7 @@ import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Map; +import java.util.Random; import java.util.Set; import java.util.UUID; import java.util.concurrent.ExecutionException; @@ -58,6 +59,11 @@ import org.apache.cassandra.cql3.functions.types.LocalDate; import org.apache.cassandra.cql3.functions.types.TypeCodec; import org.apache.cassandra.cql3.functions.types.UDTValue; import org.apache.cassandra.cql3.functions.types.UserType; +import org.apache.cassandra.db.compression.CompressionDictionary; +import org.apache.cassandra.db.compression.CompressionDictionary.DictId; +import org.apache.cassandra.db.compression.CompressionDictionaryTrainingConfig; +import org.apache.cassandra.db.compression.ZstdCompressionDictionary; +import org.apache.cassandra.db.compression.ZstdDictionaryTrainer; import org.apache.cassandra.db.marshal.FloatType; import org.apache.cassandra.db.marshal.UTF8Type; import org.apache.cassandra.dht.ByteOrderedPartitioner; @@ -84,10 +90,11 @@ import org.apache.cassandra.utils.ByteBufferUtil; import org.apache.cassandra.utils.FBUtilities; import org.apache.cassandra.utils.JavaDriverUtils; import org.apache.cassandra.utils.OutputHandler; -import org.assertj.core.api.Assertions; +import static org.apache.cassandra.db.compression.CompressionDictionary.Kind.ZSTD; import static org.apache.cassandra.utils.Clock.Global.currentTimeMillis; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; @@ -1671,7 +1678,7 @@ public abstract class CQLSSTableWriterTest writer.addRow(1, 4); - Assertions.assertThatThrownBy(() -> writer.addRow(2, 11)) + assertThatThrownBy(() -> writer.addRow(2, 11)) .describedAs("Should throw when adding a row that violates constraints") .isInstanceOf(ConstraintViolationException.class) .hasMessageContaining("Column value does not satisfy value constraint for column 'v1'. It should be v1 < 5"); @@ -1690,6 +1697,111 @@ public abstract class CQLSSTableWriterTest } } + @Test + public void testWritingWithZstdDictionaryWhenUsingInvalidCompressor() + { + // the compressor is not dictionary-aware so we will fail + final String schema = "CREATE TABLE " + qualifiedTable + " (" + + " k int," + + " v1 text," + + " PRIMARY KEY (k)" + + ") WITH compression = {'class': 'ZstdCompressor'}"; + + assertThatThrownBy(() -> CQLSSTableWriter.builder() + .inDirectory(dataDir) + .forTable(schema) + .using("INSERT INTO " + keyspace + '.' + table + " (k, v1) VALUES (?, ?)") + // does not matter, we will fail anyway + .withCompressionDictionary(new ZstdCompressionDictionary(new DictId(ZSTD, 1), new byte[0])) + .build()) + .hasMessage("Table's compressor can not accept any dictionary: {chunk_length_in_kb=16, class=org.apache.cassandra.io.compress.ZstdCompressor}") + .isInstanceOf(IllegalStateException.class); + } + + @Test + public void testWritingWithZstdDictionary() throws Exception + { + final String schema = "CREATE TABLE " + qualifiedTable + " (" + + " k int," + + " v1 text," + + " PRIMARY KEY (k)" + + ") WITH compression = {'class': 'ZstdDictionaryCompressor'}"; + + CompressionDictionary dictionary = DictionaryHelper.trainDictionary(keyspace, table); + + CQLSSTableWriter writer = CQLSSTableWriter.builder() + .inDirectory(dataDir) + .forTable(schema) + .using("INSERT INTO " + keyspace + '.' + table + " (k, v1) VALUES (?, ?)") + .withCompressionDictionary(dictionary) + .build(); + + for (int i = 0; i < 500; i++) + { + writer.addRow(i, DictionaryHelper.INSTANCE.getRandomSample()); + } + + writer.close(); + + loadSSTables(dataDir, keyspace, table); + + if (verifyDataAfterLoading) + { + UntypedResultSet resultSet = QueryProcessor.executeInternal("SELECT * FROM " + qualifiedTable); + assertNotNull(resultSet); + Iterator iter = resultSet.iterator(); + for (int i = 0; i < 500; i++) + { + UntypedResultSet.Row row = iter.next(); + assertEquals(i, row.getInt("k")); + assertNotNull(row.getString("v1")); + } + } + } + + /** + * Simple generator of random data for Zstd compression dictionary and dictionary trainer. + */ + private static class DictionaryHelper + { + public static final DictionaryHelper INSTANCE = new DictionaryHelper(); + private static final Random random = new Random(); + + private static final String[] dates = new String[] {"2025-10-20","2025-10-19","2025-10-18","2025-10-17","2025-10-16"}; + private static final String[] times = new String[] {"11:00:01","11:00:02","11:00:03","11:00:04","11:00:05"}; + private static final String[] levels = new String[] {"TRACE", "DEBUG", "INFO", "WARN", "ERROR"}; + private static final String[] services = new String[] {"com.example.UserService", "com.example.DatabasePool", "com.example.PaymentService", "com.example.OrderService"}; + + private String getRandomSample() + { + return dates[random.nextInt(dates.length)] + ' ' + + times[random.nextInt(times.length)] + ' ' + + levels[random.nextInt(levels.length)] + ' ' + + services[random.nextInt(services.length)] + ' ' + + UUID.randomUUID(); // message + } + + private static CompressionDictionary trainDictionary(String keyspace, String table) + { + CompressionDictionaryTrainingConfig config = CompressionDictionaryTrainingConfig + .builder() + .maxDictionarySize(65536) + .maxTotalSampleSize(1024 * 1024) // 1MB total + .build(); + + try (ZstdDictionaryTrainer trainer = new ZstdDictionaryTrainer(keyspace, table, config, 3)) + { + trainer.start(true); + for (int i = 0; i < 25000; i++) + { + trainer.addSample(UTF8Type.instance.fromString(DictionaryHelper.INSTANCE.getRandomSample())); + } + + return trainer.trainDictionary(false); + } + } + } + protected static void loadSSTables(File dataDir, final String ks, final String tb) throws ExecutionException, InterruptedException { SSTableLoader loader = new SSTableLoader(dataDir, new SSTableLoader.Client()