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

Add support for vectors in UDFs 

patch by Andrés de la Peña; reviewed by Berenguer Blasi and Maxwell Guo for CASSANDRA-18613
This commit is contained in:
Andrés de la Peña 2023-06-22 19:33:25 +01:00
parent bf11c9ae03
commit ddbc52990f
20 changed files with 806 additions and 56 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGES.txt
View File

@ -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)

6
src/java/org/apache/cassandra/cql3/functions/types/AbstractAddressableByIndexData.java
View File

@ -253,6 +253,12 @@ extends AbstractGettableByIndexData implements SettableByIndexData<T>
return setValue(i, codecFor(i, TypeTokens.setOf(elementsType)).serialize(v, protocolVersion));
}
@Override
public <E> T setVector(int i, List<E> v)
{
return setValue(i, codecFor(i).serialize(v, protocolVersion));
}
@Override
public T setUDTValue(int i, UDTValue v)
{

16
src/java/org/apache/cassandra/cql3/functions/types/AbstractData.java
View File

@ -540,6 +540,22 @@ implements SettableData<T>
return wrapped;
}
@Override
public <E> T setVector(int i, List<E> v)
{
return setValue(i, codecFor(i).serialize(v, protocolVersion));
}
@Override
public <E> T setVector(String name, List<E> v)
{
for (int i : getAllIndexesOf(name))
{
setVector(i, v);
}
return wrapped;
}
@Override
public T setUDTValue(int i, UDTValue v)
{

22
src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableByIndexData.java
View File

@ -365,6 +365,28 @@ abstract class AbstractGettableByIndexData implements GettableByIndexData
return codecFor(i, javaType).deserialize(value, protocolVersion);
}
/**
* {@inheritDoc}
*/
@Override
@SuppressWarnings("unchecked")
public <T> List<T> getVector(int i, Class<T> elementsClass)
{
return getList(i, TypeToken.of(elementsClass));
}
/**
* {@inheritDoc}
*/
@Override
@SuppressWarnings("unchecked")
public <T> List<T> getVector(int i, TypeToken<T> elementsType)
{
ByteBuffer value = getValue(i);
TypeToken<List<T>> javaType = TypeTokens.listOf(elementsType);
return codecFor(i, javaType).deserialize(value, protocolVersion);
}
/**
* {@inheritDoc}
*/

18
src/java/org/apache/cassandra/cql3/functions/types/AbstractGettableData.java
View File

@ -269,6 +269,24 @@ implements GettableData
return getMap(getIndexOf(name), keysType, valuesType);
}
/**
* {@inheritDoc}
*/
@Override
public <T> List<T> getVector(String name, Class<T> elementsClass)
{
return getList(getIndexOf(name), elementsClass);
}
/**
* {@inheritDoc}
*/
@Override
public <T> List<T> getVector(String name, TypeToken<T> elementsType)
{
return getList(getIndexOf(name), elementsType);
}
/**
* {@inheritDoc}
*/

69
src/java/org/apache/cassandra/cql3/functions/types/CodecRegistry.java
View File

@ -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<T>) TypeCodec.map(keyCodec, valueCodec);
}
if (cqlType instanceof VectorType
&& (javaType == null || List.class.isAssignableFrom(javaType.getRawType())))
{
VectorType type = (VectorType) cqlType;
return (TypeCodec<T>) 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<T>) TypeCodec.vector(type, findCodec(type.getSubtype(), null));
}
if ((cqlType == null || cqlType.getName() == TUPLE)
&& value instanceof TupleValue)
{
return (TypeCodec<T>)
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<T>)
TypeCodec.userType(cqlType == null ? ((UDTValue) value).getType() : (UserType) cqlType);

15
src/java/org/apache/cassandra/cql3/functions/types/DataType.java
View File

@ -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.
*

14
src/java/org/apache/cassandra/cql3/functions/types/DataTypeClassNameParser.java
View File

@ -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<String, DataType> cassTypeToDataType =
new ImmutableMap.Builder<String, DataType>()
@ -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<String> 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);

44
src/java/org/apache/cassandra/cql3/functions/types/GettableByIndexData.java
View File

@ -465,6 +465,50 @@ public interface GettableByIndexData
*/
public <K, V> Map<K, V> getMap(int i, TypeToken<K> keysType, TypeToken<V> valuesType);
/**
* Returns the {@code i}th value as a vector, represented as a Java list.
*
* <p>This method uses the {@link CodecRegistry} to find a codec to convert the underlying CQL
* type to a vector of the specified type.
*
* <p>If the type of the elements is generic, use {@link #getVector(int, TypeToken)}.
*
* <p>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 <T> List<T> getVector(int i, Class<T> elementsClass);
/**
* Returns the {@code i}th value as a vector, represented as a Java list.
*
* <p>This method uses the {@link CodecRegistry} to find a codec to convert the underlying CQL
* type to a vector of the specified type.
*
* <p>Use this variant with nested types, which produce a generic element type:
*
* <pre>
* {@code List<List<String>> l = row.getVector(1, new TypeToken<List<String>>() {});}
* </pre>
*
* <p>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 <T> List<T> getVector(int i, TypeToken<T> elementsType);
/**
* Return the {@code i}th value as a UDT value.
*

44
src/java/org/apache/cassandra/cql3/functions/types/GettableByNameData.java
View File

@ -464,6 +464,50 @@ public interface GettableByNameData
*/
public <K, V> Map<K, V> getMap(String name, TypeToken<K> keysType, TypeToken<V> valuesType);
/**
* Returns the value for {@code name} as a vector, represented as a Java list.
*
* <p>This method uses the {@link CodecRegistry} to find a codec to convert the underlying CQL
* type to a vector of the specified type.
*
* <p>If the type of the elements is generic, use {@link #getVector(String, TypeToken)}.
*
* <p>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 <T> List<T> getVector(String name, Class<T> elementsClass);
/**
* Returns the value for {@code name} as a vector, represented as a Java list.
*
* <p>This method uses the {@link CodecRegistry} to find a codec to convert the underlying CQL
* type to a vector of the specified type.
*
* <p>Use this variant with nested types, which produce a generic element type:
*
* <pre>
* {@code List<List<String>> l = row.getVector("theColumn", new TypeToken<List<String>>() {});}
* </pre>
*
* <p>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 <T> List<T> getVector(String name, TypeToken<T> elementsType);
/**
* Return the value for {@code name} as a UDT value.
*

14
src/java/org/apache/cassandra/cql3/functions/types/SettableByIndexData.java
View File

@ -485,6 +485,20 @@ public interface SettableByIndexData<T extends SettableByIndexData<T>>
*/
public <E> T setSet(int i, Set<E> v, TypeToken<E> elementsType);
/**
* Sets the {@code i}th value to the provided vector value.
*
* <p>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 <E> T setVector(int i, List<E> v);
/**
* Sets the {@code i}th value to the provided UDT value.
*

14
src/java/org/apache/cassandra/cql3/functions/types/SettableByNameData.java
View File

@ -513,6 +513,20 @@ public interface SettableByNameData<T extends SettableData<T>>
*/
public <E> T setSet(String name, Set<E> v, TypeToken<E> elementsType);
/**
* Sets the value for (all occurrences of) variable {@code name} to the provided vector value.
*
* <p>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 <E> T setVector(String name, List<E> v);
/**
* Sets the value for (all occurrences of) variable {@code name} to the provided UDT value.
*

80
src/java/org/apache/cassandra/cql3/functions/types/TypeCodec.java
View File

@ -122,6 +122,7 @@ import static org.apache.cassandra.cql3.functions.types.DataType.*;
*/
public abstract class TypeCodec<T>
{
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<T>
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 <E> TypeCodec<List<E>> vector(VectorType type, TypeCodec<E> 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<T>
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<T>
super(cqlType);
}
@Override
public int serializedSize()
{
return 8;
}
@Override
public Long parse(String value)
{
@ -1190,6 +1228,12 @@ public abstract class TypeCodec<T>
super(DataType.cboolean());
}
@Override
public int serializedSize()
{
return 1;
}
@Override
public Boolean parse(String value)
{
@ -1308,6 +1352,12 @@ public abstract class TypeCodec<T>
super(DataType.cdouble());
}
@Override
public int serializedSize()
{
return 8;
}
@Override
public Double parse(String value)
{
@ -1364,6 +1414,12 @@ public abstract class TypeCodec<T>
super(DataType.cfloat());
}
@Override
public int serializedSize()
{
return 4;
}
@Override
public Float parse(String value)
{
@ -1593,6 +1649,12 @@ public abstract class TypeCodec<T>
super(DataType.cint());
}
@Override
public int serializedSize()
{
return 4;
}
@Override
public Integer parse(String value)
{
@ -1649,6 +1711,12 @@ public abstract class TypeCodec<T>
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<T>
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<T>
super(cqlType, UUID.class);
}
@Override
public int serializedSize()
{
return 16;
}
@Override
public UUID parse(String value)
{

19
src/java/org/apache/cassandra/cql3/functions/types/TypeTokens.java
View File

@ -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 <T> The vector element type.
* @return A {@link TypeToken} that represents a {@link TypeToken} whose elements are of the given
* type.
*/
public static <T> TypeToken<List<T>> vectorOf(TypeToken<T> eltType)
{
// @formatter:off
return new TypeToken<List<T>>()
{
}.where(new TypeParameter<T>()
{
}, eltType);
// @formatter:on
}
}

204
src/java/org/apache/cassandra/cql3/functions/types/VectorCodec.java Normal file
View File

@ -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 <E> The type of the vector elements.
*/
public abstract class VectorCodec<E> extends TypeCodec<List<E>>
{
protected final VectorType type;
protected final TypeCodec<E> subtypeCodec;
private VectorCodec(VectorType type, TypeCodec<E> subtypeCodec)
{
super(type, TypeTokens.vectorOf(subtypeCodec.getJavaType()));
this.type = type;
this.subtypeCodec = subtypeCodec;
}
public static <E> VectorCodec<E> of(VectorType type, TypeCodec<E> subtypeCodec)
{
return subtypeCodec.isSerializedSizeFixed()
? new FixedLength<>(type, subtypeCodec)
: new VariableLength<>(type, subtypeCodec);
}
@Override
public List<E> parse(String value) throws InvalidTypeException
{
if (value == null || value.isEmpty() || value.equalsIgnoreCase("NULL"))
return null;
ImmutableList.Builder<E> 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<E> 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<E> extends VectorCodec<E>
{
private final int valueLength;
public FixedLength(VectorType type, TypeCodec<E> subtypeCodec)
{
super(type, subtypeCodec);
valueLength = subtypeCodec.serializedSize() * type.getDimensions();
}
@Override
public int serializedSize()
{
return valueLength;
}
@Override
public ByteBuffer serialize(List<E> value, ProtocolVersion protocolVersion) throws InvalidTypeException
{
if (value == null || type.getDimensions() <= 0)
return null;
Iterator<E> 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<E> 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<E> 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<E> extends VectorCodec<E>
{
public VariableLength(VectorType type, TypeCodec<E> subtypeCodec)
{
super(type, subtypeCodec);
}
@Override
public ByteBuffer serialize(List<E> 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<E> deserialize(ByteBuffer bytes, ProtocolVersion version) throws InvalidTypeException
{
if (bytes == null || bytes.remaining() == 0)
return null;
ByteBuffer input = bytes.duplicate();
ImmutableList.Builder<E> 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();
}
}
}

88
src/java/org/apache/cassandra/cql3/functions/types/VectorType.java Normal file
View File

@ -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<DataType> 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);
}
}

3
src/java/org/apache/cassandra/cql3/selection/Selectable.java
View File

@ -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<AbstractType<?>> expectedTypes = new ArrayList<>(selectables.size());
for (int i = 0, m = selectables.size(); i < m; i++)

12
src/java/org/apache/cassandra/utils/JavaDriverUtils.java
View File

@ -54,21 +54,9 @@ public final class JavaDriverUtils
public static TypeCodec<Object> 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<Object> codecFor(DataType dataType)
{
return codecRegistry.codecFor(dataType);

8
test/unit/org/apache/cassandra/cql3/validation/entities/UFSecurityTest.java
View File

@ -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)

171
test/unit/org/apache/cassandra/cql3/validation/operations/CQLVectorTest.java
View File

@ -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<int, 2>)");
Assertions.assertThatThrownBy(() -> createFunction(KEYSPACE,
"",
"CREATE FUNCTION %s (x vector<int, 2>) " +
"CALLED ON NULL INPUT " +
"RETURNS vector<int, 2> " +
"LANGUAGE java " +
"AS 'return x;'"))
.hasRootCauseMessage("Vectors are not supported on UDFs; given vector<int, 2>");
Vector<Integer> vector = vector(1, 2);
execute("INSERT INTO %s (pk, value) VALUES (0, ?)", vector);
Assertions.assertThatThrownBy(() -> createFunction(KEYSPACE,
"",
"CREATE FUNCTION %s (x list<vector<int, 2>>) " +
"CALLED ON NULL INPUT " +
"RETURNS list<vector<int, 2>> " +
"LANGUAGE java " +
"AS 'return x;'"))
.hasRootCauseMessage("Vectors are not supported on UDFs; given list<vector<int, 2>>");
// identitiy function
String f = createFunction(KEYSPACE,
"",
"CREATE FUNCTION %s (x vector<int, 2>) " +
"CALLED ON NULL INPUT " +
"RETURNS vector<int, 2> " +
"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<Integer>) null));
// identitiy function with nested type
f = createFunction(KEYSPACE,
"",
"CREATE FUNCTION %s (x list<vector<int, 2>>) " +
"CALLED ON NULL INPUT " +
"RETURNS list<vector<int, 2>> " +
"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<Integer>) null));
// identitiy function with elements of variable length
f = createFunction(KEYSPACE,
"",
"CREATE FUNCTION %s (x vector<text, 2>) " +
"CALLED ON NULL INPUT " +
"RETURNS vector<text, 2> " +
"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<Integer>) null));
// function accessing vector argument elements
f = createFunction(KEYSPACE,
"",
"CREATE FUNCTION %s (x vector<int, 2>, 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), %<s(value, 1) FROM %%s", f)), row(1, 2));
assertRows(execute(format("SELECT %s([2, 3], 0), %<s([2, 3], 1) FROM %%s", f)), row(2, 3));
assertRows(execute(format("SELECT %s(null, 0) FROM %%s", f)), row((Integer) null));
// function accessing vector argument dimensions
f = createFunction(KEYSPACE,
"",
"CREATE FUNCTION %s (x vector<int, 2>) " +
"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<double, 3> " +
"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<text, 3> " +
"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<int, 2>, y vector<int, 2>) " +
"CALLED ON NULL INPUT " +
"RETURNS vector<int, 4> " +
"LANGUAGE java " +
"AS '" +
"if (x == null || y == null) return null;" +
"List<Integer> l = new ArrayList<Integer>(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<Integer>) 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<int, 2> 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<int, 0>) " +
"CALLED ON NULL INPUT " +
"RETURNS vector<int, 2> " +
"LANGUAGE java " +
"AS 'return x;'");
assertInvalidThrowMessage("vectors may only have positive dimentions; given 0",
InvalidRequestException.class,
"CREATE FUNCTION %s (x vector<int, 2>) " +
"CALLED ON NULL INPUT " +
"RETURNS vector<int, 0> " +
"LANGUAGE java " +
"AS 'return x;'");
// function reading and writing a udt vector field
String udt = createType("CREATE TYPE %s (v vector<int,2>)");
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<Integer> v = new ArrayList<Integer>(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);
}
}