CNDB-14997: Fix column selection on SAI-ordered queries

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

@@ -589,7 +589,7 @@ public Columns deserialize(DataInputPlus in, TableMetadata metadata) throws IOEx
          */
         public void serializeSubset(Collection<ColumnMetadata> columns, Columns superset, DataOutputPlus out) throws IOException
         {
-            /**
+            /*
              * We weight this towards small sets, and sets where the majority of items are present, since
              * we expect this to mostly be used for serializing result sets.
              *

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

@@ -112,7 +112,7 @@ public boolean isReverseOrder()
         @Override
         public RegularAndStaticColumns columns()
         {
-            return command.metadata().regularAndStaticColumns();
+            return partitionInfo.columns;
         }
 
         @Override

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

@@ -18,12 +18,11 @@
 
 package org.apache.cassandra.index.sai.utils;
 
-import java.util.Objects;
-
 import javax.annotation.Nullable;
 
 import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.DeletionTime;
+import org.apache.cassandra.db.RegularAndStaticColumns;
 import org.apache.cassandra.db.rows.BaseRowIterator;
 import org.apache.cassandra.db.rows.EncodingStats;
 import org.apache.cassandra.db.rows.Row;
@@ -34,6 +33,7 @@ public class PartitionInfo
 {
     public final DecoratedKey key;
     public final Row staticRow;
+    public final RegularAndStaticColumns columns;
 
     // present if it's unfiltered partition iterator
     @Nullable
@@ -43,45 +43,35 @@ public class PartitionInfo
     @Nullable
     public final EncodingStats encodingStats;
 
-    public PartitionInfo(DecoratedKey key, Row staticRow)
-    {
-        this.key = key;
-        this.staticRow = staticRow;
-        this.partitionDeletion = null;
-        this.encodingStats = null;
-    }
-
-    public PartitionInfo(DecoratedKey key, Row staticRow, DeletionTime partitionDeletion, EncodingStats encodingStats)
+    private PartitionInfo(DecoratedKey key,
+                          Row staticRow,
+                          RegularAndStaticColumns columns,
+                          @Nullable DeletionTime partitionDeletion,
+                          @Nullable EncodingStats encodingStats)
     {
         this.key = key;
         this.staticRow = staticRow;
-
+        this.columns = columns;
         this.partitionDeletion = partitionDeletion;
         this.encodingStats = encodingStats;
     }
 
     public static <U extends Unfiltered, R extends BaseRowIterator<U>> PartitionInfo create(R baseRowIterator)
     {
-        return baseRowIterator instanceof UnfilteredRowIterator
-               ? new PartitionInfo(baseRowIterator.partitionKey(), baseRowIterator.staticRow(),
-                                   ((UnfilteredRowIterator) baseRowIterator).partitionLevelDeletion(),
-                                   ((UnfilteredRowIterator) baseRowIterator).stats())
-               : new PartitionInfo(baseRowIterator.partitionKey(), baseRowIterator.staticRow());
-    }
-
-    @Override
-    public boolean equals(Object o)
-    {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        PartitionInfo that = (PartitionInfo) o;
-        return Objects.equals(key, that.key) && Objects.equals(staticRow, that.staticRow)
-               && Objects.equals(partitionDeletion, that.partitionDeletion) && Objects.equals(encodingStats, that.encodingStats);
-    }
+        // only unfiltered row iterators have a partition deletion time and encoding stats
+        DeletionTime partitionDeletion = null;
+        EncodingStats encodingStats = null;
+        if (baseRowIterator instanceof UnfilteredRowIterator)
+        {
+            UnfilteredRowIterator unfilteredRowIterator = (UnfilteredRowIterator) baseRowIterator;
+            partitionDeletion = unfilteredRowIterator.partitionLevelDeletion();
+            encodingStats = unfilteredRowIterator.stats();
+        }
 
-    @Override
-    public int hashCode()
-    {
-        return Objects.hash(key, staticRow, partitionDeletion, encodingStats);
+        return new PartitionInfo(baseRowIterator.partitionKey(),
+                                 baseRowIterator.staticRow(),
+                                 baseRowIterator.columns(),
+                                 partitionDeletion,
+                                 encodingStats);
     }
 }

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

@@ -0,0 +1,83 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * 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,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.index.sai.cql;
+
+import org.junit.Test;
+
+public class ColumnSelectionTest extends VectorTester
+{
+    /**
+     * Tests that we can select any type of column in any table with SAI indexes. See CNDB-14997 for further details.
+     */
+    @Test
+    public void testColumnSelection() throws Throwable
+    {
+        createTable("CREATE TABLE %s (" +
+                    "k1 text, k2 text, " +
+                    "c1 text, c2 text, " +
+                    "s1 text static, s2 text static, " +
+                    "r1 text, r2 vector<float, 2>, " +
+                    "PRIMARY KEY((k1, k2), c1, c2))");
+        createIndex("CREATE CUSTOM INDEX ON %s(r1) USING 'StorageAttachedIndex'");
+        createIndex("CREATE CUSTOM INDEX ON %s(r2) USING 'StorageAttachedIndex'");
+
+        Object[] row = row("k1", "k2", "c1", "c2", "s1", "s2", "r1", vector(1, 2));
+        execute("INSERT INTO %s (k1, k2, c1, c2, s1, s2, r1, r2) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", row);
+
+        beforeAndAfterFlush(() -> {
+
+            // filtering
+            assertRows(execute("SELECT * FROM %s WHERE r1='r1'"), row);
+            assertRows(execute("SELECT k1 FROM %s WHERE r1='r1'"), row(row[0]));
+            assertRows(execute("SELECT k2 FROM %s WHERE r1='r1'"), row(row[1]));
+            assertRows(execute("SELECT c1 FROM %s WHERE r1='r1'"), row(row[2]));
+            assertRows(execute("SELECT c2 FROM %s WHERE r1='r1'"), row(row[3]));
+            assertRows(execute("SELECT s1 FROM %s WHERE r1='r1'"), row(row[4]));
+            assertRows(execute("SELECT s2 FROM %s WHERE r1='r1'"), row(row[5]));
+            assertRows(execute("SELECT r1 FROM %s WHERE r1='r1'"), row(row[6]));
+            assertRows(execute("SELECT r2 FROM %s WHERE r1='r1'"), row(row[7]));
+            assertRows(execute("SELECT k1, c1, s1, r1 FROM %s WHERE r1='r1'"), row(row[0], row[2], row[4], row[6]));
+            assertRows(execute("SELECT k2, c2, s2, r2 FROM %s WHERE r1='r1'"), row(row[1], row[3], row[5], row[7]));
+
+            // generic ordering
+            assertRows(execute("SELECT * FROM %s ORDER BY r1 LIMIT 10"), row);
+            assertRows(execute("SELECT k1 FROM %s ORDER BY r1 LIMIT 10"), row(row[0]));
+            assertRows(execute("SELECT k2 FROM %s ORDER BY r1 LIMIT 10"), row(row[1]));
+            assertRows(execute("SELECT c1 FROM %s ORDER BY r1 LIMIT 10"), row(row[2]));
+            assertRows(execute("SELECT c2 FROM %s ORDER BY r1 LIMIT 10"), row(row[3]));
+            assertRows(execute("SELECT s1 FROM %s ORDER BY r1 LIMIT 10"), row(row[4]));
+            assertRows(execute("SELECT s2 FROM %s ORDER BY r1 LIMIT 10"), row(row[5]));
+            assertRows(execute("SELECT r1 FROM %s ORDER BY r1 LIMIT 10"), row(row[6]));
+            assertRows(execute("SELECT r2 FROM %s ORDER BY r1 LIMIT 10"), row(row[7]));
+            assertRows(execute("SELECT k1, c1, s1, r1 FROM %s ORDER BY r1 LIMIT 10"), row(row[0], row[2], row[4], row[6]));
+            assertRows(execute("SELECT k2, c2, s2, r2 FROM %s ORDER BY r1 LIMIT 10"), row(row[1], row[3], row[5], row[7]));
+
+            // ANN ordering
+            assertRows(execute("SELECT * FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row);
+            assertRows(execute("SELECT k1 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[0]));
+            assertRows(execute("SELECT k2 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[1]));
+            assertRows(execute("SELECT c1 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[2]));
+            assertRows(execute("SELECT c2 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[3]));
+            assertRows(execute("SELECT s1 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[4]));
+            assertRows(execute("SELECT s2 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[5]));
+            assertRows(execute("SELECT r1 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[6]));
+            assertRows(execute("SELECT r2 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[7]));
+            assertRows(execute("SELECT k1, c1, s1, r1 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[0], row[2], row[4], row[6]));
+            assertRows(execute("SELECT k2, c2, s2, r2 FROM %s ORDER BY r2 ANN OF [1, 2] LIMIT 10"), row(row[1], row[3], row[5], row[7]));
+        });
+    }
+}

test/unit/org/apache/cassandra/index/sai/utils/PartitionInfoTest.java

@@ -0,0 +1,107 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * 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,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.index.sai.utils;
+
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import org.apache.cassandra.Util;
+import org.apache.cassandra.config.DatabaseDescriptor;
+import org.apache.cassandra.db.Clustering;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.DeletionTime;
+import org.apache.cassandra.db.LivenessInfo;
+import org.apache.cassandra.db.RegularAndStaticColumns;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.db.rows.BTreeRow;
+import org.apache.cassandra.db.rows.EncodingStats;
+import org.apache.cassandra.db.rows.Row;
+import org.apache.cassandra.db.rows.RowIterator;
+import org.apache.cassandra.db.rows.UnfilteredRowIterator;
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.schema.TableMetadata;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class PartitionInfoTest
+{
+    private static DecoratedKey partitionKey;
+    private static Row staticRow;
+    private static RegularAndStaticColumns columns;
+    private static DeletionTime partitionDeletion;
+    private static EncodingStats encodingStats;
+
+    @BeforeClass
+    public static void setup()
+    {
+        DatabaseDescriptor.setPartitionerUnsafe(new Murmur3Partitioner());
+        TableMetadata tableMetadata = TableMetadata.builder("test", "test")
+                                                   .partitioner(Murmur3Partitioner.instance)
+                                                   .addPartitionKeyColumn("pk", Int32Type.instance)
+                                                   .addRegularColumn("regular_col", Int32Type.instance)
+                                                   .addStaticColumn("static_col", Int32Type.instance)
+                                                   .build();
+
+        partitionKey = Util.dk("test_key");
+        staticRow = BTreeRow.emptyRow(Clustering.STATIC_CLUSTERING);
+        columns = tableMetadata.regularAndStaticColumns();
+        partitionDeletion = new DeletionTime(1000L, 10);
+        encodingStats = new EncodingStats(500L, LivenessInfo.NO_EXPIRATION_TIME, 0);
+    }
+
+    @Test
+    public void testCreateFromUnfilteredIterator()
+    {
+        UnfilteredRowIterator iterator = mock(UnfilteredRowIterator.class);
+        when(iterator.partitionKey()).thenReturn(partitionKey);
+        when(iterator.staticRow()).thenReturn(staticRow);
+        when(iterator.columns()).thenReturn(columns);
+        when(iterator.partitionLevelDeletion()).thenReturn(partitionDeletion);
+        when(iterator.stats()).thenReturn(encodingStats);
+
+        PartitionInfo info = PartitionInfo.create(iterator);
+
+        assertNotNull(info);
+        assertSame(partitionKey, info.key);
+        assertSame(staticRow, info.staticRow);
+        assertSame(columns, info.columns);
+        assertSame(partitionDeletion, info.partitionDeletion);
+        assertSame(encodingStats, info.encodingStats);
+    }
+
+    @Test
+    public void testCreateFromRowIterator()
+    {
+        RowIterator iterator = mock(RowIterator.class);
+        when(iterator.partitionKey()).thenReturn(partitionKey);
+        when(iterator.staticRow()).thenReturn(staticRow);
+        when(iterator.columns()).thenReturn(columns);
+
+        PartitionInfo info = PartitionInfo.create(iterator);
+
+        assertNotNull(info);
+        assertSame(partitionKey, info.key);
+        assertSame(staticRow, info.staticRow);
+        assertSame(columns, info.columns);
+        assertNull(info.partitionDeletion);
+        assertNull(info.encodingStats);
+    }
+}