CNDB-14343: Improve tests

Author: adelapena

src/java/org/apache/cassandra/index/sai/disk/vector/VectorMemtableIndex.java

@@ -318,6 +318,8 @@ public CloseableIterator<PrimaryKeyWithSortKey> orderResultsBy(QueryContext cont
 
         keysInRange.forEach(k ->
         {
+            // if the indexed column is static, we need to get the static row associated with the non-static row that
+            // might be referenced by the key
             if (isStatic)
                 k = k.forStaticRow();
 

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

@@ -134,26 +134,40 @@ public void testPartitionKeyComponentWithClustering() throws Throwable
     {
         createTable("CREATE TABLE %s (k1 int, k2 vector<float, 2>, c int, v int, PRIMARY KEY((k1, k2), c))");
         createIndex("CREATE CUSTOM INDEX ON %s(k2) USING 'StorageAttachedIndex'");
+        createIndex("CREATE CUSTOM INDEX ON %s(v) USING 'StorageAttachedIndex'");
 
-        String insert = "INSERT INTO %s (k1, k2, c) VALUES (?, ?, ?)";
-        execute(insert, row(1, vector(0.1f, 0.1f), 0));
-        execute(insert, row(2, vector(0.1f, 0.2f), 0));
-        execute(insert, row(3, vector(0.1f, 0.3f), 0));
-        execute(insert, row(4, vector(0.1f, 0.4f), 0));
+        // insert static rows and non-static rows all at once
+        String insert = "INSERT INTO %s (k1, k2, c, v) VALUES (?, ?, ?, ?)";
+        execute(insert, row(1, vector(0.1f, 0.1f), 10, 0));
+        execute(insert, row(2, vector(0.1f, 0.2f), 20, 0));
+        execute(insert, row(3, vector(0.1f, 0.3f), 30, 1));
+        execute(insert, row(4, vector(0.1f, 0.4f), 40, 1));
         beforeAndAfterFlush(() -> {
+            // query with ANN only
             assertRows(execute("SELECT k1 FROM %s ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 10"), row(1), row(2), row(3), row(4));
             assertRows(execute("SELECT k1 FROM %s ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 2"), row(1), row(2));
+
+            // query with hybrid search
+            assertRows(execute("SELECT k1 FROM %s WHERE v>=0 ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 10"), row(1), row(2), row(3), row(4));
+            assertRows(execute("SELECT k1 FROM %s WHERE v=0 ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 10"), row(1), row(2));
+            assertRows(execute("SELECT k1 FROM %s WHERE v>0 ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 10"), row(3), row(4));
         });
 
         // test again with a different order
         execute("TRUNCATE TABLE %s");
-        execute(insert, row(1, vector(0.1f, 0.4f), 0));
-        execute(insert, row(2, vector(0.1f, 0.3f), 0));
-        execute(insert, row(3, vector(0.1f, 0.2f), 0));
-        execute(insert, row(4, vector(0.1f, 0.1f), 0));
+        execute(insert, row(1, vector(0.1f, 0.4f), 10, 0));
+        execute(insert, row(2, vector(0.1f, 0.3f), 20, 0));
+        execute(insert, row(3, vector(0.1f, 0.2f), 30, 1));
+        execute(insert, row(4, vector(0.1f, 0.1f), 40, 1));
         beforeAndAfterFlush(() -> {
+            // query with ANN only
             assertRows(execute("SELECT k1 FROM %s ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 10"), row(4), row(3), row(2), row(1));
             assertRows(execute("SELECT k1 FROM %s ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 2"), row(4), row(3));
+
+            // query with hybrid search
+            assertRows(execute("SELECT k1 FROM %s WHERE v>=0 ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 10"), row(4), row(3), row(2), row(1));
+            assertRows(execute("SELECT k1 FROM %s WHERE v=0 ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 10"), row(2), row(1));
+            assertRows(execute("SELECT k1 FROM %s WHERE v>0 ORDER BY k2 ANN OF [0.1, 0.1] LIMIT 10"), row(4), row(3));
         });
     }