CASSANDRA-21423: CQLSSTableWriter produces empty bloom filter

CHANGES.txt

@@ -1,4 +1,5 @@
 7.0
+ * Rebuild the Bloom filter for CQLSSTableWriter-produced SSTables so it is sized for the real partition count (CASSANDRA-21423)
  * Implementation of CEP-49: Hardware-accelerated compression (CASSANDRA-20975)
  * Avoid using ObjectUtils.getFirstNonNull in Schema (CASSANDRA-21394)
  * Allow nodetool garbagecollect to take a user defined list of SSTables (CASSANDRA-16767)

src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java

@@ -25,11 +25,14 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 
+import com.google.common.base.Preconditions;
+
 import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.RegularAndStaticColumns;
@@ -40,9 +43,13 @@
 import org.apache.cassandra.index.Index;
 import org.apache.cassandra.io.sstable.format.SSTableFormat;
 import org.apache.cassandra.io.sstable.format.SSTableReader;
+import org.apache.cassandra.io.sstable.format.SSTableReaderLoadingBuilder;
+import org.apache.cassandra.io.sstable.format.SSTableReaderWithFilter;
 import org.apache.cassandra.io.util.File;
+import org.apache.cassandra.metrics.TableMetrics;
 import org.apache.cassandra.schema.TableMetadataRef;
 import org.apache.cassandra.service.ActiveRepairService;
+import org.apache.cassandra.utils.IFilter;
 
 /**
  * Base class for the sstable writers used by CQLSSTableWriter.
@@ -115,6 +122,64 @@ protected void notifySSTableProduced(Collection<SSTableReader> sstables)
         sstableProducedListener.accept(sstables);
     }
 
+    /**
+     * Rebuilds the Bloom filter of the sstable just produced by {@code txnWriter} and overwrites its {@code Filter.db}.
+     * <p>
+     * The underlying SSTable writer sizes the filter from a hardcoded partition count of {@code 0} (this writer
+     * cannot know the count up front), which yields a negligibly small, saturated filter (see CASSANDRA-21423). Once the
+     * sstable is fully written its true partition count is known, so we rebuild the filter from the on-disk index here.
+     * <p>
+     * If the sstables were opened (see {@link #shouldOpenSSTables()}), the in-memory readers still hold the original
+     * (broken) filter, so each is replaced with a clone carrying the rebuilt filter and the original is released.
+     *
+     * @param txnWriter the writer that produced the sstable
+     * @param produced  the readers returned by {@code finish} (may contain {@code null} entries when sstables were not opened)
+     * @return the readers to hand to the produced-sstable listener, with any opened readers carrying the rebuilt filter
+     */
+    protected Collection<SSTableReader> rebuildBloomFilter(SSTableTxnWriter txnWriter, Collection<SSTableReader> produced) throws IOException
+    {
+        Descriptor descriptor = SSTable.tryDescriptorFromFile(new File(txnWriter.getFilename()));
+        if (descriptor == null)
+            return produced;
+
+        TableMetrics metrics = owner != null ? owner.getMetrics() : null;
+        SSTableReaderLoadingBuilder<?, ?> loadingBuilder = descriptor.getFormat()
+                                                                     .getReaderFactory()
+                                                                     .loadingBuilder(descriptor, metadata, null);
+        IFilter filter = loadingBuilder.rebuildFilter(metrics);
+        if (filter == null || produced == null)
+            return produced;
+
+        try
+        {
+            List<SSTableReader> result = new ArrayList<>(produced.size());
+            for (SSTableReader reader : produced)
+            {
+                if (reader instanceof SSTableReaderWithFilter)
+                {
+                    // We rebuilt the filter for exactly one descriptor (derived from txnWriter.getFilename()); these
+                    // writers produce a single sstable per finish(). Guard against ever applying it to a reader for a
+                    // different sstable, which would install the wrong filter.
+                    Preconditions.checkState(descriptor.equals(reader.descriptor),
+                                             "Rebuilt bloom filter for %s but produced reader is for %s",
+                                             descriptor, reader.descriptor);
+                    result.add(((SSTableReaderWithFilter) reader).cloneAndReplace(filter.sharedCopy()));
+                    reader.selfRef().release();
+                }
+                else
+                {
+                    // sstable was not opened (null) - the on-disk filter is already corrected, nothing to swap
+                    result.add(reader);
+                }
+            }
+            return result;
+        }
+        finally
+        {
+            filter.close();
+        }
+    }
+
     protected SSTableTxnWriter createWriter(SSTable.Owner owner) throws IOException
     {
         SerializationHeader header = new SerializationHeader(true, metadata.get(), columns, EncodingStats.NO_STATS);

src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java

@@ -230,7 +230,7 @@ public void run()
                         for (Map.Entry<DecoratedKey, PartitionUpdate.Builder> entry : b.entrySet())
                             writer.append(entry.getValue().build().unfilteredIterator());
                         Collection<SSTableReader> finished = writer.finish(shouldOpenSSTables());
-                        notifySSTableProduced(finished);
+                        notifySSTableProduced(rebuildBloomFilter(writer, finished));
                     }
                 }
                 catch (Throwable e)

src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java

@@ -161,7 +161,7 @@ private void maybeCloseWriter(SSTableTxnWriter writer)
                 return;
 
             Collection<SSTableReader> finished = writer.finish(shouldOpenSSTables());
-            notifySSTableProduced(finished);
+            notifySSTableProduced(rebuildBloomFilter(writer, finished));
         }
         catch (Throwable t)
         {

src/java/org/apache/cassandra/io/sstable/format/SSTableReaderLoadingBuilder.java

@@ -35,12 +35,15 @@
 import org.apache.cassandra.io.sstable.KeyReader;
 import org.apache.cassandra.io.sstable.SSTable;
 import org.apache.cassandra.io.sstable.format.SSTableFormat.Components;
+import org.apache.cassandra.io.sstable.metadata.MetadataType;
 import org.apache.cassandra.io.sstable.metadata.ValidationMetadata;
 import org.apache.cassandra.metrics.TableMetrics;
 import org.apache.cassandra.schema.Schema;
 import org.apache.cassandra.schema.TableMetadata;
 import org.apache.cassandra.schema.TableMetadataRef;
 import org.apache.cassandra.utils.Clock;
+import org.apache.cassandra.utils.FilterFactory;
+import org.apache.cassandra.utils.IFilter;
 import org.apache.cassandra.utils.JVMStabilityInspector;
 
 import static com.google.common.base.Preconditions.checkArgument;
@@ -115,6 +118,54 @@ public R build(SSTable.Owner owner, boolean validate, boolean online)
 
     public abstract KeyReader buildKeyReader(TableMetrics tableMetrics) throws IOException;
 
+    /**
+     * Rebuilds the Bloom filter ({@code Filter.db}) of a fully written sstable by iterating its on-disk index and
+     * sizing the filter from the partition count recorded in {@code Statistics.db}. Any existing {@code Filter.db} is
+     * overwritten (and deleted if the write fails).
+     * <p>
+     * This is used to correct filters emitted by writers that cannot estimate the partition count up front - notably
+     * {@link org.apache.cassandra.io.sstable.CQLSSTableWriter}, which sizes its filter from a hardcoded count of {@code 0}
+     * and therefore produces a negligibly small, saturated filter (see CASSANDRA-21423).
+     *
+     * @param tableMetrics metrics to associate with the transient key reader, may be {@code null}
+     * @return the rebuilt filter, of which the caller takes ownership and must close, or {@code null} if the table has
+     * Bloom filtering disabled ({@code bloom_filter_fp_chance == 1.0}) and therefore has no filter component
+     * @throws IOException if the index cannot be read or the filter cannot be saved
+     */
+    public IFilter rebuildFilter(TableMetrics tableMetrics) throws IOException
+    {
+        TableMetadata metadata = tableMetadataRef.getLocal();
+        double fpChance = metadata.params.bloomFilterFpChance;
+        if (!FilterComponent.shouldUseBloomFilter(fpChance))
+            return null;
+
+        StatsComponent statsComponent = StatsComponent.load(descriptor, MetadataType.STATS);
+        long numKeys = statsComponent.statsMetadata().totalRows;
+
+        IFilter bf = null;
+        try (KeyReader keyReader = buildKeyReader(tableMetrics))
+        {
+            bf = FilterFactory.getFilter(numKeys, fpChance);
+            while (!keyReader.isExhausted())
+            {
+                bf.add(metadata.partitioner.decorateKey(keyReader.key()));
+                keyReader.advance();
+            }
+            FilterComponent.save(bf, descriptor, true);
+            return bf;
+        }
+        catch (IOException | RuntimeException | Error ex)
+        {
+            if (bf != null)
+                bf.close();
+            // A failure during the index walk (before save) would otherwise leave the original broken Filter.db in
+            // place. Remove it so the sstable has no filter rather than an invalid one - matching the deleteOnFailure
+            // path in FilterComponent.save (uses the same fileFor(Components.FILTER)).
+            descriptor.fileFor(Components.FILTER).deleteIfExists();
+            throw ex;
+        }
+    }
+
     protected abstract void openComponents(B builder, SSTable.Owner owner, boolean validate, boolean online) throws IOException;
 
     /**

test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java

@@ -203,6 +203,109 @@ private void testWritingSstableWithFormat(SSTableFormat<?, ?> format) throws Exc
         }
     }
 
+    @Test
+    public void testBloomFilterRebuiltBig() throws Exception
+    {
+        assertBloomFilterRebuilt(BigFormat.getInstance());
+    }
+
+    @Test
+    public void testBloomFilterRebuiltBti() throws Exception
+    {
+        assertBloomFilterRebuilt(new BtiFormat.BtiFormatFactory().getInstance(Collections.emptyMap()));
+    }
+
+    /**
+     * Before CASSANDRA-21423 the writer sized the Bloom filter from a hardcoded partition count of {@code 0}, producing
+     * a negligibly small (~16 byte) saturated filter regardless of the number of keys. The filter must now be rebuilt
+     * from the on-disk index and sized for the real key count.
+     */
+    private void assertBloomFilterRebuilt(SSTableFormat<?, ?> format) throws Exception
+    {
+        String schema = "CREATE TABLE " + qualifiedTable + " (k int PRIMARY KEY, v int)";
+        String insert = "INSERT INTO " + qualifiedTable + " (k, v) VALUES (?, ?)";
+        try (CQLSSTableWriter writer = CQLSSTableWriter.builder()
+                                                       .inDirectory(dataDir)
+                                                       .forTable(schema)
+                                                       .withFormat(format)
+                                                       .using(insert)
+                                                       .build())
+        {
+            for (int i = 0; i < 10_000; i++)
+                writer.addRow(i, i);
+        }
+
+        List<Path> filterFiles = filterComponentFiles();
+        assertFalse("no Filter.db was produced", filterFiles.isEmpty());
+        for (Path filterFile : filterFiles)
+            assertThat(Files.size(filterFile)).isGreaterThan(1024L);
+    }
+
+    @Test
+    public void testBloomFilterRebuildSkippedWhenDisabled() throws Exception
+    {
+        String schema = "CREATE TABLE " + qualifiedTable + " (k int PRIMARY KEY, v int) "
+                        + "WITH bloom_filter_fp_chance = 1.0";
+        String insert = "INSERT INTO " + qualifiedTable + " (k, v) VALUES (?, ?)";
+        try (CQLSSTableWriter writer = CQLSSTableWriter.builder()
+                                                       .inDirectory(dataDir)
+                                                       .forTable(schema)
+                                                       .using(insert)
+                                                       .build())
+        {
+            for (int i = 0; i < 1_000; i++)
+                writer.addRow(i, i);
+        }
+
+        // With Bloom filtering disabled the writer emits no Filter.db and the rebuild must be a no-op.
+        assertThat(filterComponentFiles()).isEmpty();
+    }
+
+    /**
+     * A single writer session can roll over into several sstables when {@code maxSSTableSizeInMiB} is exceeded. Each
+     * produced sstable is finished (and its filter rebuilt) independently, so verify that <em>every</em> emitted
+     * Filter.db is correctly sized for its own partition count - not just the first one.
+     */
+    @Test
+    public void testBloomFilterRebuiltForEachRotatedSSTable() throws Exception
+    {
+        String schema = "CREATE TABLE " + qualifiedTable + " (k int PRIMARY KEY, v text)";
+        String insert = "INSERT INTO " + qualifiedTable + " (k, v) VALUES (?, ?)";
+        // ~512 bytes of payload per row so that a 1 MiB size cap is crossed many times, forcing several rotations.
+        StringBuilder payload = new StringBuilder();
+        for (int i = 0; i < 512; i++)
+            payload.append('x');
+        String value = payload.toString();
+
+        try (CQLSSTableWriter writer = CQLSSTableWriter.builder()
+                                                       .inDirectory(dataDir)
+                                                       .forTable(schema)
+                                                       .using(insert)
+                                                       .withMaxSSTableSizeInMiB(1)
+                                                       .build())
+        {
+            for (int i = 0; i < 40_000; i++)
+                writer.addRow(i, value);
+        }
+
+        List<Path> filterFiles = filterComponentFiles();
+        assertThat(filterFiles.size())
+                .as("expected the writer to roll over into multiple sstables")
+                .isGreaterThan(1);
+        for (Path filterFile : filterFiles)
+            assertThat(Files.size(filterFile))
+                    .as("Filter.db %s should be sized for its partition count, not the broken ~16 byte filter", filterFile)
+                    .isGreaterThan(1024L);
+    }
+
+    private List<Path> filterComponentFiles() throws IOException
+    {
+        try (Stream<Path> paths = Files.list(dataDir.toPath()))
+        {
+            return paths.filter(p -> p.toString().endsWith("Filter.db")).collect(Collectors.toList());
+        }
+    }
+
     @Test
     public void testCompressedWriteAndReadBack() throws Exception
     {