Enable CQLSSTableWriter to create SSTables compressed with a dictionary

patch by Stefan Miklosovic; reviewed by Yifan Cai for CASSANDRA-20938
This commit is contained in:
Stefan Miklosovic 2025-10-20 19:22:35 +02:00
parent c9e31f297c
commit 1a96b8803a
No known key found for this signature in database
GPG Key ID: 32F35CB2F546D93E
6 changed files with 204 additions and 6 deletions

View File

@ -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)

View File

@ -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<Index.Group> indexGroups;
protected Consumer<Collection<SSTableReader>> 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<Collection<SSTableReader>> 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

View File

@ -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<Collection<SSTableReader>> 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);

View File

@ -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

View File

@ -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

View File

@ -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<UntypedResultSet.Row> 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()