Allow value/element indexing on frozen collections in SAI

patch by Sunil Ramchandra Pawar; reviewed by Caleb Rackliffe, David Capwell, and Andres de la Peña for CASSANDRA-18492

Author: sunil9977

CHANGES.txt

@@ -1,4 +1,5 @@
 5.1
+ * Allow value/element indexing on frozen collections in SAI (CASSANDRA-18492)
  * Add tool to offline dump cluster metadata and the log (CASSANDRA-21129)
  * Send client warnings when writing to a large partition (CASSANDRA-17258)
  * Harden the possible range of values for max dictionary size and max total sample size for dictionary training (CASSANDRA-21194)

src/java/org/apache/cassandra/cql3/Relation.java

@@ -44,6 +44,8 @@
  */
 public final class Relation
 {
+    public static final String FROZEN_MAP_ENTRY_PREDICATES_NOT_SUPPORTED = "Map-entry predicates on frozen map column %s are not supported";
+
     /**
      * The raw columns'expression.
      */
@@ -204,7 +206,10 @@ public SingleRestriction toRestriction(TableMetadata table, VariableSpecificatio
             AbstractType<?> baseType = column.type.unwrap();
             checkFalse(baseType instanceof ListType, "Indexes on list entries (%s[index] = value) are not supported.", column.name);
             checkTrue(baseType instanceof MapType, "Column %s cannot be used as a map", column.name);
-            checkTrue(baseType.isMultiCell(), "Map-entry predicates on frozen map column %s are not supported", column.name);
+
+            if (column.isClusteringColumn() && baseType.isCollection() && !column.type.isMultiCell())
+                throw invalidRequest(FROZEN_MAP_ENTRY_PREDICATES_NOT_SUPPORTED, column.name);
+
             columnsExpression.collectMarkerSpecification(boundNames);
         }
 

src/java/org/apache/cassandra/cql3/restrictions/MergedRestriction.java

@@ -27,6 +27,7 @@
 
 import org.apache.cassandra.cql3.Operator;
 import org.apache.cassandra.cql3.QueryOptions;
+import org.apache.cassandra.cql3.Relation;
 import org.apache.cassandra.cql3.functions.Function;
 import org.apache.cassandra.db.filter.IndexHints;
 import org.apache.cassandra.db.filter.RowFilter;
@@ -126,9 +127,21 @@ private static void validate(SimpleRestriction restriction, SimpleRestriction ot
         checkOperator(other);
 
         if (restriction.isContains() != other.isContains())
+        {
+            SimpleRestriction mapEntryRestriction = restriction.isContains() ? restriction : other;
+            if (mapEntryRestriction.isMapElementExpression())
+            {
+                ColumnMetadata column = mapEntryRestriction.firstColumn();
+                if (column.type.isFrozenCollection())
+                {
+                    throw invalidRequest(Relation.FROZEN_MAP_ENTRY_PREDICATES_NOT_SUPPORTED, column.name);
+                }
+            }
+
             throw invalidRequest("Collection column %s can only be restricted by CONTAINS, CONTAINS KEY, NOT_CONTAINS, NOT_CONTAINS_KEY" +
                                  " or map-entry equality if it already restricted by one of those",
                                  restriction.firstColumn().name);
+        }
 
         if (restriction.isSlice() && other.isSlice())
         {

src/java/org/apache/cassandra/cql3/restrictions/SimpleRestriction.java

@@ -30,6 +30,7 @@
 import org.apache.cassandra.cql3.ColumnsExpression;
 import org.apache.cassandra.cql3.Operator;
 import org.apache.cassandra.cql3.QueryOptions;
+import org.apache.cassandra.cql3.Relation;
 import org.apache.cassandra.cql3.functions.Function;
 import org.apache.cassandra.cql3.terms.Term;
 import org.apache.cassandra.cql3.terms.Terms;
@@ -157,6 +158,14 @@ public boolean isContains()
                 || columnsExpression.isMapElementExpression();
     }
 
+    /**
+     * Checks if this restriction is a map element expression (e.g., map['key'] = value).
+     */
+    public boolean isMapElementExpression()
+    {
+        return columnsExpression.isMapElementExpression();
+    }
+
     @Override
     public boolean needsFilteringOrIndexing()
     {
@@ -218,6 +227,18 @@ public boolean isSupportedBy(Index index)
         if (isOnToken())
             return false;
 
+        // For map element expressions, check if the index explicitly supports them.
+        if (columnsExpression.isMapElementExpression())
+        {
+            // If the index directly supports map element expressions, return true
+            if (index.supportsMapElementExpression())
+                return true;
+
+            // Supports post-filtering only and require ALLOW FILTERING
+            if (index.supportsFilteringOnMapElementExpression())
+                return false;
+        }
+
         for (ColumnMetadata column : columns())
         {
             if (index.supportsExpression(column, operator))
@@ -415,13 +436,28 @@ else if (isIN())
                 // TODO only map elements supported for now
                 if (columnsExpression.isMapElementExpression())
                 {
+                    // For frozen maps, check if any index on the column can support map entry predicates
+                    // either directly or via filtering. If not, throw an error.
+                    if (column.type.isFrozenCollection())
+                    {
+                        for (Index index : indexRegistry.listIndexes())
+                        {
+                            if (index.dependsOn(column)
+                                && !index.supportsMapElementExpression()
+                                && !index.supportsFilteringOnMapElementExpression())
+                            {
+                                throw invalidRequest(Relation.FROZEN_MAP_ENTRY_PREDICATES_NOT_SUPPORTED, column.name);
+                            }
+                        }
+                    }
+
                     ByteBuffer key = columnsExpression.element(options);
                     if (key == null)
-                        throw invalidRequest("Invalid null map key for column %s", firstColumn().name.toCQLString());
+                        throw invalidRequest("Invalid null map key for column %s", column.name.toCQLString());
                     if (key == ByteBufferUtil.UNSET_BYTE_BUFFER)
-                        throw invalidRequest("Invalid unset map key for column %s", firstColumn().name.toCQLString());
+                        throw invalidRequest("Invalid unset map key for column %s", column.name.toCQLString());
                     List<ByteBuffer> values = bindAndGet(options);
-                    filter.addMapEquality(firstColumn(), key, operator, values.get(0));
+                    filter.addMapEquality(column, key, operator, values.get(0));
                 }
                 break;
             default: throw new UnsupportedOperationException();

src/java/org/apache/cassandra/cql3/statements/schema/CreateIndexStatement.java

@@ -81,10 +81,13 @@ public final class CreateIndexStatement extends AlterSchemaStatement
     public static final String ONLY_PARTITION_KEY = "Cannot create secondary index on the only partition key column %s";
     public static final String CREATE_ON_FROZEN_COLUMN = "Cannot create %s() index on frozen column %s. Frozen collections are immutable and must be fully " +
                                                          "indexed by using the 'full(%s)' modifier";
-    public static final String FULL_ON_FROZEN_COLLECTIONS = "full() indexes can only be created on frozen collections";
+    public static final String FULL_ON_FROZEN_COLLECTIONS = "full() non-SAI indexes can only be created on frozen collections";
     public static final String NON_COLLECTION_SIMPLE_INDEX = "Cannot create %s() index on %s. Non-collection columns only support simple indexes";
     public static final String CREATE_WITH_NON_MAP_TYPE = "Cannot create index on %s of column %s with non-map type";
     public static final String CREATE_ON_NON_FROZEN_UDT = "Cannot create index on non-frozen UDT column %s";
+    public static final String ENTRIES_INDEX_ON_FROZEN_MAP_CLUSTERING_KEY_NOT_SUPPORTED = "Cannot create ENTRIES index on frozen map clustering column '%s'. " +
+                                                                                           "Map entry predicates (column[key] = value) are not supported on clustering columns. " +
+                                                                                           "Use FULL, KEYS, or VALUES index instead.";
     public static final String INDEX_ALREADY_EXISTS = "Index '%s' already exists";
     public static final String INDEX_DUPLICATE_OF_EXISTING = "Index %s is a duplicate of existing index %s";
     public static final String KEYSPACE_DOES_NOT_MATCH_TABLE = "Keyspace name '%s' doesn't match table name '%s'";
@@ -200,7 +203,7 @@ public Keyspaces apply(ClusterMetadata metadata)
 
         IndexMetadata.Kind kind = attrs.isCustom ? IndexMetadata.Kind.CUSTOM : IndexMetadata.Kind.COMPOSITES;
 
-        indexTargets.forEach(t -> validateIndexTarget(table, kind, t));
+        indexTargets.forEach(t -> validateIndexTarget(table, kind, t, attrs));
 
         String name = null == indexName ? generateIndexName(keyspace, indexTargets) : indexName;
 
@@ -244,7 +247,7 @@ private void validateCustomIndexColumnName(String name)
             throw ire(TOO_LONG_CUSTOM_INDEX_TARGET, name, SchemaConstants.NAME_LENGTH);
     }
 
-    private void validateIndexTarget(TableMetadata table, IndexMetadata.Kind kind, IndexTarget target)
+    private void validateIndexTarget(TableMetadata table, IndexMetadata.Kind kind, IndexTarget target, IndexAttributes attrs)
     {
         ColumnMetadata column = table.getColumn(target.column);
 
@@ -253,6 +256,8 @@ private void validateIndexTarget(TableMetadata table, IndexMetadata.Kind kind, I
 
         AbstractType<?> baseType = column.type.unwrap();
 
+        boolean isNonSAIIndex = !isSAIIndex(attrs);
+
         // TODO: this check needs to be removed with CASSANDRA-20235
         if ((kind == IndexMetadata.Kind.CUSTOM))
             validateCustomIndexColumnName(target.column.toString());
@@ -283,22 +288,41 @@ private void validateIndexTarget(TableMetadata table, IndexMetadata.Kind kind, I
         if (column.isPartitionKey() && table.partitionKeyColumns().size() == 1)
             throw ire(ONLY_PARTITION_KEY, column);
 
-        if (baseType.isFrozenCollection() && target.type != Type.FULL)
-            throw ire(CREATE_ON_FROZEN_COLUMN, target.type, column, column.name.toCQLString());
-
-        if (!baseType.isFrozenCollection() && target.type == Type.FULL)
+        if (target.type == Type.FULL && isNonSAIIndex && (!baseType.isCollection() || column.type.isMultiCell()))
             throw ire(FULL_ON_FROZEN_COLLECTIONS);
 
         if (!baseType.isCollection() && target.type != Type.SIMPLE)
             throw ire(NON_COLLECTION_SIMPLE_INDEX, target.type, column);
 
-        if (!(baseType instanceof MapType && baseType.isMultiCell()) && (target.type == Type.KEYS || target.type == Type.KEYS_AND_VALUES))
+        // Frozen collections are only supported with SAI indexes.
+        if (isNonSAIIndex && baseType.isCollection() && !column.type.isMultiCell())
+        {
+            if (target.type == Type.VALUES || target.type == Type.KEYS || target.type == Type.KEYS_AND_VALUES)
+            {
+                throw ire(CREATE_ON_FROZEN_COLUMN, target.type.toString(), column.name, column.name);
+            }
+        }
+
+        if (!(baseType instanceof MapType) && (target.type == Type.KEYS || target.type == Type.KEYS_AND_VALUES ))
             throw ire(CREATE_WITH_NON_MAP_TYPE, target.type, column);
 
+        // Can't query map[key]=value on clustering key columns, so ENTRIES index would be not queryable.
+        if (column.isClusteringColumn() && baseType instanceof MapType && !column.type.isMultiCell()
+            && target.type == Type.KEYS_AND_VALUES)
+            throw ire(ENTRIES_INDEX_ON_FROZEN_MAP_CLUSTERING_KEY_NOT_SUPPORTED, column.name);
+
         if (column.type.isUDT() && column.type.isMultiCell())
             throw ire(CREATE_ON_NON_FROZEN_UDT, column);
     }
 
+    /**
+     * Checks if the given index attributes represent a Storage Attached Index.
+     */
+    private boolean isSAIIndex(IndexAttributes attrs)
+    {
+        return attrs.isCustom && IndexMetadata.isSAIIndex(attrs.customClass);
+    }
+
     private String generateIndexName(KeyspaceMetadata keyspace, List<IndexTarget> targets)
     {
         String baseName = targets.size() == 1

src/java/org/apache/cassandra/db/filter/RowFilter.java

@@ -527,6 +527,11 @@ public Operator operator()
             return operator;
         }
 
+        public boolean isMapElementExpression()
+        {
+            return kind() == Kind.MAP_ELEMENT;
+        }
+
         /**
          * If this expression is used to query an index, the value to use as
          * partition key for that index query.

src/java/org/apache/cassandra/index/Index.java

@@ -437,6 +437,27 @@ default boolean filtersMultipleContains()
         return true;
     }
 
+    /**
+     * Return whether this index supports map element expressions on frozen map columns.
+     *
+     * @return {@code true} if this index supports map element else {@code false}.
+     */
+    default boolean supportsMapElementExpression()
+    {
+        return false;
+    }
+
+    /**
+     * Returns whether index allows filtering for map element expressions on frozen collections.
+     * SAI can handle map element predicates via post-filtering.
+     *
+     * @return {@code true} if map element expressions can be evaluated via filtering, {@code false} otherwise.
+     */
+    default boolean supportsFilteringOnMapElementExpression()
+    {
+        return false;
+    }
+
     /**
      * If the index supports custom search expressions using the
      * {@code}SELECT * FROM table WHERE expr(index_name, expression){@code} syntax, this

src/java/org/apache/cassandra/index/sai/StorageAttachedIndex.java

@@ -453,12 +453,37 @@ public boolean supportsExpression(ColumnMetadata column, Operator operator)
         return dependsOn(column) && indexTermType.supports(operator);
     }
 
+    @Override
+    public boolean supportsExpression(RowFilter.Expression expression)
+    {
+        if (expression.isMapElementExpression() &&
+            indexTermType.isFrozenCollection() &&
+            indexTermType.indexTargetType() == IndexTarget.Type.FULL)
+
+            return false;
+
+        return supportsExpression(expression.column(), expression.operator());
+    }
+
     @Override
     public boolean filtersMultipleContains()
     {
         return false;
     }
 
+    @Override
+    public boolean supportsMapElementExpression()
+    {
+        return termType().indexTargetType() == IndexTarget.Type.KEYS_AND_VALUES;
+    }
+
+    @Override
+    public boolean supportsFilteringOnMapElementExpression()
+    {
+        // SAI supports map element expressions via post-filtering on frozen collections
+        return true;
+    }
+
     @Override
     public AbstractType<?> customExpressionValueType()
     {
@@ -784,6 +809,12 @@ public void validateTermSizeForRow(DecoratedKey key, Row row, boolean isClientMu
             while (bufferIterator != null && bufferIterator.hasNext())
                 validateTermSizeForCell(analyzer, key, bufferIterator.next(), isClientMutation, state);
         }
+        else if (indexTermType.isFrozenCollection() && indexTermType.indexTargetType() != IndexTarget.Type.FULL)
+        {
+            Iterator<ByteBuffer> bufferIterator = indexTermType.valuesOfFrozenCollection(row, FBUtilities.nowInSeconds());
+            while (bufferIterator != null && bufferIterator.hasNext())
+                validateTermSizeForCell(analyzer, key, bufferIterator.next(), isClientMutation, state);
+        }
         else
         {
             ByteBuffer value = indexTermType.valueOf(key, row, FBUtilities.nowInSeconds());

src/java/org/apache/cassandra/index/sai/disk/v1/SSTableIndexWriter.java

@@ -32,6 +32,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.cassandra.cql3.statements.schema.IndexTarget;
 import org.apache.cassandra.db.rows.Row;
 import org.apache.cassandra.index.sai.StorageAttachedIndex;
 import org.apache.cassandra.index.sai.analyzer.AbstractAnalyzer;
@@ -94,6 +95,18 @@ public void addRow(PrimaryKey key, Row row, long sstableRowId) throws IOExceptio
                 }
             }
         }
+        else if (index.termType().isFrozenCollection() && index.termType().indexTargetType() != IndexTarget.Type.FULL)
+        {
+            Iterator<ByteBuffer> valueIterator = index.termType().valuesOfFrozenCollection(row, nowInSec);
+            if (valueIterator != null)
+            {
+                while (valueIterator.hasNext())
+                {
+                    ByteBuffer value = valueIterator.next();
+                    addTerm(index.termType().asIndexBytes(value.duplicate()), key, sstableRowId);
+                }
+            }
+        }
         else
         {
             ByteBuffer value = index.termType().valueOf(key.partitionKey(), row, nowInSec);

src/java/org/apache/cassandra/index/sai/memory/MemtableIndexManager.java

@@ -31,6 +31,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 
+import org.apache.cassandra.cql3.statements.schema.IndexTarget;
 import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.lifecycle.ILifecycleTransaction;
 import org.apache.cassandra.db.memtable.Memtable;
@@ -76,6 +77,18 @@ public long index(DecoratedKey key, Row row, Memtable mt)
                 }
             }
         }
+        else if (index.termType().isFrozenCollection() && index.termType().indexTargetType() != IndexTarget.Type.FULL)
+        {
+            Iterator<ByteBuffer> bufferIterator = index.termType().valuesOfFrozenCollection(row, FBUtilities.nowInSeconds());
+            if (bufferIterator != null)
+            {
+                while (bufferIterator.hasNext())
+                {
+                    ByteBuffer value = bufferIterator.next();
+                    bytes += target.index(key, row.clustering(), value);
+                }
+            }
+        }
         else
         {
             ByteBuffer value = index.termType().valueOf(key, row, FBUtilities.nowInSeconds());