CNDB-13403: Add source column to synthetic column's ColumnMetadata (#1641)

### What is the issue
Fixes riptano/cndb#13403 by implementing one
of the solutions.

### What does this PR fix and why was it fixed

The current synthetic column logic relies on convention to know that the
synthetic column's source column is the SAI ORDER BY column. This seems
fragile and will break if we find a reason to introduce a new synthetic
column. I propose that we serialize the source column name so that we
can more easily support multiple synthetic columns.

The urgency for this change is that we want to get it in before we are
locked into a given serde protocol.

Note however, that I haven't solved the primary issue described in
riptano/cndb#13402.

CNDB tests riptano/cndb#13418

Author: michaeljmarshall

src/java/org/apache/cassandra/cql3/statements/SelectStatement.java

@@ -1282,7 +1282,7 @@ private Map<ColumnMetadata, Ordering> getScoreOrdering(List<Ordering> orderings)
 
             // Create synthetic score column
             ColumnMetadata sourceColumn = expr.getColumn();
-            var cm = ColumnMetadata.syntheticColumn(sourceColumn.ksName, sourceColumn.cfName, ColumnMetadata.SYNTHETIC_SCORE_ID, FloatType.instance);
+            var cm = ColumnMetadata.syntheticScoreColumn(sourceColumn, FloatType.instance);
             return Map.of(cm, orderings.get(0));
         }
 

src/java/org/apache/cassandra/db/Columns.java

@@ -487,6 +487,7 @@ public void serialize(Columns columns, DataOutputPlus out) throws IOException
                 if (column.isSynthetic())
                 {
                     ByteBufferUtil.writeWithVIntLength(column.name.bytes, out);
+                    ByteBufferUtil.writeWithVIntLength(column.sythenticSourceColumn.bytes, out);
                     typeSerializer.serialize(column.type, out);
                 }
             }
@@ -519,6 +520,7 @@ public long serializedSize(Columns columns)
                 {
                     syntheticCount++;
                     size += ByteBufferUtil.serializedSizeWithVIntLength(column.name.bytes);
+                    size += ByteBufferUtil.serializedSizeWithVIntLength(column.sythenticSourceColumn.bytes);
                     size += typeSerializer.serializedSize(column.type);
                 }
                 else
@@ -544,12 +546,22 @@ public Columns deserialize(DataInputPlus in, TableMetadata metadata) throws IOEx
                 for (int i = 0; i < syntheticCount; i++)
                 {
                     ByteBuffer name = ByteBufferUtil.readWithVIntLength(in);
+                    ByteBuffer sourceColumnName = ByteBufferUtil.readWithVIntLength(in);
                     AbstractType<?> type = typeSerializer.deserialize(in);
 
                     if (!name.equals(ColumnMetadata.SYNTHETIC_SCORE_ID.bytes))
                         throw new IllegalStateException("Unknown synthetic column " + UTF8Type.instance.getString(name));
 
-                    ColumnMetadata column = ColumnMetadata.syntheticColumn(metadata.keyspace, metadata.name, ColumnMetadata.SYNTHETIC_SCORE_ID, type);
+                    ColumnMetadata sourceColumn = metadata.getColumn(sourceColumnName);
+                    if (sourceColumn == null)
+                    {
+                        // If we don't find the definition, it could be we have data for a dropped column
+                        sourceColumn = metadata.getDroppedColumn(name);
+                        if (sourceColumn == null)
+                            throw new RuntimeException("Unknown column " + UTF8Type.instance.getString(name) + " during deserialization of " + metadata.keyspace + '.' + metadata.name);
+                    }
+
+                    ColumnMetadata column = ColumnMetadata.syntheticScoreColumn(sourceColumn, type);
                     builder.add(column);
                 }
 

src/java/org/apache/cassandra/db/ReadCommand.java

@@ -1062,7 +1062,8 @@ public ReadCommand deserialize(DataInputPlus in, int version) throws IOException
                 // synthetic columns sort first, so when we hit the first non-synthetic, we're done
                 if (!c.isSynthetic())
                     break;
-                tmb.addColumn(ColumnMetadata.syntheticColumn(c.ksName, c.cfName, c.name, c.type));
+                assert c.sythenticSourceColumn != null;
+                tmb.addColumn(c);
             }
             metadata = tmb.build();
 

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

@@ -43,7 +43,6 @@
 import org.apache.cassandra.db.PartitionPosition;
 import org.apache.cassandra.db.ReadCommand;
 import org.apache.cassandra.db.ReadExecutionController;
-import org.apache.cassandra.db.filter.ColumnFilter;
 import org.apache.cassandra.db.marshal.FloatType;
 import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
 import org.apache.cassandra.db.rows.AbstractUnfilteredRowIterator;
@@ -124,7 +123,8 @@ public UnfilteredPartitionIterator search(ReadExecutionController executionContr
                 assert !(keysIterator instanceof KeyRangeIterator);
                 var scoredKeysIterator = (CloseableIterator<PrimaryKeyWithSortKey>) keysIterator;
                 var result = new ScoreOrderedResultRetriever(scoredKeysIterator, filterTree, controller,
-                                                             executionController, queryContext, command.limits().count());
+                                                             executionController, queryContext, command.limits().count(),
+                                                             ordering.context.getDefinition());
                 return new TopKProcessor(command).filter(result);
             }
             catch (QueryView.Builder.MissingIndexException e)
@@ -479,6 +479,10 @@ public static class ScoreOrderedResultRetriever extends AbstractIterator<Unfilte
         private final HashSet<PrimaryKey> processedKeys;
         private final Queue<UnfilteredRowIterator> pendingRows;
 
+        // Null indicates we are not sending the synthetic score column to the coordinator
+        @Nullable
+        private final ColumnMetadata syntheticScoreColumn;
+
         // The limit requested by the query. We cannot load more than softLimit rows in bulk because we only want
         // to fetch the topk rows where k is the limit. However, we allow the iterator to fetch more rows than the
         // soft limit to avoid confusing behavior. When the softLimit is reached, the iterator will fetch one row
@@ -491,7 +495,8 @@ private ScoreOrderedResultRetriever(CloseableIterator<PrimaryKeyWithSortKey> sco
                                             QueryController controller,
                                             ReadExecutionController executionController,
                                             QueryContext queryContext,
-                                            int limit)
+                                            int limit,
+                                            ColumnMetadata orderedColumn)
         {
             IndexContext context = controller.getOrderer().context;
             this.view = controller.getQueryView(context).view;
@@ -507,6 +512,13 @@ private ScoreOrderedResultRetriever(CloseableIterator<PrimaryKeyWithSortKey> sco
             this.processedKeys = new HashSet<>(limit);
             this.pendingRows = new ArrayDeque<>(limit);
             this.softLimit = limit;
+
+            // When +score is added on the coordinator side, it's represented as a PrecomputedColumnFilter
+            // even in a 'SELECT *' because WCF is not capable of representing synthetic columns.
+            // This can be simplified when we remove ANN_USE_SYNTHETIC_SCORE
+            var tempColumn = ColumnMetadata.syntheticScoreColumn(orderedColumn, FloatType.instance);
+            var isScoreFetched = controller.command().columnFilter().fetchesExplicitly(tempColumn);
+            this.syntheticScoreColumn = isScoreFetched ? tempColumn : null;
         }
 
         @Override
@@ -663,7 +675,8 @@ public UnfilteredRowIterator readAndValidatePartition(PrimaryKey pk, List<Primar
                         }
                     }
                 }
-                return isRowValid ? new PrimaryKeyIterator(partition, staticRow, row, sourceKeys, controller.command()) : null;
+                return isRowValid ? new PrimaryKeyIterator(partition, staticRow, row, sourceKeys, syntheticScoreColumn)
+                                  : null;
             }
         }
 
@@ -684,7 +697,7 @@ public static class PrimaryKeyIterator extends AbstractUnfilteredRowIterator
             private boolean consumed = false;
             private final Unfiltered row;
 
-            public PrimaryKeyIterator(UnfilteredRowIterator partition, Row staticRow, Unfiltered content, List<PrimaryKeyWithSortKey> primaryKeysWithScore, ReadCommand command)
+            public PrimaryKeyIterator(UnfilteredRowIterator partition, Row staticRow, Unfiltered content, List<PrimaryKeyWithSortKey> primaryKeysWithScore, ColumnMetadata syntheticScoreColumn)
             {
                 super(partition.metadata(),
                       partition.partitionKey(),
@@ -702,16 +715,8 @@ public PrimaryKeyIterator(UnfilteredRowIterator partition, Row staticRow, Unfilt
                     return;
                 }
 
-                // When +score is added on the coordinator side, it's represented as a PrecomputedColumnFilter
-                // even in a 'SELECT *' because WCF is not capable of representing synthetic columns.
-                // This can be simplified when we remove ANN_USE_SYNTHETIC_SCORE
-                var tm = metadata();
-                var scoreColumn = ColumnMetadata.syntheticColumn(tm.keyspace,
-                                                                 tm.name,
-                                                                 ColumnMetadata.SYNTHETIC_SCORE_ID,
-                                                                 FloatType.instance);
-                var isScoreFetched = command.columnFilter().fetchesExplicitly(scoreColumn);
-                if (!isScoreFetched)
+
+                if (syntheticScoreColumn == null)
                 {
                     this.row = content;
                     return;
@@ -724,7 +729,7 @@ public PrimaryKeyIterator(UnfilteredRowIterator partition, Row staticRow, Unfilt
 
                 // inject +score as a new column
                 var pkWithScore = (PrimaryKeyWithScore) primaryKeysWithScore.get(0);
-                columnData.add(BufferCell.live(scoreColumn,
+                columnData.add(BufferCell.live(syntheticScoreColumn,
                                                FBUtilities.nowInSeconds(),
                                                FloatType.instance.decompose(pkWithScore.indexScore)));
 

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

@@ -110,7 +110,7 @@ public TopKProcessor(ReadCommand command)
         else
             this.queryVector = null;
         this.limit = command.limits().count();
-        this.scoreColumn = ColumnMetadata.syntheticColumn(indexContext.getKeyspace(), indexContext.getTable(), ColumnMetadata.SYNTHETIC_SCORE_ID, FloatType.instance);
+        this.scoreColumn = ColumnMetadata.syntheticScoreColumn(expression.column(), FloatType.instance);
     }
 
     /**

src/java/org/apache/cassandra/schema/ColumnMetadata.java

@@ -117,6 +117,9 @@ public boolean isPrimaryKeyKind()
     private final Comparator<Object> asymmetricCellPathComparator;
     private final Comparator<? super Cell<?>> cellComparator;
 
+    // When the kind is SYNTHETIC, this is the column from which the synthetic column is derived
+    public final ColumnIdentifier sythenticSourceColumn;
+
     private int hash;
 
     /**
@@ -184,9 +187,9 @@ public static ColumnMetadata staticColumn(String keyspace, String table, String
     /**
      * Creates a new synthetic column metadata instance.
      */
-    public static ColumnMetadata syntheticColumn(String keyspace, String table, ColumnIdentifier id, AbstractType<?> type)
+    public static ColumnMetadata syntheticScoreColumn(ColumnMetadata sourceColumn, AbstractType<?> type)
     {
-        return new ColumnMetadata(keyspace, table, id, type, NO_POSITION, Kind.SYNTHETIC);
+        return new ColumnMetadata(sourceColumn.ksName, sourceColumn.cfName, SYNTHETIC_SCORE_ID, type, NO_POSITION, Kind.SYNTHETIC, false, sourceColumn.name);
     }
 
     /**
@@ -236,6 +239,18 @@ public ColumnMetadata(String ksName,
                           int position,
                           Kind kind,
                           boolean isDropped)
+    {
+        this(ksName, cfName, name, type, position, kind, isDropped, null);
+    }
+
+    public ColumnMetadata(String ksName,
+                          String cfName,
+                          ColumnIdentifier name,
+                          AbstractType<?> type,
+                          int position,
+                          Kind kind,
+                          boolean isDropped,
+                          ColumnIdentifier sythenticSourceColumnName)
     {
         super(ksName, cfName, name, type);
         assert name != null && type != null && kind != null;
@@ -263,6 +278,10 @@ public int compare(Object a, Object b)
         };
         this.comparisonOrder = comparisonOrder(kind, isComplex(), Math.max(0, position), name);
         this.isDropped = isDropped;
+
+        // Synthetic columns are the only ones that can have a source column
+        assert kind == Kind.SYNTHETIC || sythenticSourceColumnName == null;
+        this.sythenticSourceColumn = sythenticSourceColumnName;
     }
 
     private static Comparator<CellPath> makeCellPathComparator(Kind kind, AbstractType<?> type)