CNDB-13446: Make search-then-sort BM25 queries apply BM25 filter (#1659)

### What is the issue
Fixes riptano/cndb#13446

### What does this PR fix and why was it fixed
This pr fixes an issue in the search-then-sort path where we never apply
the query term filter. The change is quite simple because we were
already scoring these terms, so we just need to filter out the ones that
do not have frequency of at least 1 for each query term.

Author: michaeljmarshall

src/java/org/apache/cassandra/index/sai/disk/v1/InvertedIndexSearcher.java

@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -267,7 +268,10 @@ public CloseableIterator<PrimaryKeyWithSortKey> orderResultsBy(SSTableReader rea
             documentFrequencies.put(term, matches);
         }
         var analyzer = indexContext.getAnalyzerFactory().create();
-        var it = keys.stream().map(pk -> DocTF.createFromDocument(pk, readColumn(sstable, pk), analyzer, queryTerms)).iterator();
+        var it = keys.stream()
+                     .map(pk -> DocTF.createFromDocument(pk, readColumn(sstable, pk), analyzer, queryTerms))
+                     .filter(Objects::nonNull)
+                     .iterator();
         return bm25Internal(CloseableIterator.wrap(it), queryTerms, documentFrequencies);
     }
 

src/java/org/apache/cassandra/index/sai/memory/TrieMemtableIndex.java

@@ -33,6 +33,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterators;
+import com.google.common.collect.Streams;
 import com.google.common.util.concurrent.Runnables;
 
 import org.apache.cassandra.cql3.Operator;
@@ -322,14 +323,10 @@ public List<CloseableIterator<PrimaryKeyWithSortKey>> orderBy(QueryContext query
         // Compute BM25 scores
         var docStats = computeDocumentFrequencies(queryContext, queryTerms);
         var analyzer = indexContext.getAnalyzerFactory().create();
-        var it = Iterators.filter(Iterators.transform(intersectedIterator, pk ->
-        {
-            Cell<?> cellForKey = getCellForKey(pk);
-            if (cellForKey == null)
-                // skip deleted rows
-                return null;
-            return BM25Utils.DocTF.createFromDocument(pk, cellForKey, analyzer, queryTerms);
-        }), Objects::nonNull);
+        var it = Streams.stream(intersectedIterator)
+                 .map(pk -> BM25Utils.DocTF.createFromDocument(pk, getCellForKey(pk), analyzer, queryTerms))
+                 .filter(Objects::nonNull)
+                 .iterator();
 
         return List.of(BM25Utils.computeScores(CloseableIterator.wrap(it),
                                                queryTerms,
@@ -395,13 +392,10 @@ public CloseableIterator<PrimaryKeyWithSortKey> orderResultsBy(QueryContext quer
         var analyzer = indexContext.getAnalyzerFactory().create();
         var queryTerms = orderer.getQueryTerms();
         var docStats = computeDocumentFrequencies(queryContext, queryTerms);
-        var it = keys.stream().map(pk -> {
-            Cell<?> cellForKey = getCellForKey(pk);
-            if (cellForKey == null)
-                // skip deleted rows
-                return null;
-            return BM25Utils.DocTF.createFromDocument(pk, cellForKey, analyzer, queryTerms);
-        }).filter(Objects::nonNull).iterator();
+        var it = keys.stream()
+                     .map(pk -> BM25Utils.DocTF.createFromDocument(pk, getCellForKey(pk), analyzer, queryTerms))
+                     .filter(Objects::nonNull)
+                     .iterator();
         return BM25Utils.computeScores(CloseableIterator.wrap(it),
                                        queryTerms,
                                        docStats,

src/java/org/apache/cassandra/index/sai/utils/BM25Utils.java

@@ -27,6 +27,8 @@
 import java.util.List;
 import java.util.Map;
 
+import javax.annotation.Nullable;
+
 import org.apache.cassandra.db.memtable.Memtable;
 import org.apache.cassandra.db.rows.Cell;
 import org.apache.cassandra.index.sai.IndexContext;
@@ -78,12 +80,14 @@ public int getTermFrequency(ByteBuffer term)
             return frequencies.getOrDefault(term, 0);
         }
 
+        @Nullable
         public static DocTF createFromDocument(PrimaryKey pk,
-                                             Cell<?> cell,
-                                             AbstractAnalyzer docAnalyzer,
-                                             Collection<ByteBuffer> queryTerms)
+                                               Cell<?> cell,
+                                               AbstractAnalyzer docAnalyzer,
+                                               Collection<ByteBuffer> queryTerms)
         {
-            assert cell != null : "Cannot find a cell for pk " + pk;
+            if (cell == null)
+                return null;
 
             int count = 0;
             Map<ByteBuffer, Integer> frequencies = new HashMap<>();
@@ -103,6 +107,10 @@ public static DocTF createFromDocument(PrimaryKey pk,
                 docAnalyzer.end();
             }
 
+            // Every query term must be present in the document
+            if (queryTerms.size() > frequencies.size())
+                return null;
+
             return new DocTF(pk, count, frequencies);
         }
     }

test/unit/org/apache/cassandra/index/sai/cql/BM25Test.java

@@ -416,7 +416,12 @@ private void createSimpleTable()
 
     private String createAnalyzedIndex()
     {
-        return createIndex("CREATE CUSTOM INDEX ON %s(v) " +
+        return createAnalyzedIndex("v");
+    }
+
+    private String createAnalyzedIndex(String column)
+    {
+        return createIndex("CREATE CUSTOM INDEX ON %s(" + column + ") " +
                            "USING 'org.apache.cassandra.index.sai.StorageAttachedIndex' " +
                            "WITH OPTIONS = {" +
                            "'index_analyzer': '{" +
@@ -628,4 +633,63 @@ public void cannotHaveAggregationOnBM25Query()
                 .isInstanceOf(InvalidRequestException.class)
                 .hasMessage(SelectStatement.TOPK_AGGREGATION_ERROR);
     }
+
+    @Test
+    public void testBM25andFilterz() throws Throwable
+    {
+        createTable("CREATE TABLE %s (id int PRIMARY KEY, category text, score int, title text, body text)");
+        createAnalyzedIndex("body");
+        createIndex("CREATE CUSTOM INDEX ON %s (score) USING 'StorageAttachedIndex'");
+        insertArticle();
+        beforeAndAfterFlush(
+                () -> {
+                    // 10 docs have score 3 and 3 of those have "health"
+                    var result = execute("SELECT * FROM %s WHERE score = 3 ORDER BY body BM25 OF ? LIMIT 10",
+                                         "health");
+                    assertThat(result).hasSize(3);
+
+                    // 4 docs have score 2 and one of those has "discussed"
+                    result = execute("SELECT * FROM %s WHERE score = 2 ORDER BY body BM25 OF ? LIMIT 10",
+                                         "discussed");
+                    assertThat(result).hasSize(1);
+                });
+    }
+
+    private void insertArticle()
+    {
+        Object[][] dataset = {
+        { 1, "Climate", 5, "Climate change is a pressing issue. Climate patterns are shifting globally. Scientists study climate data daily." },
+        { 2, "Technology", 3, "Technology is advancing. New technology in AI and robotics is groundbreaking." },
+        { 3, "Economy", 4, "The economy is recovering. Economy experts are optimistic. However, the global economy still faces risks." },
+        { 4, "Health", 3, "Health is wealth. Health policies need to be improved to ensure better public health outcomes." },
+        { 5, "Education", 2, "Education is the foundation of success. Online education is booming." },
+        { 6, "Climate", 4, "Climate and health are closely linked. Climate affects air quality and health outcomes." },
+        { 7, "Education", 3, "Technology and education go hand in hand. EdTech is revolutionizing education through technology." },
+        { 8, "Economy", 3, "The global economy is influenced by technology. Fintech is a key part of the economy today." },
+        { 9, "Health", 3, "Education and health programs must be prioritized. Health education is vital in schools." },
+        { 10, "Mixed", 3, "Technology, economy, and education are pillars of development." },
+        { 11, "Climate", 5, "Climate climate climate. It's everywhere. Climate drives political and economic decisions." },
+        { 12, "Health", 2, "Health concerns rise with climate issues. Health organizations are sounding the alarm." },
+        { 13, "Economy", 3, "The economy is fluctuating. Uncertainty looms over the economy." },
+        { 14, "Health", 3, "Cutting-edge technology is transforming healthcare. Healthtech merges health and technology." },
+        { 15, "Education", 2, "Education reforms are underway. Education experts suggest holistic changes." },
+        { 16, "Climate", 4, "Climate affects the economy and health. Climate events cost billions annually." },
+        { 17, "Technology", 3, "Technology is the backbone of the modern economy. Without technology, economic growth stagnates." },
+        { 18, "Health", 2, "Health is discussed less than economy or climate, but health matters deeply." },
+        { 19, "Climate", 5, "Climate change, climate policies, climate research—climate is the buzzword of our time." },
+        { 20, "Mixed", 3, "Investments in education and technology will shape the future of the global economy." }
+        };
+
+        for (Object[] article : dataset)
+        {
+            execute(
+            "INSERT INTO %s (id, category, score, body) VALUES (?, ?, ?, ?)",
+            article[0],
+            article[1],
+            article[2],
+            article[3]
+            );
+        }
+    }
+
 }