diff --git a/CHANGES.txt b/CHANGES.txt index d79d78bad2..caea458186 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -1,4 +1,5 @@ 5.0 + * Add support for vectors in UDFs (CASSANDRA-18613) * Improve vector value validation errors (CASSANDRA-18652) * Upgrade Guava to 32.0.1 (CASSANDRA-18645) * Add duration and count of partition keys to sstablemetadata (CASSANDRA-18639) diff --git a/src/java/org/apache/cassandra/cql3/functions/types/AbstractAddressableByIndexData.java b/src/java/org/apache/cassandra/cql3/functions/types/AbstractAddressableByIndexData.java index 0e98c8948f..2b1fd84a9e 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/AbstractAddressableByIndexData.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/AbstractAddressableByIndexData.java @@ -253,6 +253,12 @@ extends AbstractGettableByIndexData implements SettableByIndexData return setValue(i, codecFor(i, TypeTokens.setOf(elementsType)).serialize(v, protocolVersion)); } + @Override + public T setVector(int i, List v) + { + return setValue(i, codecFor(i).serialize(v, protocolVersion)); + } + @Override public T setUDTValue(int i, UDTValue v) { diff --git a/src/java/org/apache/cassandra/cql3/functions/types/AbstractData.java b/src/java/org/apache/cassandra/cql3/functions/types/AbstractData.java index adb7c0e0bd..45ce309332 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/AbstractData.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/AbstractData.java @@ -540,6 +540,22 @@ implements SettableData return wrapped; } + @Override + public T setVector(int i, List v) + { + return setValue(i, codecFor(i).serialize(v, protocolVersion)); + } + + @Override + public T setVector(String name, List v) + { + for (int i : getAllIndexesOf(name)) + { + setVector(i, v); + } + return wrapped; + } + @Override public T setUDTValue(int i, UDTValue v) { diff --git a/src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableByIndexData.java b/src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableByIndexData.java index 1552309edb..ca8827b0d7 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableByIndexData.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableByIndexData.java @@ -365,6 +365,28 @@ abstract class AbstractGettableByIndexData implements GettableByIndexData return codecFor(i, javaType).deserialize(value, protocolVersion); } + /** + * {@inheritDoc} + */ + @Override + @SuppressWarnings("unchecked") + public List getVector(int i, Class elementsClass) + { + return getList(i, TypeToken.of(elementsClass)); + } + + /** + * {@inheritDoc} + */ + @Override + @SuppressWarnings("unchecked") + public List getVector(int i, TypeToken elementsType) + { + ByteBuffer value = getValue(i); + TypeToken> javaType = TypeTokens.listOf(elementsType); + return codecFor(i, javaType).deserialize(value, protocolVersion); + } + /** * {@inheritDoc} */ diff --git a/src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableData.java b/src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableData.java index 2ac7d980f4..f26abd296c 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableData.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableData.java @@ -269,6 +269,24 @@ implements GettableData return getMap(getIndexOf(name), keysType, valuesType); } + /** + * {@inheritDoc} + */ + @Override + public List getVector(String name, Class elementsClass) + { + return getList(getIndexOf(name), elementsClass); + } + + /** + * {@inheritDoc} + */ + @Override + public List getVector(String name, TypeToken elementsType) + { + return getList(getIndexOf(name), elementsType); + } + /** * {@inheritDoc} */ diff --git a/src/java/org/apache/cassandra/cql3/functions/types/CodecRegistry.java b/src/java/org/apache/cassandra/cql3/functions/types/CodecRegistry.java index a979a57269..706ef41a41 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/CodecRegistry.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/CodecRegistry.java @@ -165,27 +165,27 @@ public final class CodecRegistry static { - BUILT_IN_CODECS_MAP.put(DataType.Name.ASCII, TypeCodec.ascii()); - BUILT_IN_CODECS_MAP.put(DataType.Name.BIGINT, TypeCodec.bigint()); - BUILT_IN_CODECS_MAP.put(DataType.Name.BLOB, TypeCodec.blob()); - BUILT_IN_CODECS_MAP.put(DataType.Name.BOOLEAN, TypeCodec.cboolean()); - BUILT_IN_CODECS_MAP.put(DataType.Name.COUNTER, TypeCodec.counter()); - BUILT_IN_CODECS_MAP.put(DataType.Name.DECIMAL, TypeCodec.decimal()); - BUILT_IN_CODECS_MAP.put(DataType.Name.DOUBLE, TypeCodec.cdouble()); - BUILT_IN_CODECS_MAP.put(DataType.Name.FLOAT, TypeCodec.cfloat()); - BUILT_IN_CODECS_MAP.put(DataType.Name.INET, TypeCodec.inet()); - BUILT_IN_CODECS_MAP.put(DataType.Name.INT, TypeCodec.cint()); - BUILT_IN_CODECS_MAP.put(DataType.Name.TEXT, TypeCodec.varchar()); - BUILT_IN_CODECS_MAP.put(DataType.Name.TIMESTAMP, TypeCodec.timestamp()); - BUILT_IN_CODECS_MAP.put(DataType.Name.UUID, TypeCodec.uuid()); - BUILT_IN_CODECS_MAP.put(DataType.Name.VARCHAR, TypeCodec.varchar()); - BUILT_IN_CODECS_MAP.put(DataType.Name.VARINT, TypeCodec.varint()); - BUILT_IN_CODECS_MAP.put(DataType.Name.TIMEUUID, TypeCodec.timeUUID()); - BUILT_IN_CODECS_MAP.put(DataType.Name.SMALLINT, TypeCodec.smallInt()); - BUILT_IN_CODECS_MAP.put(DataType.Name.TINYINT, TypeCodec.tinyInt()); - BUILT_IN_CODECS_MAP.put(DataType.Name.DATE, TypeCodec.date()); - BUILT_IN_CODECS_MAP.put(DataType.Name.TIME, TypeCodec.time()); - BUILT_IN_CODECS_MAP.put(DataType.Name.DURATION, TypeCodec.duration()); + BUILT_IN_CODECS_MAP.put(ASCII, TypeCodec.ascii()); + BUILT_IN_CODECS_MAP.put(BIGINT, TypeCodec.bigint()); + BUILT_IN_CODECS_MAP.put(BLOB, TypeCodec.blob()); + BUILT_IN_CODECS_MAP.put(BOOLEAN, TypeCodec.cboolean()); + BUILT_IN_CODECS_MAP.put(COUNTER, TypeCodec.counter()); + BUILT_IN_CODECS_MAP.put(DECIMAL, TypeCodec.decimal()); + BUILT_IN_CODECS_MAP.put(DOUBLE, TypeCodec.cdouble()); + BUILT_IN_CODECS_MAP.put(FLOAT, TypeCodec.cfloat()); + BUILT_IN_CODECS_MAP.put(INET, TypeCodec.inet()); + BUILT_IN_CODECS_MAP.put(INT, TypeCodec.cint()); + BUILT_IN_CODECS_MAP.put(TEXT, TypeCodec.varchar()); + BUILT_IN_CODECS_MAP.put(TIMESTAMP, TypeCodec.timestamp()); + BUILT_IN_CODECS_MAP.put(UUID, TypeCodec.uuid()); + BUILT_IN_CODECS_MAP.put(VARCHAR, TypeCodec.varchar()); + BUILT_IN_CODECS_MAP.put(VARINT, TypeCodec.varint()); + BUILT_IN_CODECS_MAP.put(TIMEUUID, TypeCodec.timeUUID()); + BUILT_IN_CODECS_MAP.put(SMALLINT, TypeCodec.smallInt()); + BUILT_IN_CODECS_MAP.put(TINYINT, TypeCodec.tinyInt()); + BUILT_IN_CODECS_MAP.put(DATE, TypeCodec.date()); + BUILT_IN_CODECS_MAP.put(TIME, TypeCodec.time()); + BUILT_IN_CODECS_MAP.put(DURATION, TypeCodec.duration()); } // roughly sorted by popularity @@ -318,6 +318,16 @@ public final class CodecRegistry } return weight; } + case VECTOR: + { + int weight = level; + DataType eltType = cqlType.getTypeArguments().get(0); + if (eltType != null) + { + weight += weigh(eltType, level + 1); + } + return weight == 0 ? 1 : weight; + } case UDT: { int weight = level; @@ -746,6 +756,13 @@ public final class CodecRegistry return (TypeCodec) TypeCodec.map(keyCodec, valueCodec); } + if (cqlType instanceof VectorType + && (javaType == null || List.class.isAssignableFrom(javaType.getRawType()))) + { + VectorType type = (VectorType) cqlType; + return (TypeCodec) TypeCodec.vector(type, findCodec(type.getSubtype(), null)); + } + if (cqlType instanceof TupleType && (javaType == null || TupleValue.class.isAssignableFrom(javaType.getRawType()))) { @@ -847,14 +864,20 @@ public final class CodecRegistry } } - if ((cqlType == null || cqlType.getName() == DataType.Name.TUPLE) + if ((cqlType == null || cqlType.getName() == VECTOR) && value instanceof List) + { + VectorType type = (VectorType) cqlType; + return (TypeCodec) TypeCodec.vector(type, findCodec(type.getSubtype(), null)); + } + + if ((cqlType == null || cqlType.getName() == TUPLE) && value instanceof TupleValue) { return (TypeCodec) TypeCodec.tuple(cqlType == null ? ((TupleValue) value).getType() : (TupleType) cqlType); } - if ((cqlType == null || cqlType.getName() == DataType.Name.UDT) && value instanceof UDTValue) + if ((cqlType == null || cqlType.getName() == UDT) && value instanceof UDTValue) { return (TypeCodec) TypeCodec.userType(cqlType == null ? ((UDTValue) value).getType() : (UserType) cqlType); diff --git a/src/java/org/apache/cassandra/cql3/functions/types/DataType.java b/src/java/org/apache/cassandra/cql3/functions/types/DataType.java index 54127202b1..fafc721b99 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/DataType.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/DataType.java @@ -74,7 +74,8 @@ public abstract class DataType MAP(33), SET(34), UDT(48, ProtocolVersion.V3), - TUPLE(49, ProtocolVersion.V3); + TUPLE(49, ProtocolVersion.V3), + VECTOR(50, ProtocolVersion.V5); final int protocolId; @@ -437,6 +438,18 @@ public abstract class DataType return map(keyType, valueType, false); } + /** + * Returns the type of vector of {@code elementType} elements with {@code dimensions} dimensions. + * + * @param elementType the type of the vector elements. + * @param dimensions the number of dimensions of the vector. + * @return the type of vectors of {@code elementType} elements and {@code dimensions} dimensions. + */ + public static VectorType vector(DataType elementType, int dimensions) + { + return new VectorType(elementType, dimensions); + } + /** * Returns a Custom type. * diff --git a/src/java/org/apache/cassandra/cql3/functions/types/DataTypeClassNameParser.java b/src/java/org/apache/cassandra/cql3/functions/types/DataTypeClassNameParser.java index 7064ba22c7..1d50edfd86 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/DataTypeClassNameParser.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/DataTypeClassNameParser.java @@ -53,6 +53,7 @@ public class DataTypeClassNameParser private static final String UDT_TYPE = "org.apache.cassandra.db.marshal.UserType"; private static final String TUPLE_TYPE = "org.apache.cassandra.db.marshal.TupleType"; private static final String DURATION_TYPE = "org.apache.cassandra.db.marshal.DurationType"; + private static final String VECTOR_TYPE = "org.apache.cassandra.db.marshal.VectorType"; private static final ImmutableMap cassTypeToDataType = new ImmutableMap.Builder() @@ -118,6 +119,14 @@ public class DataTypeClassNameParser "Got o.a.c.db.marshal.FrozenType for something else than a collection, " + "this driver version might be too old for your version of Cassandra"); + if (isVectorType(next)) + { + List parameters = parser.getTypeParameters(); + DataType subtype = parseOne(parameters.get(0), protocolVersion, codecRegistry); + int dimensions = Integer.parseInt(parameters.get(1)); + return DataType.vector(subtype, dimensions); + } + if (isUserType(next)) { ++parser.idx; // skipping '(' @@ -173,6 +182,11 @@ public class DataTypeClassNameParser return className; } + private static boolean isVectorType(String className) + { + return className.startsWith(VECTOR_TYPE); + } + private static boolean isUserType(String className) { return className.startsWith(UDT_TYPE); diff --git a/src/java/org/apache/cassandra/cql3/functions/types/GettableByIndexData.java b/src/java/org/apache/cassandra/cql3/functions/types/GettableByIndexData.java index bcc5585df4..abbe32b29b 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/GettableByIndexData.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/GettableByIndexData.java @@ -465,6 +465,50 @@ public interface GettableByIndexData */ public Map getMap(int i, TypeToken keysType, TypeToken valuesType); + /** + * Returns the {@code i}th value as a vector, represented as a Java list. + * + *

This method uses the {@link CodecRegistry} to find a codec to convert the underlying CQL + * type to a vector of the specified type. + * + *

If the type of the elements is generic, use {@link #getVector(int, TypeToken)}. + * + *

Implementation note: the actual {@link List} implementation representing the vector will depend on the {@link + * TypeCodec codec} being used; therefore, callers should make no assumptions concerning its mutability nor its + * thread-safety. + * + * @param i the index ({@code 0 <= i < size()}) to retrieve. + * @param elementsClass the class for the elements of the list to retrieve. + * @return the value of the {@code i}th element as a list of {@code T} objects. + * @throws IndexOutOfBoundsException if {@code i} is not a valid index for this object. + * @throws CodecNotFoundException if there is no registered codec to convert the element's CQL type to a vector. + */ + public List getVector(int i, Class elementsClass); + + /** + * Returns the {@code i}th value as a vector, represented as a Java list. + * + *

This method uses the {@link CodecRegistry} to find a codec to convert the underlying CQL + * type to a vector of the specified type. + * + *

Use this variant with nested types, which produce a generic element type: + * + * 

+     * {@code List> l = row.getVector(1, new TypeToken>() {});}
+     *
+ * + *

Implementation note: the actual {@link List} implementation representing the vector will depend on the {@link + * TypeCodec codec} being used; therefore, callers should make no assumptions concerning its mutability nor its + * thread-safety. + * + * @param i the index ({@code 0 <= i < size()}) to retrieve. + * @param elementsType the type of the elements of the vector to retrieve. + * @return the value of the {@code i}th element as a vector of {@code T} objects. + * @throws IndexOutOfBoundsException if {@code i} is not a valid index for this object. + * @throws CodecNotFoundException if there is no registered codec to convert the element's CQL type to a vector. + */ + public List getVector(int i, TypeToken elementsType); + /** * Return the {@code i}th value as a UDT value. * diff --git a/src/java/org/apache/cassandra/cql3/functions/types/GettableByNameData.java b/src/java/org/apache/cassandra/cql3/functions/types/GettableByNameData.java index c578ff4cff..d65bf36d19 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/GettableByNameData.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/GettableByNameData.java @@ -464,6 +464,50 @@ public interface GettableByNameData */ public Map getMap(String name, TypeToken keysType, TypeToken valuesType); + /** + * Returns the value for {@code name} as a vector, represented as a Java list. + * + *

This method uses the {@link CodecRegistry} to find a codec to convert the underlying CQL + * type to a vector of the specified type. + * + *

If the type of the elements is generic, use {@link #getVector(String, TypeToken)}. + * + *

Implementation note: the actual {@link List} implementation representing the vector will depend on the {@link + * TypeCodec codec} being used; therefore, callers should make no assumptions concerning its mutability nor its + * thread-safety. + * + * @param name the name to retrieve. + * @param elementsClass the class for the elements of the vector to retrieve. + * @return the value of the {@code i}th element as a vector of {@code T} objects. + * @throws IllegalArgumentException if {@code name} is not valid name for this object. + * @throws CodecNotFoundException if there is no registered codec to convert the underlying CQL type to a vector. + */ + public List getVector(String name, Class elementsClass); + + /** + * Returns the value for {@code name} as a vector, represented as a Java list. + * + *

This method uses the {@link CodecRegistry} to find a codec to convert the underlying CQL + * type to a vector of the specified type. + * + *

Use this variant with nested types, which produce a generic element type: + * + * 

+     * {@code List> l = row.getVector("theColumn", new TypeToken>() {});}
+     *
+ * + *

Implementation note: the actual {@link List} implementation representing the vector will depend on the {@link + * TypeCodec codec} being used; therefore, callers should make no assumptions concerning its mutability nor its + * thread-safety. + * + * @param name the name to retrieve. + * @param elementsType the type for the elements of the vector to retrieve. + * @return the value of the {@code i}th element as a vector of {@code T} objects. + * @throws IllegalArgumentException if {@code name} is not valid name for this object. + * @throws CodecNotFoundException if there is no registered codec to convert the underlying CQL type to a vector. + */ + public List getVector(String name, TypeToken elementsType); + /** * Return the value for {@code name} as a UDT value. * diff --git a/src/java/org/apache/cassandra/cql3/functions/types/SettableByIndexData.java b/src/java/org/apache/cassandra/cql3/functions/types/SettableByIndexData.java index a9d08981e7..4e59a27550 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/SettableByIndexData.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/SettableByIndexData.java @@ -485,6 +485,20 @@ public interface SettableByIndexData> */ public T setSet(int i, Set v, TypeToken elementsType); + /** + * Sets the {@code i}th value to the provided vector value. + * + *

This method uses the {@link CodecRegistry} to find a codec to handle the conversion of {@code v} to the + * underlying CQL type. + * + * @param i the index of the value to set. + * @param v the value to set. + * @return this object. + * @throws IndexOutOfBoundsException if {@code i} is not a valid index for this object. + * @throws CodecNotFoundException if there is no registered codec to convert the value to the underlying CQL type. + */ + public T setVector(int i, List v); + /** * Sets the {@code i}th value to the provided UDT value. * diff --git a/src/java/org/apache/cassandra/cql3/functions/types/SettableByNameData.java b/src/java/org/apache/cassandra/cql3/functions/types/SettableByNameData.java index 514ba60154..376cf471cd 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/SettableByNameData.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/SettableByNameData.java @@ -513,6 +513,20 @@ public interface SettableByNameData> */ public T setSet(String name, Set v, TypeToken elementsType); + /** + * Sets the value for (all occurrences of) variable {@code name} to the provided vector value. + * + *

This method uses the {@link CodecRegistry} to find a codec to handle the conversion of {@code v} to the + * underlying CQL type. + * + * @param name the name of the value to set; if {@code name} is present multiple times, all its values are set. + * @param v the value to set. + * @return this object. + * @throws IllegalArgumentException if {@code name} is not a valid name for this object. + * @throws CodecNotFoundException if there is no registered codec to convert the value to the underlying CQL type. + */ + public T setVector(String name, List v); + /** * Sets the value for (all occurrences of) variable {@code name} to the provided UDT value. * diff --git a/src/java/org/apache/cassandra/cql3/functions/types/TypeCodec.java b/src/java/org/apache/cassandra/cql3/functions/types/TypeCodec.java index 54ff540789..1212e8e2ab 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/TypeCodec.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/TypeCodec.java @@ -122,6 +122,7 @@ import static org.apache.cassandra.cql3.functions.types.DataType.*; */ public abstract class TypeCodec { + private final static int VARIABLE_LENGTH = -1; /** * Return the default codec for the CQL type {@code boolean}. The returned codec maps the CQL type @@ -375,6 +376,19 @@ public abstract class TypeCodec return new MapCodec<>(keyCodec, valueCodec); } + /** + * Return a newly-created codec for the given CQL vector type. The returned codec maps the vector + * type into the Java type {@link List}. This method does not cache returned instances and + * returns a newly-allocated object at each invocation. + * + * @param type the vector type this codec should handle. + * @return A newly-created codec for the given CQL tuple type. + */ + public static TypeCodec> vector(VectorType type, TypeCodec valueCodec) + { + return VectorCodec.of(type, valueCodec); + } + /** * Return a newly-created codec for the given user-defined CQL type. The returned codec maps the * user-defined type into the Java type {@link UDTValue}. This method does not cache returned @@ -488,6 +502,24 @@ public abstract class TypeCodec return cqlType; } + /** + * @return the length of values for this type if all values are of fixed length, {@link #VARIABLE_LENGTH} otherwise + */ + public int serializedSize() + { + return VARIABLE_LENGTH; + } + + /** + * Checks if all values are of fixed length. + * + * @return {@code true} if all values are of fixed length, {@code false} otherwise. + */ + public final boolean isSerializedSizeFixed() + { + return serializedSize() != VARIABLE_LENGTH; + } + /** * Serialize the given value according to the CQL type handled by this codec. * @@ -1020,6 +1052,12 @@ public abstract class TypeCodec super(cqlType); } + @Override + public int serializedSize() + { + return 8; + } + @Override public Long parse(String value) { @@ -1190,6 +1228,12 @@ public abstract class TypeCodec super(DataType.cboolean()); } + @Override + public int serializedSize() + { + return 1; + } + @Override public Boolean parse(String value) { @@ -1308,6 +1352,12 @@ public abstract class TypeCodec super(DataType.cdouble()); } + @Override + public int serializedSize() + { + return 8; + } + @Override public Double parse(String value) { @@ -1364,6 +1414,12 @@ public abstract class TypeCodec super(DataType.cfloat()); } + @Override + public int serializedSize() + { + return 4; + } + @Override public Float parse(String value) { @@ -1593,6 +1649,12 @@ public abstract class TypeCodec super(DataType.cint()); } + @Override + public int serializedSize() + { + return 4; + } + @Override public Integer parse(String value) { @@ -1649,6 +1711,12 @@ public abstract class TypeCodec super(DataType.timestamp(), Date.class); } + @Override + public int serializedSize() + { + return 8; + } + @Override public Date parse(String value) { @@ -1719,6 +1787,12 @@ public abstract class TypeCodec super(DataType.date(), LocalDate.class); } + @Override + public int serializedSize() + { + return 8; + } + @Override public LocalDate parse(String value) { @@ -1856,6 +1930,12 @@ public abstract class TypeCodec super(cqlType, UUID.class); } + @Override + public int serializedSize() + { + return 16; + } + @Override public UUID parse(String value) { diff --git a/src/java/org/apache/cassandra/cql3/functions/types/TypeTokens.java b/src/java/org/apache/cassandra/cql3/functions/types/TypeTokens.java index aba74668a0..024c55e6c6 100644 --- a/src/java/org/apache/cassandra/cql3/functions/types/TypeTokens.java +++ b/src/java/org/apache/cassandra/cql3/functions/types/TypeTokens.java @@ -154,4 +154,23 @@ public final class TypeTokens }, valueType); // @formatter:on } + + /** + * Create a {@link TypeToken} that represents a {@link TypeToken} whose elements are of the given type. + * + * @param eltType The vector element type. + * @param The vector element type. + * @return A {@link TypeToken} that represents a {@link TypeToken} whose elements are of the given + * type. + */ + public static TypeToken> vectorOf(TypeToken eltType) + { + // @formatter:off + return new TypeToken>() + { + }.where(new TypeParameter() + { + }, eltType); + // @formatter:on + } } diff --git a/src/java/org/apache/cassandra/cql3/functions/types/VectorCodec.java b/src/java/org/apache/cassandra/cql3/functions/types/VectorCodec.java new file mode 100644 index 0000000000..c54b8da7b8 --- /dev/null +++ b/src/java/org/apache/cassandra/cql3/functions/types/VectorCodec.java @@ -0,0 +1,204 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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.cql3.functions.types; + +import java.nio.ByteBuffer; +import java.util.Iterator; +import java.util.List; + +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; + +import org.apache.cassandra.cql3.functions.types.exceptions.InvalidTypeException; +import org.apache.cassandra.transport.ProtocolVersion; +import org.apache.cassandra.utils.vint.VIntCoding; + +/** + * {@link TypeCodec} for the CQL type {@code vector}. Vectors are represented as {@link List}s for convenience, since + * it's probably easier for UDFs trying to return a newly created vector to create it as a standard Java list, rather + * than using a custom, not-standard vector class. + * + * @param The type of the vector elements. + */ +public abstract class VectorCodec extends TypeCodec> +{ + protected final VectorType type; + protected final TypeCodec subtypeCodec; + + private VectorCodec(VectorType type, TypeCodec subtypeCodec) + { + super(type, TypeTokens.vectorOf(subtypeCodec.getJavaType())); + this.type = type; + this.subtypeCodec = subtypeCodec; + } + + public static VectorCodec of(VectorType type, TypeCodec subtypeCodec) + { + return subtypeCodec.isSerializedSizeFixed() + ? new FixedLength<>(type, subtypeCodec) + : new VariableLength<>(type, subtypeCodec); + } + + @Override + public List parse(String value) throws InvalidTypeException + { + if (value == null || value.isEmpty() || value.equalsIgnoreCase("NULL")) + return null; + + ImmutableList.Builder values = ImmutableList.builder(); + for (String element : Splitter.on(", ").split(value.substring(1, value.length() - 1))) + { + values.add(subtypeCodec.parse(element)); + } + + return values.build(); + } + + @Override + public String format(List value) throws InvalidTypeException + { + return value == null ? "NULL" : Iterables.toString(value); + } + + /** + * {@link VectorCodec} for vectors of elements using a fixed-length encoding. + */ + private static class FixedLength extends VectorCodec + { + private final int valueLength; + + public FixedLength(VectorType type, TypeCodec subtypeCodec) + { + super(type, subtypeCodec); + valueLength = subtypeCodec.serializedSize() * type.getDimensions(); + } + + @Override + public int serializedSize() + { + return valueLength; + } + + @Override + public ByteBuffer serialize(List value, ProtocolVersion protocolVersion) throws InvalidTypeException + { + if (value == null || type.getDimensions() <= 0) + return null; + + Iterator values = value.iterator(); + ByteBuffer rv = ByteBuffer.allocate(valueLength); + for (int i = 0; i < type.getDimensions(); ++i) + { + ByteBuffer valueBuff = subtypeCodec.serialize(values.next(), protocolVersion); + valueBuff.rewind(); + rv.put(valueBuff); + } + rv.flip(); + return rv; + } + + @Override + public List deserialize(ByteBuffer bytes, ProtocolVersion protocolVersion) throws InvalidTypeException + { + if (bytes == null || bytes.remaining() == 0) + return null; + + // Determine element size by dividing count of remaining bytes by number of elements. + // This should have a remainder of zero since all elements are of the same fixed size. + int elementSize = Math.floorDiv(bytes.remaining(), type.getDimensions()); + assert bytes.remaining() % type.getDimensions() == 0 + : String.format("Expected elements of uniform size, observed %d elements with total bytes %d", + type.getDimensions(), bytes.remaining()); + + ImmutableList.Builder values = ImmutableList.builder(); + for (int i = 0; i < type.getDimensions(); ++i) + { + ByteBuffer slice = bytes.slice(); + slice.limit(elementSize); + values.add(subtypeCodec.deserialize(slice, protocolVersion)); + bytes.position(bytes.position() + elementSize); + } + + // Restore the input ByteBuffer to its original state + bytes.rewind(); + + return values.build(); + } + } + + /** + * {@link VectorCodec} for vectors of elements using a varaible-length encoding. + */ + private static class VariableLength extends VectorCodec + { + public VariableLength(VectorType type, TypeCodec subtypeCodec) + { + super(type, subtypeCodec); + } + + @Override + public ByteBuffer serialize(List values, ProtocolVersion version) throws InvalidTypeException + { + if (values == null) + return null; + + assert values.size() == type.getDimensions(); + + int i = 0; + int outputSize = 0; + ByteBuffer[] buffers = new ByteBuffer[values.size()]; + for (E value : values) + { + ByteBuffer bb = subtypeCodec.serialize(value, version); + buffers[i++] = bb; + int elemSize = bb.remaining(); + outputSize += elemSize + VIntCoding.computeUnsignedVIntSize(elemSize); + } + + ByteBuffer output = ByteBuffer.allocate(outputSize); + for (ByteBuffer bb : buffers) + { + VIntCoding.writeUnsignedVInt32(bb.remaining(), output); + output.put(bb.duplicate()); + } + return (ByteBuffer) output.flip(); + } + + @Override + public List deserialize(ByteBuffer bytes, ProtocolVersion version) throws InvalidTypeException + { + if (bytes == null || bytes.remaining() == 0) + return null; + + ByteBuffer input = bytes.duplicate(); + ImmutableList.Builder values = ImmutableList.builder(); + + for (int i = 0; i < type.getDimensions(); i++) + { + int size = VIntCoding.getUnsignedVInt32(input, input.position()); + input.position(input.position() + VIntCoding.computeUnsignedVIntSize(size)); + + ByteBuffer value = size < 0 ? null : CodecUtils.readBytes(input, size); + values.add(subtypeCodec.deserialize(value, version)); + } + return values.build(); + } + } +} diff --git a/src/java/org/apache/cassandra/cql3/functions/types/VectorType.java b/src/java/org/apache/cassandra/cql3/functions/types/VectorType.java new file mode 100644 index 0000000000..8d22322424 --- /dev/null +++ b/src/java/org/apache/cassandra/cql3/functions/types/VectorType.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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.cql3.functions.types; + +import java.util.Arrays; +import java.util.List; + +/** + * A vector type. + */ +public class VectorType extends DataType +{ + private final DataType subtype; + private final int dimensions; + + VectorType(DataType subtype, int dimensions) + { + super(Name.VECTOR); + assert dimensions > 0 : "vectors may only have positive dimentions; given " + dimensions; + this.subtype = subtype; + this.dimensions = dimensions; + } + + public DataType getSubtype() + { + return this.subtype; + } + + public int getDimensions() + { + return this.dimensions; + } + + @Override + public List getTypeArguments() + { + return Arrays.asList(subtype, DataType.cint()); + } + + @Override + public boolean isFrozen() + { + return true; + } + + @Override + public boolean equals(Object o) + { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + + VectorType that = (VectorType) o; + + if (dimensions != that.dimensions) return false; + if (!subtype.equals(that.subtype)) return false; + return name == that.name; + } + + @Override + public int hashCode() + { + int result = name.hashCode(); + result = 31 * result + dimensions; + result = 31 * result + subtype.hashCode(); + return result; + } + + @Override + public String toString() + { + return String.format("vector<%s, %d>", subtype, dimensions); + } +} diff --git a/src/java/org/apache/cassandra/cql3/selection/Selectable.java b/src/java/org/apache/cassandra/cql3/selection/Selectable.java index dbde53fc4d..5c06da64e8 100644 --- a/src/java/org/apache/cassandra/cql3/selection/Selectable.java +++ b/src/java/org/apache/cassandra/cql3/selection/Selectable.java @@ -910,7 +910,8 @@ public interface Selectable extends AssignmentTestable } VectorType vectorType = (VectorType) type; - assert vectorType.dimension == selectables.size() : String.format("Unable to create a vector selector of type %s from %d elements", vectorType.asCQL3Type(), selectables.size()); + if (vectorType.dimension != selectables.size()) + throw invalidRequest("Unable to create a vector selector of type %s from %d elements", vectorType.asCQL3Type(), selectables.size()); List> expectedTypes = new ArrayList<>(selectables.size()); for (int i = 0, m = selectables.size(); i < m; i++) diff --git a/src/java/org/apache/cassandra/utils/JavaDriverUtils.java b/src/java/org/apache/cassandra/utils/JavaDriverUtils.java index d58a26c359..6c0567d03f 100644 --- a/src/java/org/apache/cassandra/utils/JavaDriverUtils.java +++ b/src/java/org/apache/cassandra/utils/JavaDriverUtils.java @@ -54,21 +54,9 @@ public final class JavaDriverUtils public static TypeCodec codecFor(AbstractType abstractType) { - if (containsVector(abstractType.unwrap())) - throw new IllegalArgumentException("Vectors are not supported on UDFs; given " + abstractType.asCQL3Type()); return codecFor(driverType(abstractType)); } - private static boolean containsVector(AbstractType type) - { - if (type.isVector()) return true; - for (AbstractType child : type.subTypes()) - { - if (containsVector(child)) return true; - } - return false; - } - public static TypeCodec codecFor(DataType dataType) { return codecRegistry.codecFor(dataType); diff --git a/test/unit/org/apache/cassandra/cql3/validation/entities/UFSecurityTest.java b/test/unit/org/apache/cassandra/cql3/validation/entities/UFSecurityTest.java index 05b353acf2..7716e8b4e5 100644 --- a/test/unit/org/apache/cassandra/cql3/validation/entities/UFSecurityTest.java +++ b/test/unit/org/apache/cassandra/cql3/validation/entities/UFSecurityTest.java @@ -118,7 +118,13 @@ public class UFSecurityTest extends CQLTester " java.lang.Runtime.getRuntime().exec(\"/tmp/foo\"); return 0d;" + "} catch (Exception t) {" + " throw new RuntimeException(t);" + - '}'} + '}'}, + {"org.apache.cassandra.utils.vint.VIntCoding", + "try {" + + " org.apache.cassandra.utils.vint.VIntCoding.computeUnsignedVIntSize(0L); return 0d;" + + "} catch (Exception t) {" + + " throw new RuntimeException(t);" + + '}'} }; for (String[] typeAndSource : typesAndSources) diff --git a/test/unit/org/apache/cassandra/cql3/validation/operations/CQLVectorTest.java b/test/unit/org/apache/cassandra/cql3/validation/operations/CQLVectorTest.java index ab39feb952..773a08cf38 100644 --- a/test/unit/org/apache/cassandra/cql3/validation/operations/CQLVectorTest.java +++ b/test/unit/org/apache/cassandra/cql3/validation/operations/CQLVectorTest.java @@ -36,6 +36,8 @@ import org.apache.cassandra.exceptions.InvalidRequestException; import org.apache.cassandra.transport.ProtocolVersion; import org.assertj.core.api.Assertions; +import static java.lang.String.format; + public class CQLVectorTest extends CQLTester.InMemory { @Test @@ -338,25 +340,158 @@ public class CQLVectorTest extends CQLTester.InMemory @Test public void udf() throws Throwable { - // For future authors, if this test starts to fail as vectors become supported in UDFs, please update this test - // to test the integration and remove the requirement that we reject UDFs all together createTable(KEYSPACE, "CREATE TABLE %s (pk int primary key, value vector)"); - Assertions.assertThatThrownBy(() -> createFunction(KEYSPACE, - "", - "CREATE FUNCTION %s (x vector) " + - "CALLED ON NULL INPUT " + - "RETURNS vector " + - "LANGUAGE java " + - "AS 'return x;'")) - .hasRootCauseMessage("Vectors are not supported on UDFs; given vector"); + Vector vector = vector(1, 2); + execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector); - Assertions.assertThatThrownBy(() -> createFunction(KEYSPACE, - "", - "CREATE FUNCTION %s (x list>) " + - "CALLED ON NULL INPUT " + - "RETURNS list> " + - "LANGUAGE java " + - "AS 'return x;'")) - .hasRootCauseMessage("Vectors are not supported on UDFs; given list>"); + // identitiy function + String f = createFunction(KEYSPACE, + "", + "CREATE FUNCTION %s (x vector) " + + "CALLED ON NULL INPUT " + + "RETURNS vector " + + "LANGUAGE java " + + "AS 'return x;'"); + assertRows(execute(format("SELECT %s(value) FROM %%s", f)), row(vector)); + assertRows(execute(format("SELECT %s([2, 3]) FROM %%s", f)), row(vector(2, 3))); + assertRows(execute(format("SELECT %s(null) FROM %%s", f)), row((Vector) null)); + + // identitiy function with nested type + f = createFunction(KEYSPACE, + "", + "CREATE FUNCTION %s (x list>) " + + "CALLED ON NULL INPUT " + + "RETURNS list> " + + "LANGUAGE java " + + "AS 'return x;'"); + assertRows(execute(format("SELECT %s([value]) FROM %%s", f)), row(list(vector))); + assertRows(execute(format("SELECT %s([[2, 3]]) FROM %%s", f)), row(list(vector(2, 3)))); + assertRows(execute(format("SELECT %s(null) FROM %%s", f)), row((Vector) null)); + + // identitiy function with elements of variable length + f = createFunction(KEYSPACE, + "", + "CREATE FUNCTION %s (x vector) " + + "CALLED ON NULL INPUT " + + "RETURNS vector " + + "LANGUAGE java " + + "AS 'return x;'"); + assertRows(execute(format("SELECT %s(['abc', 'defghij']) FROM %%s", f)), row(vector("abc", "defghij"))); + assertRows(execute(format("SELECT %s(null) FROM %%s", f)), row((Vector) null)); + + // function accessing vector argument elements + f = createFunction(KEYSPACE, + "", + "CREATE FUNCTION %s (x vector, i int) " + + "CALLED ON NULL INPUT " + + "RETURNS int " + + "LANGUAGE java " + + "AS 'return x == null ? null : x.get(i);'"); + assertRows(execute(format("SELECT %s(value, 0), %) " + + "CALLED ON NULL INPUT " + + "RETURNS int " + + "LANGUAGE java " + + "AS 'return x == null ? 0 : x.size();'"); + assertRows(execute(format("SELECT %s(value) FROM %%s", f)), row(2)); + assertRows(execute(format("SELECT %s([2, 3]) FROM %%s", f)), row(2)); + assertRows(execute(format("SELECT %s(null) FROM %%s", f)), row(0)); + + // build vector with elements of fixed length + f = createFunction(KEYSPACE, + "", + "CREATE FUNCTION %s () " + + "CALLED ON NULL INPUT " + + "RETURNS vector " + + "LANGUAGE java " + + "AS 'return Arrays.asList(1.3, 2.2, 3.1);'"); + assertRows(execute(format("SELECT %s() FROM %%s", f)), row(vector(1.3, 2.2, 3.1))); + + // build vector with elements of variable length + f = createFunction(KEYSPACE, + "", + "CREATE FUNCTION %s () " + + "CALLED ON NULL INPUT " + + "RETURNS vector " + + "LANGUAGE java " + + "AS 'return Arrays.asList(\"a\", \"bc\", \"def\");'"); + assertRows(execute(format("SELECT %s() FROM %%s", f)), row(vector("a", "bc", "def"))); + + // concat vectors, just to put it all together + f = createFunction(KEYSPACE, + "", + "CREATE FUNCTION %s (x vector, y vector) " + + "CALLED ON NULL INPUT " + + "RETURNS vector " + + "LANGUAGE java " + + "AS '" + + "if (x == null || y == null) return null;" + + "List l = new ArrayList(x); " + + "l.addAll(y); " + + "return l;'"); + assertRows(execute(format("SELECT %s(value, [3, 4]) FROM %%s", f)), row(vector(1, 2, 3, 4))); + assertRows(execute(format("SELECT %s([2, 3], value) FROM %%s", f)), row(vector(2, 3, 1, 2))); + assertRows(execute(format("SELECT %s(null, null) FROM %%s", f)), row((Vector) null)); + + // Test wrong arguments on function call + assertInvalidThrowMessage("cannot be passed as argument 0 of function " + f, + InvalidRequestException.class, + format("SELECT %s((int) 0, [3, 4]) FROM %%s", f)); + assertInvalidThrowMessage("cannot be passed as argument 1 of function " + f, + InvalidRequestException.class, + format("SELECT %s([1, 2], (int) 0) FROM %%s", f)); + assertInvalidThrowMessage("Invalid number of arguments in call to function " + f, + InvalidRequestException.class, + format("SELECT %s([1, 2]) FROM %%s", f)); + assertInvalidThrowMessage("Invalid number of arguments in call to function " + f, + InvalidRequestException.class, + format("SELECT %s([1, 2], [3, 4], [5, 6]) FROM %%s", f)); + assertInvalidThrowMessage("Unable to create a vector selector of type vector from 3 elements", + InvalidRequestException.class, + format("SELECT %s([1, 2, 3], [4, 5, 6]) FROM %%s", f)); + + // Test wrong types on function creation + assertInvalidThrowMessage("vectors may only have positive dimentions; given 0", + InvalidRequestException.class, + "CREATE FUNCTION %s (x vector) " + + "CALLED ON NULL INPUT " + + "RETURNS vector " + + "LANGUAGE java " + + "AS 'return x;'"); + assertInvalidThrowMessage("vectors may only have positive dimentions; given 0", + InvalidRequestException.class, + "CREATE FUNCTION %s (x vector) " + + "CALLED ON NULL INPUT " + + "RETURNS vector " + + "LANGUAGE java " + + "AS 'return x;'"); + + // function reading and writing a udt vector field + String udt = createType("CREATE TYPE %s (v vector)"); + alterTable("ALTER TABLE %s ADD udt " + udt); + execute("INSERT INTO %s (pk, udt) VALUES (0, ?)", userType("v", vector)); + f = createFunction(KEYSPACE, + "", + "CREATE FUNCTION %s (udt " + udt + ") " + + "CALLED ON NULL INPUT " + + "RETURNS " + udt + ' ' + + "LANGUAGE java " + + "AS '" + + "if (udt == null) return null;" + + "List v = new ArrayList(udt.getVector(\"v\", Integer.class));" + + "v.set(0, 7);" + + "return udt.setVector(\"v\", v);'"); + assertRows(execute(format("SELECT %s(udt) FROM %%s", f)), row(userType("v", vector(7, 2)))); + assertRows(execute(format("SELECT %s({v: [10, 20]}) FROM %%s", f)), row(userType("v", vector(7, 20)))); + assertRows(execute(format("SELECT %s(null) FROM %%s", f)), row((Object) null)); + + // make sure the function referencing the UDT is dropped before dropping the UDT at cleanup + execute("DROP FUNCTION " + f); } }