CASSANDRA-21484 - Support IF, ELSE IF, and ELSE for transaction statements

.gitmodules

@@ -1,4 +1,4 @@
 [submodule "modules/accord"]
 	path = modules/accord
-	url = https://github.com/apache/cassandra-accord.git
-	branch = trunk
+	url = https://github.com/alanwang67/cassandra-accord.git
+	branch = linearIntersectionBug

src/antlr/Lexer.g

@@ -128,6 +128,7 @@ K_ALLOW:       A L L O W;
 K_FILTERING:   F I L T E R I N G;
 K_IF:          I F;
 K_THEN:        T H E N;
+K_ELSE:        E L S E;
 K_END:         E N D;
 K_IS:          I S;
 K_CONTAINS:    C O N T A I N S;

src/antlr/Parser.g

@@ -29,8 +29,6 @@ options {
 
     // enables parsing txn specific syntax when true
     protected boolean isParsingTxn = false;
-    // tracks whether a txn has conditional updates
-    protected boolean isTxnConditional = false;
 
     protected List<RowDataReference.Raw> references;
 
@@ -783,20 +781,33 @@ batchTxnStatement returns [TransactionStatement.Parsed expr]
         List<SelectStatement.RawStatement> assignments = new ArrayList<>();
         SelectStatement.RawStatement select = null;
         List<RowDataReference.Raw> returning = null;
-        List<ModificationStatement.Parsed> updates = new ArrayList<>();
+        List<List<ConditionStatement.Raw>> conditions = new ArrayList<>();
+        List<List<ModificationStatement.Parsed>> updates = new ArrayList<>();
     }
     : K_BEGIN K_TRANSACTION
       ( let=letStatement ';' { assignments.add(let); })*
       ( ( (selectStatement) => s=selectStatement ';' { select = s; }) | ( K_SELECT drs=rowDataReferences ';' { returning = drs; }) )?
-      ( K_IF conditions=txnConditions K_THEN { isTxnConditional = true; } )?
-      ( upd=batchStatementObjective ';' { updates.add(upd); } )*
-      ( {!isTxnConditional}? (K_COMMIT K_TRANSACTION) | {isTxnConditional}? (K_END K_IF K_COMMIT K_TRANSACTION))
+      (
+          K_IF c=txnConditions K_THEN u=updateStatements { conditions.add(c); updates.add(u); }
+          ( K_ELSE K_IF c=txnConditions K_THEN u=updateStatements { conditions.add(c); updates.add(u); } )*
+          ( K_ELSE u=updateStatements { conditions.add(Collections.singletonList(new ConditionStatement.Raw(null, ConditionStatement.Kind.ELSE, null))); updates.add(u); } )?
+          K_END K_IF
+      )?
+      ( (updateStatements) => u=updateStatements { updates.add(u); } )?
+      (K_COMMIT K_TRANSACTION)
     {
         $expr = new TransactionStatement.Parsed(assignments, select, returning, updates, conditions, references);
     }
     ;
     finally { isParsingTxn = false; }
 
+updateStatements returns [List<ModificationStatement.Parsed> updates]
+    @init {
+        updates = new ArrayList<ModificationStatement.Parsed>();
+    }
+    : (upd=batchStatementObjective ';' { updates.add(upd); })*
+    ;
+
 rowDataReferences returns [List<RowDataReference.Raw> refs]
     : r1=rowDataReference { refs = new ArrayList<RowDataReference.Raw>(); refs.add(r1); } (',' rN=rowDataReference { refs.add(rN); })*
     ;

src/java/org/apache/cassandra/cql3/transactions/ConditionStatement.java

@@ -40,13 +40,14 @@ public enum Kind
         GT(TxnCondition.Kind.GREATER_THAN, TxnCondition.Kind.LESS_THAN),
         GTE(TxnCondition.Kind.GREATER_THAN_OR_EQUAL, TxnCondition.Kind.LESS_THAN_OR_EQUAL),
         LT(TxnCondition.Kind.LESS_THAN, TxnCondition.Kind.GREATER_THAN),
-        LTE(TxnCondition.Kind.LESS_THAN_OR_EQUAL, TxnCondition.Kind.GREATER_THAN_OR_EQUAL);
-
+        LTE(TxnCondition.Kind.LESS_THAN_OR_EQUAL, TxnCondition.Kind.GREATER_THAN_OR_EQUAL),
+        ELSE(TxnCondition.Kind.ELSE, null);
+
         // TODO: Support for IN, CONTAINS, CONTAINS KEY
 
         private final TxnCondition.Kind kind;
         private final TxnCondition.Kind reversedKind;
-        
+
         Kind(TxnCondition.Kind kind, TxnCondition.Kind reversedKind)
         {
             this.kind = kind;
@@ -80,15 +81,18 @@ public static class Raw
 
         public Raw(Term.Raw lhs, Kind kind, Term.Raw rhs)
         {
-            Preconditions.checkArgument(lhs != null);
-            Preconditions.checkArgument((rhs == null) == (kind == Kind.IS_NOT_NULL || kind == Kind.IS_NULL));
+            Preconditions.checkArgument(lhs != null || kind == Kind.ELSE);
+            Preconditions.checkArgument((rhs == null) == (kind == Kind.IS_NOT_NULL || kind == Kind.IS_NULL || kind == Kind.ELSE));
             this.lhs = lhs;
             this.kind = kind;
             this.rhs = rhs;
         }
 
         public ConditionStatement prepare(String keyspace, VariableSpecifications bindVariables)
         {
+            if (kind == Kind.ELSE)
+                return new ConditionStatement(null, kind, null, false);
+
             if (rhs == null)
             {
                 // In the IS NULL/IS NOT NULL case, the reference will always be on the LHS
@@ -134,6 +138,8 @@ public TxnCondition createCondition(QueryOptions options)
     {
         switch (kind)
         {
+            case ELSE:
+                return TxnCondition.Else.instance;
             case IS_NOT_NULL:
             case IS_NULL:
                 return new TxnCondition.Exists(reference.toTxnReference(options), kind.toTxnKind(reversed));

src/java/org/apache/cassandra/service/accord/txn/TxnCondition.java

@@ -117,7 +117,8 @@ public enum Kind
         GREATER_THAN_OR_EQUAL(">=", Operator.GTE),
         LESS_THAN("<", Operator.LT),
         LESS_THAN_OR_EQUAL("<=", Operator.LTE),
-        COLUMN_CONDITIONS("COLUMN_CONDITIONS", null);
+        COLUMN_CONDITIONS("COLUMN_CONDITIONS", null),
+        ELSE("ELSE", null);
 
         @Nonnull
         private final String symbol;
@@ -152,6 +153,8 @@ private ConditionSerializer serializer()
                     return None.serializer;
                 case COLUMN_CONDITIONS:
                     return ColumnConditionsAdapter.serializer;
+                case ELSE:
+                    return Else.serializer;
                 default:
                     throw new IllegalArgumentException("No serializer exists for kind " + this);
             }
@@ -231,6 +234,43 @@ public static TxnCondition none()
         return None.instance;
     }
 
+    public static class Else extends TxnCondition
+    {
+        public static final Else instance = new Else();
+
+        private Else()
+        {
+            super(Kind.ELSE);
+        }
+
+        @Override
+        public String toString()
+        {
+            return kind.toString();
+        }
+
+        @Override
+        public void collect(TableMetadatas.Collector collector)
+        {
+        }
+
+        @Override
+        public boolean applies(TxnData data)
+        {
+            return true;
+        }
+
+        private static final ConditionSerializer<Else> serializer = new ConditionSerializer<>()
+        {
+            @Override
+            public void serialize(Else condition, TableMetadatas tables, DataOutputPlus out) {}
+            @Override
+            public Else deserialize(TableMetadatas tables, DataInputPlus in, Kind kind) { return instance; }
+            @Override
+            public long serializedSize(Else condition, TableMetadatas tables) { return 0; }
+        };
+    }
+
     public static class Exists extends TxnCondition
     {
         private static final Set<Kind> KINDS = ImmutableSet.of(Kind.IS_NOT_NULL, Kind.IS_NULL);

src/java/org/apache/cassandra/service/accord/txn/TxnUpdate.java

@@ -253,7 +253,7 @@ public long estimatedSizeOnHeap()
         }
     }
 
-    static class Block
+    public static class Block
     {
         private static final long EMPTY_SIZE = ObjectSizes.measure(new Block(null, null));
         public static final ParameterisedUnversionedSerializer<Block, TableMetadatasAndKeys> serializer = new ParameterisedUnversionedSerializer<>()
@@ -293,6 +293,11 @@ public long serializedSize(Block t, TableMetadatasAndKeys p)
         final BlockFragment[] fragments;
         final ConditionalBlock[] conditionalBlocks;
 
+        // Must ensure the following invariants when constructing a Block
+        // (1) Every fragment in fragments is sorted by key order
+        // (2) conditionalBlocks occurs in the same order as the IF/ELSE IF/ELSE statements because the order of the array
+        // determines the order in which the conditions get evaluated
+        // (3) Within each ConditionalBlock, fragmentIds is in sorted order
         Block(BlockFragment[] fragments, ConditionalBlock[] conditionalBlocks)
         {
             this.fragments = fragments;
@@ -372,26 +377,16 @@ public Block select(Keys keys)
 
             ConditionalBlock[] outConditions;
             {
-                List<ConditionalBlock> collect = null;
+                List<ConditionalBlock> collect = new ArrayList<>(conditionalBlocks.length - 1);
                 for (int i = 0 ; i < conditionalBlocks.length ; ++i)
                 {
                     ConditionalBlock cb = conditionalBlocks[i];
                     int[] cbOutFragmentIds = SortedArrays.linearIntersection(cb.fragmentIds, 0, cb.fragmentIds.length, outFragmentIds, 0, count, cachedInts());
-                    //noinspection ArrayEquality
-                    if (cbOutFragmentIds != cb.fragmentIds) // when arrays are equal the cb.fragmentIds gets returned unchanged, so can do a pointer check to detect a change
-                    {
-                        if (collect == null)
-                        {
-                            collect = new ArrayList<>(conditionalBlocks.length - 1);
-                            for (int j = 0 ; j < i ; ++j) //TODO (review): why do we include the previous blocks that "should" have empty fragments, but we provide them without empty fragments?
-                                collect.add(conditionalBlocks[j]);
-                        }
-                        if (cbOutFragmentIds.length > 0)
-                            collect.add(new ConditionalBlock(cb.id, cb.condition, cbOutFragmentIds));
-                    }
+
+                    if (cbOutFragmentIds.length > 0)
+                        collect.add(new ConditionalBlock(cb.id, cb.condition, cbOutFragmentIds));
                 }
-                if (collect == null) outConditions = conditionalBlocks;
-                else if (collect.isEmpty()) outConditions = NO_CONDITIONAL_BLOCKS;
+                if (collect.isEmpty()) outConditions = NO_CONDITIONAL_BLOCKS;
                 else outConditions = collect.toArray(ConditionalBlock[]::new);
             }
 
@@ -576,6 +571,127 @@ public TxnUpdate(TableMetadatas tables, List<Fragment> fragments, TxnCondition c
         this.preserveTimestamps = preserveTimestamps;
     }
 
+    public TxnUpdate(TableMetadatas tables, List<PreTransformedBlock> preTransformedBlocks, @Nullable ConsistencyLevel cassandraCommitCL, PreserveTimestamp preserveTimestamps)
+    {
+        requireArgument(cassandraCommitCL == null || IAccordService.SUPPORTED_COMMIT_CONSISTENCY_LEVELS.contains(cassandraCommitCL));
+        this.tables = tables;
+        this.keys = sortedFragmentKeys(preTransformedBlocks);
+
+        List<Block> blocks = new ArrayList<>();
+        int nextConditionIndex = 0;
+        for (PreTransformedBlock preTransformedBlock : preTransformedBlocks)
+        {
+            Pair<Integer, Block> pair = preTransformedBlock.generateBlock(nextConditionIndex, tables);
+            nextConditionIndex = pair.left();
+            blocks.add(pair.right());
+        }
+
+        this.blocks = blocks;
+        this.cassandraCommitCL = cassandraCommitCL;
+        this.preserveTimestamps = preserveTimestamps;
+    }
+
+    public Keys sortedFragmentKeys(List<PreTransformedBlock> preTransformedBlocks)
+    {
+        List<Key> keys = new ArrayList<>();
+        for (PreTransformedBlock preTransformedBlock : preTransformedBlocks)
+            keys.addAll(preTransformedBlock.getKeys());
+
+        keys.sort(RoutableKey::compareTo);
+        return Keys.of(keys, k -> k);
+    }
+
+    public static class PreTransformedBlock
+    {
+        public TxnCondition[] conditions;
+        @Nullable
+        public List<Pair<TxnWrite.Fragment, Integer>> fragmentConditionIndexPair;
+        @Nullable
+        public List<TxnWrite.Fragment> noneConditions;
+
+        public PreTransformedBlock(TxnCondition[] conditions, List<Pair<TxnWrite.Fragment, Integer>> fragmentConditionIndexPair, List<TxnWrite.Fragment> noneConditions)
+        {
+            Invariants.require((fragmentConditionIndexPair == null && noneConditions != null) || (fragmentConditionIndexPair != null && noneConditions == null));
+            this.conditions = conditions;
+            if (fragmentConditionIndexPair != null)
+                fragmentConditionIndexPair.sort((l, r) -> TxnWrite.Fragment.compareKeys(l.left(), r.left()));
+            this.fragmentConditionIndexPair = fragmentConditionIndexPair;
+            if (noneConditions != null)
+                noneConditions.sort(Fragment::compareKeys);
+            this.noneConditions = noneConditions;
+        }
+
+        private boolean isTrailingUpdate()
+        {
+            return noneConditions != null;
+        }
+
+        public List<Key> getKeys()
+        {
+            List<Key> keys = new ArrayList<>();
+            if (isTrailingUpdate())
+            {
+                for (int i = 0; i < noneConditions.size(); i++)
+                    keys.add(noneConditions.get(i).key);
+            }
+            else {
+                for (int i = 0; i < fragmentConditionIndexPair.size(); i++)
+                    keys.add(fragmentConditionIndexPair.get(i).left().key);
+            }
+            return keys;
+        }
+
+        public Pair<Integer, Block> generateBlock(int conditionalBlockIndex, TableMetadatas tables)
+        {
+            if (isTrailingUpdate())
+            {
+                BlockFragment[] blockFragments = new BlockFragment[noneConditions.size()];
+                int[] fragmentIds = new int[noneConditions.size()];
+                for (int i = 0 ; i < noneConditions.size() ; ++i)
+                {
+                    Fragment fragment = noneConditions.get(i);
+                    blockFragments[i] = new BlockFragment(i, fragment.key, Fragment.FragmentSerializer.serialize(fragment, tables, Version.LATEST));
+                    fragmentIds[i] = i;
+                }
+
+                SerializedTxnCondition serializedCondition = new SerializedTxnCondition(TxnCondition.none(), tables);
+                ConditionalBlock[] conditionalBlock = new ConditionalBlock[] { new ConditionalBlock(conditionalBlockIndex++, serializedCondition, fragmentIds) };
+                return Pair.create(conditionalBlockIndex, new Block(blockFragments, conditionalBlock));
+            }
+            else
+            {
+                List<List<Integer>> conditionsIndexToFragmentId = new ArrayList<>(conditions.length);
+                for (int i = 0; i < conditions.length; i++)
+                    conditionsIndexToFragmentId.add(new ArrayList<>());
+
+                BlockFragment[] blockFragments = new BlockFragment[fragmentConditionIndexPair.size()];
+                ConditionalBlock[] conditionalBlocks = new ConditionalBlock[conditions.length];
+
+                int blockFragmentIdx = 0;
+                for (Pair<Fragment, Integer> pair : fragmentConditionIndexPair)
+                {
+                    Fragment fragment = pair.left;
+                    int conditionsIndex = pair.right;
+                    blockFragments[blockFragmentIdx] = new BlockFragment(blockFragmentIdx, fragment.key, Fragment.FragmentSerializer.serialize(fragment, tables, Version.LATEST));
+                    conditionsIndexToFragmentId.get(conditionsIndex).add(blockFragmentIdx);
+                    blockFragmentIdx++;
+                }
+
+                int conditionalBlocksIdx = 0;
+
+                for (int i = 0; i < conditions.length; i++)
+                {
+                    int[] fragmentIds = conditionsIndexToFragmentId.get(i).stream().mapToInt(x -> x).toArray();
+                    SerializedTxnCondition serializedCondition = new SerializedTxnCondition(conditions[i], tables);
+                    conditionalBlocks[conditionalBlocksIdx] = new ConditionalBlock(conditionalBlocksIdx, serializedCondition, fragmentIds);
+                    conditionalBlocksIdx++;
+                }
+
+                return Pair.create(conditionalBlocksIdx, new Block(blockFragments, conditionalBlocks));
+            }
+        }
+    }
+
     private TxnUpdate(TableMetadatas tables, Keys keys, List<Block> blocks, ConsistencyLevel cassandraCommitCL, PreserveTimestamp preserveTimestamps)
     {
         this.tables = tables;