CNDB-13408 add BM25 tests to compare with CNDB (#1663)

k-rus · web-flow · commit cb3de82121b4 · 2025-04-02T17:05:04.000+02:00
Adds CC's unit tests to help with
testing BM25  in integration CNDB's tests and 
expect the same result.

Also adds tests to reproduce minor issues 
with error messages on BM25.

Fixes few warnings in the affected file.
diff --git a/test/unit/org/apache/cassandra/index/sai/cql/BM25Test.java b/test/unit/org/apache/cassandra/index/sai/cql/BM25Test.java
@@ -1,13 +1,11 @@
 /*
- * 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
+ * Copyright DataStax, Inc.
  *
- *     http://www.apache.org/licenses/LICENSE-2.0
+ * Licensed 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,
@@ -19,10 +17,16 @@
 package org.apache.cassandra.index.sai.cql;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.stream.Collectors;
 
+import org.assertj.core.api.Assertions;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -124,15 +128,15 @@ public void testTwoIndexesAmbiguousPredicate() throws Throwable
         // be rejected
         beforeAndAfterFlush(() -> {
             // Single predicate
-            assertInvalidMessage(String.format(EQ_AMBIGUOUS_ERROR, "v", getIndex(0), getIndex(1)),
+            assertInvalidMessage(String.format(EQ_AMBIGUOUS_ERROR, 'v', getIndex(0), getIndex(1)),
                                  "SELECT k FROM %s WHERE v = 'apple'");
 
             // AND
-            assertInvalidMessage(String.format(EQ_AMBIGUOUS_ERROR, "v", getIndex(0), getIndex(1)),
+            assertInvalidMessage(String.format(EQ_AMBIGUOUS_ERROR, 'v', getIndex(0), getIndex(1)),
                                  "SELECT k FROM %s WHERE v = 'apple' AND v : 'juice'");
 
             // OR
-            assertInvalidMessage(String.format(EQ_AMBIGUOUS_ERROR, "v", getIndex(0), getIndex(1)),
+            assertInvalidMessage(String.format(EQ_AMBIGUOUS_ERROR, 'v', getIndex(0), getIndex(1)),
                                  "SELECT k FROM %s WHERE v = 'apple' OR v : 'juice'");
         });
     }
@@ -178,14 +182,7 @@ public void testComplexQueriesWithMultipleIndexes() throws Throwable
 
         // Create mix of analyzed, unanalyzed, and non-text indexes
         createIndex("CREATE CUSTOM INDEX ON %s(v1) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex'");
-        createIndex("CREATE CUSTOM INDEX ON %s(v2) " +
-                    "USING 'org.apache.cassandra.index.sai.StorageAttachedIndex' " +
-                    "WITH OPTIONS = {" +
-                    "'index_analyzer': '{" +
-                    "\"tokenizer\" : {\"name\" : \"standard\"}, " +
-                    "\"filters\" : [{\"name\" : \"porterstem\"}]" +
-                    "}'" +
-                    "}");
+        createAnalyzedIndex("v2");
         createIndex("CREATE CUSTOM INDEX ON %s(v3) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex'");
 
         execute("INSERT INTO %s (k, v1, v2, v3) VALUES (1, 'apple', 'orange juice', 5)");
@@ -263,10 +260,8 @@ public void testEmptyQuery() throws Throwable
         execute("INSERT INTO %s (k, v) VALUES (1, 'apple')");
 
         beforeAndAfterFlush(() ->
-                            {
-                                assertInvalidMessage("BM25 query must contain at least one term (perhaps your analyzer is discarding tokens you didn't expect)",
-                                                     "SELECT k FROM %s ORDER BY v BM25 OF '+' LIMIT 1");
-                            });
+                            assertInvalidMessage("BM25 query must contain at least one term (perhaps your analyzer is discarding tokens you didn't expect)",
+                                                 "SELECT k FROM %s ORDER BY v BM25 OF '+' LIMIT 1"));
     }
 
     @Test
@@ -420,14 +415,20 @@ private String createAnalyzedIndex()
     }
 
     private String createAnalyzedIndex(String column)
+    {
+        return createAnalyzedIndex(column, false);
+    }
+
+    private String createAnalyzedIndex(String column, boolean lowercase)
     {
         return createIndex("CREATE CUSTOM INDEX ON %s(" + column + ") " +
                            "USING 'org.apache.cassandra.index.sai.StorageAttachedIndex' " +
                            "WITH OPTIONS = {" +
                            "'index_analyzer': '{" +
                            "\"tokenizer\" : {\"name\" : \"standard\"}, " +
-                           "\"filters\" : [{\"name\" : \"porterstem\"}]" +
-                           "}'}"
+                           "\"filters\" : [{\"name\" : \"porterstem\"}" +
+                           (lowercase ? ", {\"name\" : \"lowercase\"}]" : "]")
+                           + "}'}"
         );
     }
 
@@ -640,7 +641,7 @@ 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();
+        insertPrimitiveData();
         beforeAndAfterFlush(
                 () -> {
                     // 10 docs have score 3 and 3 of those have "health"
@@ -655,41 +656,158 @@ public void testBM25andFilterz() throws Throwable
                 });
     }
 
-    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)
+    @Test
+    public void testErrorMessages()
+    {
+        createTable("CREATE TABLE %s (id int PRIMARY KEY, category text, score int, " +
+                    "title text, body text, bodyset set<text>, " +
+                    "map_category map<int, text>, map_body map<text, text>)");
+        createAnalyzedIndex("body", true);
+        createAnalyzedIndex("bodyset", true);
+        createAnalyzedIndex("map_body", true);
+
+        // Improve message issue CNDB-13514
+        assertInvalidMessage("BM25 ordering on column bodyset requires an analyzed index",
+                             "SELECT * FROM %s ORDER BY bodyset BM25 OF ? LIMIT 10");
+
+        // Discussion of message incosistency CNDB-13526
+        assertInvalidMessage("Ordering on non-clustering column requires each restricted column to be indexed except for fully-specified partition keys",
+                             "SELECT * FROM %s WHERE map_body CONTAINS KEY 'Climate' ORDER BY body BM25 OF ? LIMIT 10");
+    }
+
+    @Test
+    public void testCollections() throws Throwable
+    {
+        createTable("CREATE TABLE %s (id int PRIMARY KEY, category text, score int, tie int," +
+                    "title text, body text, bodyset set<text>, " +
+                    "map_category map<int, text>, map_body map<text, text>)");
+        createAnalyzedIndex("body", true);
+        createAnalyzedIndex("bodyset", true);
+        createAnalyzedIndex("map_body", true);
+        createIndex("CREATE CUSTOM INDEX ON %s (score) USING 'StorageAttachedIndex'");
+        createIndex("CREATE CUSTOM INDEX ON %s (category) USING 'StorageAttachedIndex'");
+        createIndex("CREATE CUSTOM INDEX ON %s (tie) USING 'StorageAttachedIndex'");
+        createIndex("CREATE CUSTOM INDEX ON %s (map_category) USING 'StorageAttachedIndex'");
+        createIndex("CREATE CUSTOM INDEX ON %s (KEYS(map_body)) USING 'StorageAttachedIndex'");
+        insertCollectionData();
+
+        beforeAndAfterFlush(
+        () -> {
+            executeQuery(Arrays.asList(11, 1, 16, 18), "SELECT * FROM %s WHERE tie = 1 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+            executeQuery(Arrays.asList(11, 1), "SELECT * FROM %s WHERE score = 5 AND tie = 1 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+            executeQuery(Arrays.asList(6, 16), "SELECT * FROM %s WHERE score > 3 ORDER BY body BM25 OF ? LIMIT 10",
+                         "health");
+            executeQuery(Arrays.asList(4, 18, 14), "SELECT * FROM %s WHERE category = 'Health' AND tie = 1 " +
+                                                          "ORDER BY body BM25 OF ? LIMIT 10",
+                         "Health");
+            executeQuery(Arrays.asList(4, 18, 14), "SELECT * FROM %s WHERE score <= 3 AND tie = 1 AND category = 'Health' " +
+                                                          "ORDER BY body BM25 OF ? LIMIT 10",
+                         "health");
+            executeQuery(Arrays.asList(11, 1, 16, 18), "SELECT * FROM %s WHERE bodyset CONTAINS 'climate' AND tie <= 1 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+            executeQuery(Arrays.asList(6, 12), "SELECT * FROM %s WHERE bodyset CONTAINS 'health' AND tie > 1 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+            executeQuery(Arrays.asList(11, 1, 16, 18), "SELECT * FROM %s WHERE map_category CONTAINS 'Climate' AND tie <= 1 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+            executeQuery(Arrays.asList(19, 6, 12), "SELECT * FROM %s WHERE map_category CONTAINS 'Health' AND tie > 1 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+            executeQuery(Arrays.asList(11, 1, 16, 18), "SELECT * FROM %s WHERE map_body CONTAINS 'Climate' AND tie <= 1 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+            executeQuery(Arrays.asList(11, 16, 18), "SELECT * FROM %s WHERE map_body CONTAINS 'health' AND tie < 2 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+            executeQuery(Arrays.asList(19, 6, 12), "SELECT * FROM %s WHERE map_body CONTAINS KEY 'Health' AND tie >= 2 ORDER BY body BM25 OF ? LIMIT 10",
+                         "climate");
+        });
+    }
+
+    private final static Object[][] DATASET =
+    {
+    { 1, "Climate", 5, "Climate change is a pressing issue. Climate patterns are shifting globally. Scientists study climate data daily.", 1 },
+    { 2, "Technology", 3, "Technology is advancing. New technology in AI and robotics is groundbreaking.", 1 },
+    { 3, "Economy", 4, "The economy is recovering. Economy experts are optimistic. However, the global economy still faces risks.", 1 },
+    { 4, "Health", 3, "Health is wealth. Health policies need to be improved to ensure better public health outcomes.", 1 },
+    { 5, "Education", 2, "Education is the foundation of success. Online education is booming.", 4 },
+    { 6, "Climate", 4, "Climate and health are closely linked. Climate affects air quality and health outcomes.", 2 },
+    { 7, "Education", 3, "Technology and education go hand in hand. EdTech is revolutionizing education through technology.", 3 },
+    { 8, "Economy", 3, "The global economy is influenced by technology. Fintech is a key part of the economy today.", 2 },
+    { 9, "Health", 3, "Education and health programs must be prioritized. Health education is vital in schools.", 2 },
+    { 10, "Mixed", 3, "Technology, economy, and education are pillars of development.", 2 },
+    { 11, "Climate", 5, "Climate climate climate. It's everywhere. Climate drives political and economic decisions.", 1 },
+    { 12, "Health", 2, "Health concerns rise with climate issues. Health organizations are sounding the alarm.", 2 },
+    { 13, "Economy", 3, "The economy is fluctuating. Uncertainty looms over the economy.", 1 },
+    { 14, "Health", 3, "Cutting-edge technology is transforming healthcare. Healthtech merges health and technology.", 1 },
+    { 15, "Education", 2, "Education reforms are underway. Education experts suggest holistic changes.", 1 },
+    { 16, "Climate", 4, "Climate affects the economy and health. Climate events cost billions annually.", 1 },
+    { 17, "Technology", 3, "Technology is the backbone of the modern economy. Without technology, economic growth stagnates.", 2 },
+    { 18, "Health", 2, "Health is discussed less than economy or climate, but health matters deeply.", 1 },
+    { 19, "Climate", 5, "Climate change, climate policies, climate research—climate is the buzzword of our time.", 2 },
+    { 20, "Mixed", 3, "Investments in education and technology will shape the future of the global economy.", 1 }
+    };
+
+    private void insertPrimitiveData()
+    {
+        for (Object[] row : DATASET)
         {
             execute(
             "INSERT INTO %s (id, category, score, body) VALUES (?, ?, ?, ?)",
-            article[0],
-            article[1],
-            article[2],
-            article[3]
+            row[0],
+            row[1],
+            row[2],
+            row[3]
+            );
+        }
+    }
+
+    private void insertCollectionData()
+    {
+        int setsize = 1;
+        for (int row = 0; row < DATASET.length; row++)
+        {
+            var set = new HashSet<String>();
+            for (int j = 0; j < setsize; j++)
+                set.add((String) DATASET[row - j][3]);
+            if (setsize >= 3)
+                setsize -= 2;
+            else
+                setsize++;
+            var map = new HashMap<Integer, String>();
+            var map_text = new HashMap<String, String>();
+            for (int j = 0; j <= row && j < 3; j++)
+            {
+                map.putIfAbsent((Integer) DATASET[row - j][2], (String) DATASET[row - j][1]);
+                map_text.putIfAbsent((String) DATASET[row - j][1], (String) DATASET[row - j][3]);
+            }
+
+            execute(
+            "INSERT INTO %s (id, category, score, body, tie, bodyset, map_category, map_body) " +
+            "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+            DATASET[row][0],
+            DATASET[row][1],
+            DATASET[row][2],
+            DATASET[row][3],
+            DATASET[row][4],
+            set,
+            map,
+            map_text
             );
         }
     }
 
+    private void executeQuery(List<Integer> expected, String query, Object... values) throws Throwable
+    {
+        assertResult(execute(query, values), expected);
+        prepare(query);
+        assertResult(execute(query, values), expected);
+    }
+
+    private void assertResult(UntypedResultSet result, List<Integer> expected)
+    {
+        Assertions.assertThat(result).hasSize(expected.size());
+        var ids = result.stream()
+                        .map(row -> row.getInt("id"))
+                        .collect(Collectors.toList());
+        Assertions.assertThat(ids).isEqualTo(expected);
+    }
 }
