natas_hw
/
cassandra
mirror of https://github.com/apache/cassandra
0
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

added support for let evaluations

This commit is contained in:
Alan Wang 2026-06-26 12:18:34 -07:00
parent 4aebfde38c
commit 3b4258b978
10 changed files with 166 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/antlr/Parser.g
View File

@ -2111,7 +2111,7 @@ normalColumnOperation[UpdateStatement.OperationCollector operations, ColumnIdent
addRecognitionError("Only expressions of the form X = X " + ($i.text.charAt(0) == '-' ? '-' : '+') + " <value> are supported.");
addRawUpdate(operations, key, new Operation.Addition(Constants.Literal.integer($i.text)));
}
| {isParsingTxn}? r=dottedRowDataReference (sig=('+'|'-') t=term)?
| {isParsingTxn}? r=dottedRowDataReference (sig=('+'|'-') t=term)? // Todo: This parses row1.v + 2, but we still need to handle the case for 2 + row1.v and also key + row1.v
{
if (t == null)
{

6
src/java/org/apache/cassandra/cql3/Operation.java
View File

@ -17,6 +17,7 @@
*/
package org.apache.cassandra.cql3;
import java.nio.ByteBuffer;
import java.util.List;
import org.apache.cassandra.cql3.functions.Function;
@ -120,6 +121,11 @@ public abstract class Operation
*/
public abstract void execute(DecoratedKey partitionKey, RowUpdateBuilder builder) throws InvalidRequestException;
public void execute(DecoratedKey partitionKey, RowUpdateBuilder builder, ByteBuffer term) throws InvalidRequestException
{
throw new UnsupportedOperationException();
}
/**
* A parsed raw UPDATE operation.
*

2
src/java/org/apache/cassandra/cql3/statements/UpdateStatement.java
View File

@ -214,7 +214,7 @@ public class UpdateStatement extends ModificationStatement
{
ReferenceValue.Raw raw = (ReferenceValue.Raw) value;
ReferenceValue referenceValue = raw.prepare(def, bindVariables);
ReferenceOperation operation = new ReferenceOperation(def, metadata, TxnReferenceOperation.Kind.setterFor(def), null, null, referenceValue);
ReferenceOperation operation = new ReferenceOperation(def, metadata, TxnReferenceOperation.Kind.setterFor(def), null, null, null, referenceValue);
operations.add(def, operation);
}
else

40
src/java/org/apache/cassandra/cql3/terms/Constants.java
View File

@ -21,6 +21,8 @@ import java.math.BigDecimal;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import accord.utils.Invariants;
import org.apache.cassandra.cql3.AssignmentTestable;
import org.apache.cassandra.cql3.CQL3Type;
import org.apache.cassandra.cql3.ColumnSpecification;
@ -550,6 +552,30 @@ public abstract class Constants
builder.addCell(column, newValue);
}
}
public void execute(DecoratedKey partitionKey, RowUpdateBuilder builder, ByteBuffer constant) throws InvalidRequestException
{
Invariants.require(constant != null);
if (column.type instanceof NumberType<?>)
{
@SuppressWarnings("unchecked") NumberType<Number> type = (NumberType<Number>) column.type;
ByteBuffer increment = type.sanitize(t.bindAndGet(builder));
if (increment == null)
return;
ByteBuffer newValue = type.add(type.compose(constant), type.compose(increment));
builder.addCell(column, newValue);
}
else if (column.type instanceof StringType)
{
ByteBuffer left = t.bindAndGet(builder);
if (left == null)
return;
ByteBuffer newValue = ByteBuffer.allocate(left.remaining() + constant.remaining());
FastByteOperations.copy(left, left.position(), newValue, newValue.position(), left.remaining());
FastByteOperations.copy(constant, constant.position(), newValue, newValue.position() + left.remaining(), constant.remaining());
builder.addCell(column, newValue);
}
}
}
public static class Substracter extends Operation
@ -594,6 +620,20 @@ public abstract class Constants
builder.addCell(column, newValue);
}
}
public void execute(DecoratedKey partitionKey, RowUpdateBuilder builder, ByteBuffer constant) throws InvalidRequestException
{
Invariants.require(constant != null);
if (column.type instanceof NumberType<?>)
{
@SuppressWarnings("unchecked") NumberType<Number> type = (NumberType<Number>) column.type;
ByteBuffer increment = type.sanitize(t.bindAndGet(builder));
if (increment == null)
return;
ByteBuffer newValue = type.substract(type.compose(increment), type.compose(constant));
builder.addCell(column, newValue);
}
}
}
// This happens to also handle collection because it doesn't felt worth

14
src/java/org/apache/cassandra/cql3/transactions/ReferenceOperation.java
View File

@ -48,16 +48,18 @@ public class ReferenceOperation
private final TableMetadata table;
private final TxnReferenceOperation.Kind kind;
private final FieldIdentifier field;
private final Term constant;
private final Term key;
private final ReferenceValue value;
public ReferenceOperation(ColumnMetadata receiver, TableMetadata table, TxnReferenceOperation.Kind kind, Term key, FieldIdentifier field, ReferenceValue value)
public ReferenceOperation(ColumnMetadata receiver, TableMetadata table, TxnReferenceOperation.Kind kind, Term key, FieldIdentifier field, Term constant, ReferenceValue value)
{
this.receiver = receiver;
this.table = table;
this.kind = kind;
this.key = key;
this.field = field;
this.constant = constant;
this.value = value;
}
@ -75,7 +77,8 @@ public class ReferenceOperation
ReferenceValue value = new ReferenceValue.Constant(operation.term());
Term key = extractKeyOrIndex(operation);
FieldIdentifier field = extractField(operation);
return new ReferenceOperation(receiver, table, kind, key, field, value);
// Chore: Maybe we should change the semantics of this?
return new ReferenceOperation(receiver, table, kind, key, field, null, value);
}
public TxnReferenceOperation.Kind getKind()
@ -105,6 +108,7 @@ public class ReferenceOperation
receiver, table,
key != null ? key.bindAndGet(options) : null,
field != null ? field.bytes : null,
constant != null ? constant.bindAndGet(options) : null,
value.bindAndGet(options));
}
@ -157,7 +161,11 @@ public class ReferenceOperation
}
}
return new ReferenceOperation(receiver, metadata, kind, key, field, value.prepare(valueReceiver, bindVariables));
ReferenceValue referenceValue = value.prepare(valueReceiver, bindVariables);
// operation.term().equals(referenceValue.getTerm()) covers the case where we have
// v += row1.c,
return new ReferenceOperation(receiver, metadata, kind, key, field, operation.term().equals(referenceValue.getTerm()) ? null : operation.term(), referenceValue);
}
}

14
src/java/org/apache/cassandra/cql3/transactions/ReferenceValue.java
View File

@ -33,6 +33,8 @@ public abstract class ReferenceValue
{
public abstract TxnReferenceValue bindAndGet(FunctionContext context);
public abstract Term getTerm();
public static abstract class Raw extends Term.Raw
{
public abstract ReferenceValue prepare(ColumnMetadata receiver, VariableSpecifications bindVariables);
@ -53,6 +55,12 @@ public abstract class ReferenceValue
return new TxnReferenceValue.Constant(term.bindAndGet(context));
}
@Override
public Term getTerm()
{
return term;
}
public static class Raw extends ReferenceValue.Raw
{
private final Term.Raw term;
@ -109,6 +117,12 @@ public abstract class ReferenceValue
return new TxnReferenceValue.Substitution(reference.toTxnReference(context).asColumn());
}
@Override
public Term getTerm()
{
return reference;
}
public static class Raw extends ReferenceValue.Raw
{
private final RowDataReference.Raw reference;

14
src/java/org/apache/cassandra/cql3/transactions/RowDataReference.java
View File

@ -20,6 +20,7 @@ package org.apache.cassandra.cql3.transactions;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import com.google.common.base.Preconditions;
@ -101,6 +102,19 @@ public class RowDataReference extends Term.NonTerminal
throw new UnsupportedOperationException("Functions are not currently supported w/ reference terms.");
}
@Override
public boolean equals(Object o)
{
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
RowDataReference that = (RowDataReference) o;
return Objects.equals(txnDataName, that.txnDataName)
&& Objects.equals(column, that.column)
&& Objects.equals(table, that.table)
&& Objects.equals(elementPath, that.elementPath)
&& Objects.equals(fieldPath, that.fieldPath);
}
public ColumnMetadata toResultMetadata()
{
ColumnIdentifier fullName = getFullyQualifiedName();

20
src/java/org/apache/cassandra/service/accord/txn/TxnReferenceOperation.java
View File

@ -170,18 +170,20 @@ public class TxnReferenceOperation
public final TableMetadata table;
private final @Nullable ByteBuffer keyOrIndex;
private final @Nullable ByteBuffer field;
private final @Nullable ByteBuffer constant;
private final TxnReferenceValue value;
private final @Nullable AbstractType<?> keyOrIndexType;
private final AbstractType<?> valueType;
public TxnReferenceOperation(Kind kind, ColumnMetadata receiver, TableMetadata table,
@Nullable ByteBuffer keyOrIndex, @Nullable ByteBuffer field, TxnReferenceValue value)
@Nullable ByteBuffer keyOrIndex, @Nullable ByteBuffer field, @Nullable ByteBuffer constant, TxnReferenceValue value)
{
this.kind = kind;
this.receiver = receiver;
this.table = table;
this.keyOrIndex = keyOrIndex;
this.field = field;
this.constant = constant;
// We don't expect operators on clustering keys, but unwrap just in case.
AbstractType<?> receiverType = receiver.type.unwrap();
@ -272,7 +274,11 @@ public class TxnReferenceOperation
public void apply(TxnData data, DecoratedKey key, RowUpdateBuilder up)
{
Operation operation = toOperation(data);
operation.execute(key, up);
// When constant != null, we are performing a computation with a LET variable (i.e. row1.v + 2)
if (constant != null)
operation.execute(key, up, constant);
else
operation.execute(key, up);
}
@VisibleForTesting
@ -321,6 +327,9 @@ public class TxnReferenceOperation
out.writeBoolean(operation.field != null);
if (operation.field != null)
ByteBufferUtil.writeWithVIntLength(operation.field, out);
out.writeBoolean(operation.constant != null);
if (operation.constant != null)
ByteBufferUtil.writeWithVIntLength(operation.constant, out);
}
@Override
@ -332,7 +341,8 @@ public class TxnReferenceOperation
TxnReferenceValue value = TxnReferenceValue.serializer.deserialize(tables, in);
ByteBuffer key = in.readBoolean() ? ByteBufferUtil.readWithVIntLength(in) : null;
ByteBuffer field = in.readBoolean() ? ByteBufferUtil.readWithVIntLength(in) : null;
return new TxnReferenceOperation(kind, receiver, table, key, field, value);
ByteBuffer constant = in.readBoolean() ? ByteBufferUtil.readWithVIntLength(in) : null;
return new TxnReferenceOperation(kind, receiver, table, key, field, constant, value);
}
@Override
@ -351,6 +361,10 @@ public class TxnReferenceOperation
if (operation.field != null)
size += ByteBufferUtil.serializedSizeWithVIntLength(operation.field);
size += TypeSizes.sizeof(operation.constant != null);
if (operation.constant != null)
size += ByteBufferUtil.serializedSizeWithVIntLength(operation.constant);
return size;
}
};

80
test/distributed/org/apache/cassandra/distributed/test/accord/AccordCQLTestBase.java
View File

@ -84,6 +84,7 @@ import org.apache.cassandra.distributed.shared.AssertUtils;
import org.apache.cassandra.distributed.test.sai.SAIUtil;
import org.apache.cassandra.distributed.util.QueryResultUtil;
import org.apache.cassandra.exceptions.InvalidRequestException;
import org.apache.cassandra.exceptions.SyntaxException;
import org.apache.cassandra.exceptions.WriteTimeoutException;
import org.apache.cassandra.io.util.DataInputBuffer;
import org.apache.cassandra.schema.SchemaConstants;
@ -3606,41 +3607,66 @@ public abstract class AccordCQLTestBase extends AccordTestBase
}
@Test
public void testRowReferenceDoesNotNPE() throws Exception
public void testLETVariableReferenceInUpdateFails() throws Exception
{
test("CREATE TABLE " + qualifiedAccordTableName + " (k int, c int, v int, PRIMARY KEY (k, c)) WITH CLUSTERING ORDER BY (c DESC) AND transactional_mode='" + transactionalMode + "'",
cluster ->
{
cluster.coordinator(1).execute("INSERT INTO " + qualifiedAccordTableName + " (k, c, v) VALUES (1, 1, 1)", ConsistencyLevel.ALL);
// Regression test for prior NPE
test("CREATE TABLE " + qualifiedAccordTableName + " (k int, c int, v int, PRIMARY KEY (k, c)) WITH " + transactionalMode.asCqlParam(), cluster -> {
try
{
String txn = "BEGIN TRANSACTION\n" +
" LET r = (SELECT v FROM " + qualifiedAccordTableName + " WHERE k = 1 AND c = 1);\n" +
" UPDATE " + qualifiedAccordTableName + " SET v = r WHERE k=1 AND c=1;\n" +
"COMMIT TRANSACTION";
String update = "BEGIN TRANSACTION\n" +
" LET r = (SELECT v FROM " + qualifiedAccordTableName + " WHERE k = 1 AND c = 1);\n" +
" UPDATE " + qualifiedAccordTableName + " SET v = r WHERE k=1 AND c=1;\n" +
"COMMIT TRANSACTION";
assertRowEqualsWithPreemptedRetry(cluster, new Object[]{1, 1, 1}, update);
});
cluster.coordinator(1).executeWithResult(txn, ConsistencyLevel.SERIAL);
fail("Expected exception");
}
catch (Throwable t)
{
assertEquals(SyntaxException.class.getName(), t.getClass().getName());
}
});
}
@Test
public void testUseLetVariableForEvaluationInUpdateStatement() throws Exception
public void testUseLetVariableForEvaluationWithInt() throws Exception
{
test("CREATE TABLE " + qualifiedAccordTableName + " (k int, c int, v int, PRIMARY KEY (k, c)) WITH CLUSTERING ORDER BY (c DESC) AND transactional_mode='" + transactionalMode + "'",
cluster ->
{
cluster.coordinator(1).execute("INSERT INTO " + qualifiedAccordTableName + " (k, c, v) VALUES (1, 1, 1)", ConsistencyLevel.ALL);
test("CREATE TABLE " + qualifiedAccordTableName + " (k int, c int, v int, PRIMARY KEY (k, c)) WITH " + transactionalMode.asCqlParam(), cluster -> {
cluster.coordinator(1).execute("INSERT INTO " + qualifiedAccordTableName + " (k, c, v) VALUES (1, 1, 6)", ConsistencyLevel.ALL);
String update = "BEGIN TRANSACTION\n" +
" LET row1 = (SELECT * FROM " + qualifiedAccordTableName + " WHERE k = 1 AND c = 1);\n" +
" UPDATE " + qualifiedAccordTableName + " SET v = row1.v + 2 WHERE k=1 AND c=1;\n" +
"COMMIT TRANSACTION";
cluster.coordinator(1).executeWithResult(update, ConsistencyLevel.ALL);
String update = "BEGIN TRANSACTION\n" +
" LET row1 = (SELECT * FROM " + qualifiedAccordTableName + " WHERE k = 1 AND c = 1);\n" +
" UPDATE " + qualifiedAccordTableName + " SET v = row1.v + 3 WHERE k = 1 AND c = 1;\n" +
"COMMIT TRANSACTION";
cluster.coordinator(1).executeWithResult(update, ConsistencyLevel.ALL);
String read = "BEGIN TRANSACTION\n" +
"SELECT * FROM " + qualifiedAccordTableName + " WHERE k = 1;\n" +
"COMMIT TRANSACTION";
String read = "BEGIN TRANSACTION\n" +
"SELECT * FROM " + qualifiedAccordTableName + " WHERE k = 1;\n" +
"COMMIT TRANSACTION";
SimpleQueryResult result = cluster.coordinator(1).executeWithResult(read, ConsistencyLevel.SERIAL);
assertThat(result).hasSize(1).contains(1, 3);
});
SimpleQueryResult result = cluster.coordinator(1).executeWithResult(read, ConsistencyLevel.SERIAL);
assertThat(result).hasSize(1).contains(1, 1, 11);
});
}
@Test
public void testUseLetVariableForEvaluationWithString() throws Exception
{
test("CREATE TABLE " + qualifiedAccordTableName + " (k int, c int, v text, PRIMARY KEY (k, c)) WITH " + transactionalMode.asCqlParam(), cluster -> {
cluster.coordinator(1).execute("INSERT INTO " + qualifiedAccordTableName + " (k, c, v) VALUES (?, ?, ?)", ConsistencyLevel.ALL, 1, 1, "a");
String update = "BEGIN TRANSACTION\n" +
" LET row1 = (SELECT * FROM " + qualifiedAccordTableName + " WHERE k = 1 AND c = 1);\n" +
" UPDATE " + qualifiedAccordTableName + " SET v = row1.v + ? WHERE k = 1 AND c = 1;\n" +
"COMMIT TRANSACTION";
cluster.coordinator(1).executeWithResult(update, ConsistencyLevel.ALL, "m");
String read = "BEGIN TRANSACTION\n" +
"SELECT * FROM " + qualifiedAccordTableName + " WHERE k = 1;\n" +
"COMMIT TRANSACTION";
SimpleQueryResult result = cluster.coordinator(1).executeWithResult(read, ConsistencyLevel.SERIAL);
assertThat(result).hasSize(1).contains(1, 1, "am");
});
}
}

14
test/unit/org/apache/cassandra/service/accord/txn/TxnReferenceOperationTest.java
View File

@ -55,9 +55,9 @@ import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.service.accord.serializers.TableMetadatas;
import org.apache.cassandra.utils.AbstractTypeGenerators;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.cassandra.utils.Generators;
import static accord.utils.Property.qt;
import static org.apache.cassandra.utils.Generators.toGen;
public class TxnReferenceOperationTest
{
@ -110,7 +110,7 @@ public class TxnReferenceOperationTest
private static Gen<TxnReferenceValue> valueGen(AbstractType<?> type)
{
return Generators.toGen(AbstractTypeGenerators.getTypeSupport(type)
return toGen(AbstractTypeGenerators.getTypeSupport(type)
.bytesGen())
.map(TxnReferenceValue.Constant::new);
}
@ -123,6 +123,7 @@ public class TxnReferenceOperationTest
TableMetadata table;
@Nullable ByteBuffer keyOrIndex = null;
@Nullable ByteBuffer field = null;
@Nullable ByteBuffer constant = null;
TxnReferenceValue value;
Group group = rs.pick(Group.values());
switch (group)
@ -139,7 +140,7 @@ public class TxnReferenceOperationTest
break;
case Discarder:
{
CollectionType<?> type = (CollectionType<?>) Generators.toGen(AbstractTypeGenerators.builder()
CollectionType<?> type = (CollectionType<?>) toGen(AbstractTypeGenerators.builder()
.withMaxDepth(1)
.withTypeKinds(AbstractTypeGenerators.TypeKind.LIST,
AbstractTypeGenerators.TypeKind.SET,
@ -195,13 +196,15 @@ public class TxnReferenceOperationTest
table = table(type);
receiver = table.getColumn(ColumnIdentifier.getInterned("col", true));
value = valueGen(type).next(rs);
if (rs.nextBoolean())
constant = toGen(AbstractTypeGenerators.getTypeSupport(type).bytesGen()).next(rs);
kind = group == Group.Adder ? TxnReferenceOperation.Kind.ConstantAdder : TxnReferenceOperation.Kind.ConstantSubtracter;
}
}
break;
case Setter:
{
var type = Generators.toGen(AbstractTypeGenerators.builder()
var type = toGen(AbstractTypeGenerators.builder()
.withoutUnsafeEquality()
.withMaxDepth(1)
.build())
@ -276,7 +279,8 @@ public class TxnReferenceOperationTest
default:
throw new UnsupportedOperationException();
}
return new TxnReferenceOperation(kind, receiver, table, keyOrIndex, field, value);
return new TxnReferenceOperation(kind, receiver, table, keyOrIndex, field, constant, value);
};
}