CASSANDRA-20460: Expand TableWalk tests to include collections and add support for += and -= for these types

test/distributed/org/apache/cassandra/distributed/test/cql3/SingleNodeTableWalkTest.java

@@ -66,10 +66,11 @@
 import org.apache.cassandra.utils.ASTGenerators;
 import org.apache.cassandra.utils.AbstractTypeGenerators;
 import org.apache.cassandra.utils.AbstractTypeGenerators.TypeGenBuilder;
+import org.apache.cassandra.utils.AbstractTypeGenerators.TypeKind;
 import org.apache.cassandra.utils.ByteBufferUtil;
 import org.apache.cassandra.utils.CassandraGenerators.TableMetadataBuilder;
+import org.apache.cassandra.utils.Generators;
 import org.apache.cassandra.utils.ImmutableUniqueList;
-import org.quicktheories.generators.SourceDSL;
 
 import static accord.utils.Property.commands;
 import static accord.utils.Property.stateful;
@@ -78,6 +79,17 @@
 
 public class SingleNodeTableWalkTest extends StatefulASTBase
 {
+    private static final Gen<Gen<Boolean>> BOOLEAN_DISTRIBUTION = Gens.bools().mixedDistribution();
+    //TODO (coverage): COMPOSITE, DYNAMIC_COMPOSITE
+    private static Gen<Gen<TypeKind>> TYPE_KIND_DISTRIBUTION = Gens.mixedDistribution(TypeKind.PRIMITIVE,
+                                                                                      TypeKind.SET, TypeKind.LIST, TypeKind.MAP,
+                                                                                      TypeKind.TUPLE, TypeKind.UDT,
+                                                                                      TypeKind.VECTOR
+    );
+    private static Gen<Gen<AbstractType<?>>> PRIMITIVE_DISTRIBUTION = Gens.mixedDistribution(AbstractTypeGenerators.knownPrimitiveTypes()
+                                                                                                                   .stream()
+                                                                                                                   .filter(t -> !AbstractTypeGenerators.isUnsafeEquality(t))
+                                                                                                                   .collect(Collectors.toList()));
     private static final Logger logger = LoggerFactory.getLogger(SingleNodeTableWalkTest.class);
 
     protected void preCheck(Cluster cluster, Property.StatefulBuilder builder)
@@ -88,10 +100,20 @@ protected void preCheck(Cluster cluster, Property.StatefulBuilder builder)
         // CQL_DEBUG_APPLY_OPERATOR = true;
     }
 
-    protected TypeGenBuilder supportedTypes()
+    protected TypeGenBuilder supportedTypes(RandomSource rs)
     {
-        return AbstractTypeGenerators.withoutUnsafeEquality(AbstractTypeGenerators.builder()
-                                                                                  .withTypeKinds(AbstractTypeGenerators.TypeKind.PRIMITIVE));
+        return AbstractTypeGenerators.builder()
+                                     .withTypeKinds(Generators.fromGen(TYPE_KIND_DISTRIBUTION.next(rs)))
+                                     .withPrimitives(Generators.fromGen(PRIMITIVE_DISTRIBUTION.next(rs)))
+                                     .withUserTypeFields(AbstractTypeGenerators.UserTypeFieldsGen.simpleNames())
+                                     .withMaxDepth(1);
+    }
+
+    protected TypeGenBuilder supportedPrimaryColumnTypes(RandomSource rs)
+    {
+        return AbstractTypeGenerators.builder()
+                                     .withTypeKinds(TypeKind.PRIMITIVE)
+                                     .withPrimitives(Generators.fromGen(PRIMITIVE_DISTRIBUTION.next(rs)));
     }
 
     protected List<CreateIndexDDL.Indexer> supportedIndexers()
@@ -206,7 +228,7 @@ protected List<CreateIndexDDL.Indexer> supportedIndexers()
             builder.value(pk, key.bufferAt(pks.indexOf(pk)));
 
 
-        List<Symbol> searchableColumns = state.nonPartitionColumns;
+        List<Symbol> searchableColumns = state.searchableNonPartitionColumns;
         Symbol symbol = rs.pick(searchableColumns);
 
         TreeMap<ByteBuffer, List<BytesPartitionState.PrimaryKey>> universe = state.model.index(ref, symbol);
@@ -363,7 +385,8 @@ protected TableMetadata defineTable(RandomSource rs, String ks)
                                        .withCompression())
                      .withKeyspaceName(ks).withTableName("tbl")
                      .withSimpleColumnNames()
-                     .withDefaultTypeGen(supportedTypes())
+                     .withDefaultTypeGen(supportedTypes(rs))
+                     .withPrimaryColumnTypeGen(supportedPrimaryColumnTypes(rs))
                      .withPartitioner(Murmur3Partitioner.instance)
                      .build())
                .next(rs);
@@ -393,7 +416,7 @@ public class State extends CommonState
     {
         protected final LinkedHashMap<Symbol, IndexedColumn> indexes;
         private final Gen<Mutation> mutationGen;
-        private final List<Symbol> nonPartitionColumns;
+        private final List<Symbol> searchableNonPartitionColumns;
         private final List<Symbol> searchableColumns;
         private final List<Symbol> nonPkIndexedColumns;
 
@@ -424,7 +447,8 @@ public State(RandomSource rs, Cluster cluster)
                                                                   .withoutTransaction()
                                                                   .withoutTtl()
                                                                   .withoutTimestamp()
-                                                                  .withPartitions(SourceDSL.arbitrary().pick(uniquePartitions));
+                                                                  .withPartitions(Generators.fromGen(Gens.mixedDistribution(uniquePartitions).next(rs)))
+                                                                  .withColumnExpressions(e -> e.withOperators(Generators.fromGen(BOOLEAN_DISTRIBUTION.next(rs))));
             if (IGNORED_ISSUES.contains(KnownIssue.SAI_EMPTY_TYPE))
             {
                 model.factory.regularAndStaticColumns.stream()
@@ -438,16 +462,30 @@ public State(RandomSource rs, Cluster cluster)
             }
             this.mutationGen = toGen(mutationGenBuilder.build());
 
-            nonPartitionColumns = ImmutableList.<Symbol>builder()
-                                               .addAll(model.factory.clusteringColumns)
-                                               .addAll(model.factory.staticColumns)
-                                               .addAll(model.factory.regularColumns)
-                                               .build();
+            var nonPartitionColumns = ImmutableList.<Symbol>builder()
+                                                   .addAll(model.factory.clusteringColumns)
+                                                   .addAll(model.factory.staticColumns)
+                                                   .addAll(model.factory.regularColumns)
+                                                   .build();
+            searchableNonPartitionColumns = nonPartitionColumns.stream()
+                                                               .filter(this::isSearchable)
+                                                               .collect(Collectors.toList());
             nonPkIndexedColumns = nonPartitionColumns.stream()
                                                      .filter(indexes::containsKey)
                                                      .collect(Collectors.toList());
 
-            searchableColumns = metadata.partitionKeyColumns().size() > 1 ?  model.factory.selectionOrder : nonPartitionColumns;
+            searchableColumns = (metadata.partitionKeyColumns().size() > 1 ? model.factory.selectionOrder : nonPartitionColumns)
+                                .stream()
+                                .filter(this::isSearchable)
+                                .collect(Collectors.toList());
+        }
+
+        private boolean isSearchable(Symbol symbol)
+        {
+            // See org.apache.cassandra.cql3.Operator.validateFor
+            // multi cell collections can only be searched if you search their elements, not the collection as a whole
+            //TODO (coverage): can you query for UDT fields?  its a single cell so you "should"?
+            return !(symbol.type().isMultiCell() && (symbol.type().isCollection() || symbol.type().isUDT()));
         }
 
         @Override
@@ -523,6 +561,8 @@ private List<Symbol> multiColumnQueryColumns()
             List<Symbol> allowedColumns = searchableColumns;
             if (hasMultiNodeMultiColumnAllowFilteringWithLocalWritesIssue())
                 allowedColumns = nonPkIndexedColumns;
+            if (IGNORED_ISSUES.contains(KnownIssue.SAI_AND_VECTOR_COLUMNS) && !indexes.isEmpty())
+                allowedColumns = allowedColumns.stream().filter(s -> !s.type().isVector()).collect(Collectors.toList());
             return allowedColumns;
         }
 
@@ -533,7 +573,7 @@ private boolean hasMultiNodeMultiColumnAllowFilteringWithLocalWritesIssue()
 
         public boolean allowPartitionQuery()
         {
-            return !(model.isEmpty() || nonPartitionColumns.isEmpty());
+            return !(model.isEmpty() || searchableNonPartitionColumns.isEmpty());
         }
 
         @Override

test/harry/main/org/apache/cassandra/harry/model/ASTSingleTableModel.java

@@ -21,6 +21,7 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.BitSet;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -43,13 +44,16 @@
 import com.google.common.collect.Sets;
 
 import accord.utils.Invariants;
+import org.apache.cassandra.cql3.ast.AssignmentOperator;
 import org.apache.cassandra.cql3.ast.Conditional;
 import org.apache.cassandra.cql3.ast.Conditional.Where.Inequality;
 import org.apache.cassandra.cql3.ast.Element;
 import org.apache.cassandra.cql3.ast.Expression;
 import org.apache.cassandra.cql3.ast.ExpressionEvaluator;
 import org.apache.cassandra.cql3.ast.FunctionCall;
+import org.apache.cassandra.cql3.ast.Literal;
 import org.apache.cassandra.cql3.ast.Mutation;
+import org.apache.cassandra.cql3.ast.Operator;
 import org.apache.cassandra.cql3.ast.Select;
 import org.apache.cassandra.cql3.ast.StandardVisitors;
 import org.apache.cassandra.cql3.ast.Symbol;
@@ -244,7 +248,10 @@ public void update(Mutation.Update update)
                 // static columns to add in.  If we are doing something like += to a row that doesn't exist, we still update statics...
                 Map<Symbol, ByteBuffer> write = new HashMap<>();
                 for (Symbol col : Sets.intersection(factory.staticColumns.asSet(), set.keySet()))
-                    write.put(col, eval(set.get(col)));
+                {
+                    ByteBuffer current = partition.staticRow().get(col);
+                    write.put(col, eval(col, current, set.get(col)));
+                }
                 partition.setStaticColumns(write);
             }
             // table has clustering but non are in the write, so only pk/static can be updated
@@ -254,7 +261,10 @@ public void update(Mutation.Update update)
             {
                 Map<Symbol, ByteBuffer> write = new HashMap<>();
                 for (Symbol col : Sets.intersection(factory.regularColumns.asSet(), set.keySet()))
-                    write.put(col, eval(set.get(col)));
+                {
+                    ByteBuffer current = partition.get(cd, col);
+                    write.put(col, eval(col, current, set.get(col)));
+                }
 
                 partition.setColumns(cd, write, false);
             }
@@ -493,6 +503,45 @@ private static void validateAnyOrder(ImmutableUniqueList<Symbol> columns, Set<Ro
             if (actual.isEmpty()) sb.append("No rows returned");
             else sb.append("Missing rows:\n").append(table(columns, missing));
         }
+        if (!unexpected.isEmpty() && unexpected.size() == missing.size())
+        {
+            // good chance a column differs
+            StringBuilder finalSb = sb;
+            Runnable runOnce = new Runnable()
+            {
+                boolean ran = false;
+                @Override
+                public void run()
+                {
+                    if (ran) return;
+                    finalSb.append("\nPossible column conflicts:");
+                    ran = true;
+                }
+            };
+            for (var e : missing)
+            {
+                Row smallest = null;
+                BitSet smallestDiff = null;
+                for (var a : unexpected)
+                {
+                    BitSet diff = e.diff(a);
+                    if (smallestDiff == null || diff.cardinality() < smallestDiff.cardinality())
+                    {
+                        smallest = a;
+                        smallestDiff = diff;
+                    }
+                }
+                // if every column differs then ignore
+                if (smallestDiff.cardinality() == e.values.length)
+                    continue;
+                runOnce.run();
+                sb.append("\n\tExpected: ").append(e);
+                sb.append("\n\tDiff (expected over actual):\n");
+                Row eSmall = e.select(smallestDiff);
+                Row aSmall = smallest.select(smallestDiff);
+                sb.append(table(eSmall.columns, Arrays.asList(eSmall, aSmall)));
+            }
+        }
         if (sb != null)
         {
             sb.append("\nExpected:\n").append(table(columns, expected));
@@ -731,7 +780,7 @@ private static boolean matches(ByteBuffer value, List<? extends Expression> cond
         for (Expression e : conditions)
         {
             ByteBuffer expected = eval(e);
-            if (expected.equals(value))
+            if (expected != null && expected.equals(value))
                 return true;
         }
         return false;
@@ -893,28 +942,61 @@ private List<Clustering<ByteBuffer>> keys(Map<Symbol, List<? extends Expression>
         return current.stream().map(BufferClustering::new).collect(Collectors.toList());
     }
 
+    private static ByteBuffer eval(Symbol col, @Nullable ByteBuffer current, Expression e)
+    {
+        if (!(e instanceof AssignmentOperator)) return eval(e);
+        // multi cell collections have the property that they do update even if the current value is null
+        boolean isFancy = col.type().isCollection() && col.type().isMultiCell();
+        if (current == null && !isFancy) return null; // null + ? == null
+        var assignment = (AssignmentOperator) e;
+        if (isFancy && current == null)
+        {
+            return assignment.kind == AssignmentOperator.Kind.SUBTRACT
+                   // if it doesn't exist, then there is nothing to subtract
+                   ? null
+                   : eval(assignment.right);
+        }
+        switch (assignment.kind)
+        {
+            case ADD:
+                return eval(new Operator(Operator.Kind.ADD, new Literal(current, e.type()), assignment.right));
+            case SUBTRACT:
+                return eval(new Operator(Operator.Kind.SUBTRACT, new Literal(current, e.type()), assignment.right));
+            default:
+                throw new UnsupportedOperationException(assignment.kind + ": " + assignment.toCQL());
+        }
+    }
+
+    @Nullable
     private static ByteBuffer eval(Expression e)
     {
-        return ExpressionEvaluator.tryEvalEncoded(e).get();
+        return ExpressionEvaluator.evalEncoded(e);
     }
 
     private static class Row
     {
+        private static final Row EMPTY = new Row(ImmutableUniqueList.empty(), ByteBufferUtil.EMPTY_ARRAY);
+
         private final ImmutableUniqueList<Symbol> columns;
         private final ByteBuffer[] values;
 
         private Row(ImmutableUniqueList<Symbol> columns, ByteBuffer[] values)
         {
             this.columns = columns;
             this.values = values;
+            if (columns.size() != values.length)
+                throw new IllegalArgumentException("Columns " + columns + " should have the same size as values, but had " + values.length);
         }
 
         public String asCQL(Symbol symbol)
         {
             int offset = columns.indexOf(symbol);
             assert offset >= 0;
             ByteBuffer b = values[offset];
-            return (b == null || ByteBufferUtil.EMPTY_BYTE_BUFFER.equals(b)) ? "null" : symbol.type().asCQL3Type().toCQLLiteral(b);
+            if (b == null) return "null";
+            if (ByteBufferUtil.UNSET_BYTE_BUFFER == b) return "<unset>";
+            if (ByteBufferUtil.EMPTY_BYTE_BUFFER.equals(b)) return "<empty>";
+            return symbol.type().asCQL3Type().toCQLLiteral(b);
         }
 
         public List<String> asCQL()
@@ -925,6 +1007,40 @@ public List<String> asCQL()
             return human;
         }
 
+        public BitSet diff(Row other)
+        {
+            if (!columns.equals(other.columns))
+                throw new UnsupportedOperationException("Columns do not match: expected " + columns + " but given " + other.columns);
+            int maxLength = Math.max(values.length, other.values.length);
+            int minLength = Math.min(values.length, other.values.length);
+            BitSet set = new BitSet(maxLength);
+            for (int i = 0; i < minLength; i++)
+            {
+                ByteBuffer a = values[i];
+                ByteBuffer b = other.values[i];
+                if (!Objects.equals(a, b))
+                    set.set(i);
+            }
+            for (int i = minLength; i < maxLength; i++)
+                set.set(i);
+            return set;
+        }
+
+        public Row select(BitSet selection)
+        {
+            if (selection.isEmpty()) return EMPTY;
+            var names = ImmutableUniqueList.<Symbol>builder(selection.cardinality());
+            ByteBuffer[] copy = new ByteBuffer[selection.cardinality()];
+            int offset = 0;
+            for (int i = 0; i < this.values.length; i++)
+            {
+                if (!selection.get(i)) continue;
+                names.add(this.columns.get(i));
+                copy[offset++] = this.values[i];
+            }
+            return new Row(names.build(), copy);
+        }
+
         @Override
         public boolean equals(Object o)
         {