CNDB-14343: Fix ANN queries on primary key columns

Author: adelapena

src/java/org/apache/cassandra/index/sai/plan/StorageAttachedIndexSearcher.java

@@ -666,19 +666,19 @@ public UnfilteredRowIterator readAndValidatePartition(PrimaryKey pk, List<Primar
             {
                 queryContext.addPartitionsRead(1);
                 queryContext.checkpoint();
-                var staticRow = partition.staticRow();
                 UnfilteredRowIterator clusters = applyIndexFilter(partition, filterTree, queryContext);
 
-                if (clusters == null || !clusters.hasNext())
-                {
-                    processedKeys.add(pk);
+                if (clusters == null)
                     return null;
-                }
 
+                if (!clusters.hasNext())
+                    throw new UnsupportedOperationException("Static columns are not supported");
+
+                var row = clusters.next();
+                var staticRow = partition.staticRow();
                 var now = FBUtilities.nowInSeconds();
                 boolean isRowValid = false;
-                var row = clusters.next();
-                assert !clusters.hasNext() : "Expected only one row per partition";
+
                 if (!row.isRangeTombstoneMarker())
                 {
                     for (PrimaryKeyWithSortKey sourceKey : sourceKeys)

src/java/org/apache/cassandra/index/sai/plan/TopKProcessor.java

@@ -303,7 +303,7 @@ private float getScoreForRow(DecoratedKey key, Row row)
     {
         ColumnMetadata column = indexContext.getDefinition();
 
-        if (column.isPrimaryKeyColumn() && key == null)
+        if (column.isPartitionKey() && key == null)
             return 0;
 
         if (column.isStatic() && !row.isStatic())

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

@@ -28,6 +28,7 @@
 import org.apache.cassandra.dht.Token;
 import org.apache.cassandra.index.sai.IndexContext;
 import org.apache.cassandra.io.sstable.SSTableId;
+import org.apache.cassandra.schema.ColumnMetadata;
 import org.apache.cassandra.utils.bytecomparable.ByteComparable;
 import org.apache.cassandra.utils.bytecomparable.ByteSource;
 
@@ -51,7 +52,7 @@ protected PrimaryKeyWithSortKey(IndexContext context, Memtable sourceTable, Prim
         this.primaryKey = primaryKey;
     }
 
-    protected PrimaryKeyWithSortKey(IndexContext context, SSTableId sourceTable, PrimaryKey primaryKey)
+    protected PrimaryKeyWithSortKey(IndexContext context, SSTableId<?> sourceTable, PrimaryKey primaryKey)
     {
         this.context = context;
         this.sourceTable = sourceTable;
@@ -65,10 +66,14 @@ public PrimaryKey primaryKey()
 
     public boolean isIndexDataValid(Row row, int nowInSecs)
     {
-        assert context.getDefinition().isRegular() : "Only regular columns are supported, got " + context.getDefinition();
-        var cell = row.getCell(context.getDefinition());
+        ColumnMetadata column = context.getDefinition();
+        if (!column.isRegular())
+            return true;
+
+        var cell = row.getCell(column);
         if (!cell.isLive(nowInSecs))
             return false;
+
         assert cell instanceof CellWithSourceTable : "Expected CellWithSource, got " + cell.getClass();
         return sourceTable.equals(((CellWithSourceTable<?>) cell).sourceTable())
                && isIndexDataEqualToLiveData(cell.buffer());