mirror of https://github.com/apache/cassandra
Add vector similarity functions
patch by Andrés de la Peña; reviewed by Ekaterina Dimitrova and Maxwell Guo for CASSANDRA-18640 Co-authored-by: Jonathan Ellis <jbellis@datastax.com> Co-authored-by: Andrés de la Peña <a.penya.garcia@gmail.com>
This commit is contained in:
Andrés de la Peña
parent
4bbfd64fcd
commit
5a82c04fd3
|
|
@ -1,4 +1,5 @@
|
|||
5.0
|
||||
* Add vector similarity functions (CASSANDRA-18640)
|
||||
* Lift MessagingService.minimum_version to 40 (CASSANDRA-18314)
|
||||
* Introduce pluggable crypto providers and default to AmazonCorrettoCryptoProvider (CASSANDRA-18624)
|
||||
* Improved DeletionTime serialization (CASSANDRA-18648)
|
||||
|
|
|
|||
1
NEWS.txt
1
NEWS.txt
|
|
@ -84,6 +84,7 @@ New features
|
|||
class. Further details and documentation can be found in
|
||||
src/java/org/apache/cassandra/db/compaction/UnifiedCompactionStrategy.md
|
||||
- New `VectorType` (cql `vector<element_type, dimension>`) which adds new fixed-length element arrays. See CASSANDRA-18504
|
||||
- Added new vector similarity functions `similarity_cosine`, `similarity_euclidean` and `similarity_dot_product`.
|
||||
- Removed UDT type migration logic for 3.6+ clusters upgrading to 4.0. If migration has been disabled, it must be
|
||||
enabled before upgrading to 5.0 if the cluster used UDTs. See CASSANDRA-18504
|
||||
- Entended max expiration time from 2038-01-19T03:14:06+00:00 to 2106-02-07T06:28:13+00:00
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ The following describes the changes in each version of CQL.
|
|||
|
||||
== 3.4.7
|
||||
|
||||
* Add vector similarity functions (`18640`)
|
||||
* Remove deprecated functions `dateOf` and `unixTimestampOf`, replaced by `toTimestamp` and `toUnixTimestamp` (`18328`)
|
||||
* Added support for attaching masking functions to table columns (`18068`)
|
||||
* Add UNMASK permission (`18069`)
|
||||
|
|
|
|||
|
|
@ -284,6 +284,13 @@ A number of functions allow to obscure the real contents of a column containing
|
|||
|
||||
include::partial$masking_functions.adoc[]
|
||||
|
||||
[[vector-similarity-functions]]
|
||||
===== Vector similarity functions
|
||||
|
||||
A number of functions allow to obtain the similarity score between vectors of floats.
|
||||
|
||||
include::partial$vector_functions.adoc[]
|
||||
|
||||
[[user-defined-scalar-functions]]
|
||||
==== User-defined functions
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,41 @@
|
|||
[cols=",",options="header",]
|
||||
|===
|
||||
|Function | Description
|
||||
|
||||
| `similarity_cosine(vector, vector)` | Calculates the cosine similarity score between two float vectors of the same dimension.
|
||||
|
||||
Examples:
|
||||
|
||||
`similarity_cosine([0.1, 0.2], null)` -> `null`
|
||||
|
||||
`similarity_cosine([0.1, 0.2], [0.1, 0.2])` -> `1`
|
||||
|
||||
`similarity_cosine([0.1, 0.2], [-0.1, -0.2])` -> `0`
|
||||
|
||||
`similarity_cosine([0.1, 0.2], [0.9, 0.8])` -> `0.964238`
|
||||
|
||||
| `similarity_euclidean(vector, vector)` | Calculates the euclidian distance between two float vectors of the same dimension.
|
||||
|
||||
Examples:
|
||||
|
||||
`similarity_euclidean([0.1, 0.2], null)` -> `null`
|
||||
|
||||
`similarity_euclidean([0.1, 0.2], [0.1, 0.2])` -> `1`
|
||||
|
||||
`similarity_euclidean([0.1, 0.2], [-0.1, -0.2])` -> `0.833333`
|
||||
|
||||
`similarity_euclidean([0.1, 0.2], [0.9, 0.8])` -> `0.5`
|
||||
|
||||
| `similarity_dot_product(vector, vector)` | Calculates the dot product between two float vectors of the same dimension.
|
||||
|
||||
Examples:
|
||||
|
||||
`similarity_dot_product([0.1, 0.2], null)` -> `null`
|
||||
|
||||
`similarity_dot_product([0.1, 0.2], [0.1, 0.2])` -> `0.525`
|
||||
|
||||
`similarity_dot_product([0.1, 0.2], [-0.1, -0.2])` -> `0.475`
|
||||
|
||||
`similarity_dot_product([0.1, 0.2], [0.9, 0.8])` -> `0.625`
|
||||
|
||||
|===
|
||||
|
|
@ -834,13 +834,13 @@ public interface CQL3Type
|
|||
private static class RawVector extends Raw
|
||||
{
|
||||
private final CQL3Type.Raw element;
|
||||
private final int dimention;
|
||||
private final int dimension;
|
||||
|
||||
private RawVector(Raw element, int dimention)
|
||||
private RawVector(Raw element, int dimension)
|
||||
{
|
||||
super(true);
|
||||
this.element = element;
|
||||
this.dimention = dimention;
|
||||
this.dimension = dimension;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -865,7 +865,13 @@ public interface CQL3Type
|
|||
public CQL3Type prepare(String keyspace, Types udts) throws InvalidRequestException
|
||||
{
|
||||
CQL3Type type = element.prepare(keyspace, udts);
|
||||
return new Vector(type.getType(), dimention);
|
||||
return new Vector(type.getType(), dimension);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString()
|
||||
{
|
||||
return "vector<" + element.toString() + ", " + dimension + '>';
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -90,9 +90,19 @@ public abstract class FunctionFactory
|
|||
if (numArgs < numMandatoryParameters || numArgs > numParameters)
|
||||
throw invalidNumberOfArgumentsException();
|
||||
|
||||
// try to infer the types of the arguments
|
||||
// Do a first pass trying to infer the types of the arguments individually, without any context about the types
|
||||
// of the other arguments. We don't do any validation during this first pass.
|
||||
List<AbstractType<?>> types = new ArrayList<>(args.size());
|
||||
for (int i = 0; i < args.size(); i++)
|
||||
{
|
||||
AssignmentTestable arg = args.get(i);
|
||||
FunctionParameter parameter = parameters.get(i);
|
||||
types.add(parameter.inferType(SchemaConstants.SYSTEM_KEYSPACE_NAME, arg, receiverType, null));
|
||||
}
|
||||
|
||||
// Do a second pass trying to infer the types of the arguments considering the types of other inferred types.
|
||||
// We can validate the inferred types during this second pass.
|
||||
for (int i = 0; i < args.size(); i++)
|
||||
{
|
||||
AssignmentTestable arg = args.get(i);
|
||||
FunctionParameter parameter = parameters.get(i);
|
||||
|
|
@ -103,7 +113,7 @@ public abstract class FunctionFactory
|
|||
arg, this));
|
||||
parameter.validateType(name, arg, type);
|
||||
type = type.udfType();
|
||||
types.add(type);
|
||||
types.set(i, type);
|
||||
}
|
||||
|
||||
return doGetOrCreateFunction(types, receiverType);
|
||||
|
|
|
|||
|
|
@ -25,12 +25,14 @@ import javax.annotation.Nullable;
|
|||
|
||||
import org.apache.cassandra.cql3.AssignmentTestable;
|
||||
import org.apache.cassandra.cql3.CQL3Type;
|
||||
import org.apache.cassandra.cql3.selection.Selectable;
|
||||
import org.apache.cassandra.db.marshal.AbstractType;
|
||||
import org.apache.cassandra.db.marshal.CollectionType;
|
||||
import org.apache.cassandra.db.marshal.ListType;
|
||||
import org.apache.cassandra.db.marshal.MapType;
|
||||
import org.apache.cassandra.db.marshal.NumberType;
|
||||
import org.apache.cassandra.db.marshal.SetType;
|
||||
import org.apache.cassandra.db.marshal.VectorType;
|
||||
import org.apache.cassandra.exceptions.InvalidRequestException;
|
||||
|
||||
import static java.lang.String.format;
|
||||
|
|
@ -48,13 +50,14 @@ public interface FunctionParameter
|
|||
* @param keyspace the current keyspace
|
||||
* @param arg a parameter value in a specific function call
|
||||
* @param receiverType the type of the object that will receive the result of the function call
|
||||
* @param inferredTypes the types that have been inferred for the other parameters
|
||||
* @return the inferred data type of the parameter, or {@link null} it isn't possible to infer it
|
||||
*/
|
||||
@Nullable
|
||||
default AbstractType<?> inferType(String keyspace,
|
||||
AssignmentTestable arg,
|
||||
@Nullable AbstractType<?> receiverType,
|
||||
List<AbstractType<?>> previousTypes)
|
||||
@Nullable List<AbstractType<?>> inferredTypes)
|
||||
{
|
||||
return arg.getCompatibleTypeIfKnown(keyspace);
|
||||
}
|
||||
|
|
@ -83,9 +86,9 @@ public interface FunctionParameter
|
|||
public AbstractType<?> inferType(String keyspace,
|
||||
AssignmentTestable arg,
|
||||
@Nullable AbstractType<?> receiverType,
|
||||
List<AbstractType<?>> previousTypes)
|
||||
@Nullable List<AbstractType<?>> inferredTypes)
|
||||
{
|
||||
return wrapped.inferType(keyspace, arg, receiverType, previousTypes);
|
||||
return wrapped.inferType(keyspace, arg, receiverType, inferredTypes);
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -130,7 +133,7 @@ public interface FunctionParameter
|
|||
public AbstractType<?> inferType(String keyspace,
|
||||
AssignmentTestable arg,
|
||||
@Nullable AbstractType<?> receiverType,
|
||||
List<AbstractType<?>> previousTypes)
|
||||
@Nullable List<AbstractType<?>> inferredTypes)
|
||||
{
|
||||
AbstractType<?> inferred = arg.getCompatibleTypeIfKnown(keyspace);
|
||||
return inferred != null ? inferred : types[0].getType();
|
||||
|
|
@ -168,7 +171,7 @@ public interface FunctionParameter
|
|||
public AbstractType<?> inferType(String keyspace,
|
||||
AssignmentTestable arg,
|
||||
@Nullable AbstractType<?> receiverType,
|
||||
List<AbstractType<?>> previousTypes)
|
||||
@Nullable List<AbstractType<?>> inferredTypes)
|
||||
{
|
||||
AbstractType<?> type = arg.getCompatibleTypeIfKnown(keyspace);
|
||||
return type == null && inferFromReceiver ? receiverType : type;
|
||||
|
|
@ -189,9 +192,12 @@ public interface FunctionParameter
|
|||
}
|
||||
|
||||
/**
|
||||
* @return a function parameter definition that accepts values with the same type as the first parameter
|
||||
* @param index the index of the function argument that this parameter is associated with
|
||||
* @param preferOther whether the parameter should prefer the type of the other parameter over its own type
|
||||
* @param parameter the type of this parameter when the type of the associated parameter is unknown
|
||||
* @return a function parameter definition that is expected to have the same type as another parameter
|
||||
*/
|
||||
static FunctionParameter sameAsFirst()
|
||||
static FunctionParameter sameAs(int index, boolean preferOther, FunctionParameter parameter)
|
||||
{
|
||||
return new FunctionParameter()
|
||||
{
|
||||
|
|
@ -199,21 +205,28 @@ public interface FunctionParameter
|
|||
public AbstractType<?> inferType(String keyspace,
|
||||
AssignmentTestable arg,
|
||||
@Nullable AbstractType<?> receiverType,
|
||||
List<AbstractType<?>> previousTypes)
|
||||
@Nullable List<AbstractType<?>> inferredTypes)
|
||||
{
|
||||
return previousTypes.get(0);
|
||||
if (preferOther)
|
||||
{
|
||||
AbstractType<?> other = inferredTypes == null ? null : inferredTypes.get(index);
|
||||
return other == null ? parameter.inferType(keyspace, arg, receiverType, inferredTypes) : other;
|
||||
}
|
||||
|
||||
AbstractType<?> inferred = parameter.inferType(keyspace, arg, receiverType, inferredTypes);
|
||||
return inferred == null && inferredTypes != null ? inferredTypes.get(index) : inferred;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void validateType(FunctionName name, AssignmentTestable arg, AbstractType<?> argType)
|
||||
{
|
||||
// nothing to do here, all types are accepted
|
||||
parameter.validateType(name, arg, argType);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString()
|
||||
{
|
||||
return "same";
|
||||
return parameter.toString();
|
||||
}
|
||||
};
|
||||
}
|
||||
|
|
@ -336,4 +349,55 @@ public interface FunctionParameter
|
|||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* @param type the type of the vector elements
|
||||
* @return a function parameter definition that accepts values of type {@link VectorType} with elements of the
|
||||
* specified {@code type} and any dimensions.
|
||||
*/
|
||||
static FunctionParameter vector(CQL3Type type)
|
||||
{
|
||||
return new FunctionParameter()
|
||||
{
|
||||
@Override
|
||||
public AbstractType<?> inferType(String keyspace,
|
||||
AssignmentTestable arg,
|
||||
@Nullable AbstractType<?> receiverType,
|
||||
@Nullable List<AbstractType<?>> inferredTypes)
|
||||
{
|
||||
if (arg instanceof Selectable.WithArrayLiteral)
|
||||
return VectorType.getInstance(type.getType(), ((Selectable.WithArrayLiteral) arg).getSize());
|
||||
|
||||
AbstractType<?> inferred = arg.getCompatibleTypeIfKnown(keyspace);
|
||||
return inferred == null ? receiverType : inferred;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void validateType(FunctionName name, AssignmentTestable arg, AbstractType<?> argType)
|
||||
{
|
||||
if (argType.isVector())
|
||||
{
|
||||
VectorType<?> vectorType = (VectorType<?>) argType;
|
||||
if (vectorType.elementType.asCQL3Type() == type)
|
||||
return;
|
||||
}
|
||||
else if (argType instanceof ListType) // if it's terminal it will be a list
|
||||
{
|
||||
ListType<?> listType = (ListType<?>) argType;
|
||||
if (listType.getElementsType().testAssignment(type.getType()) == NOT_ASSIGNABLE)
|
||||
return;
|
||||
}
|
||||
|
||||
throw new InvalidRequestException(format("Function %s requires a %s vector argument, " +
|
||||
"but found argument %s of type %s",
|
||||
name, type, arg, argType.asCQL3Type()));
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString()
|
||||
{
|
||||
return format("vector<%s, n>", type);
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ public class NativeFunctions
|
|||
BytesConversionFcts.addFunctionsTo(this);
|
||||
MathFcts.addFunctionsTo(this);
|
||||
MaskingFcts.addFunctionsTo(this);
|
||||
VectorFcts.addFunctionsTo(this);
|
||||
}
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,105 @@
|
|||
/*
|
||||
* 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;
|
||||
|
||||
import java.nio.ByteBuffer;
|
||||
import java.util.List;
|
||||
|
||||
import org.apache.cassandra.cql3.CQL3Type;
|
||||
import org.apache.cassandra.db.marshal.AbstractType;
|
||||
import org.apache.cassandra.db.marshal.FloatType;
|
||||
import org.apache.cassandra.db.marshal.VectorType;
|
||||
import org.apache.cassandra.exceptions.InvalidRequestException;
|
||||
import org.apache.cassandra.transport.ProtocolVersion;
|
||||
import org.apache.lucene.index.VectorSimilarityFunction;
|
||||
|
||||
public class VectorFcts
|
||||
{
|
||||
public static void addFunctionsTo(NativeFunctions functions)
|
||||
{
|
||||
functions.add(createSimilarityFunctionFactory("similarity_cosine", VectorSimilarityFunction.COSINE, false));
|
||||
functions.add(createSimilarityFunctionFactory("similarity_euclidean", VectorSimilarityFunction.EUCLIDEAN, true));
|
||||
functions.add(createSimilarityFunctionFactory("similarity_dot_product", VectorSimilarityFunction.DOT_PRODUCT, true));
|
||||
}
|
||||
|
||||
private static FunctionFactory createSimilarityFunctionFactory(String name,
|
||||
VectorSimilarityFunction luceneFunction,
|
||||
boolean supportsZeroVectors)
|
||||
{
|
||||
return new FunctionFactory(name,
|
||||
FunctionParameter.sameAs(1, false, FunctionParameter.vector(CQL3Type.Native.FLOAT)),
|
||||
FunctionParameter.sameAs(0, false, FunctionParameter.vector(CQL3Type.Native.FLOAT)))
|
||||
{
|
||||
@Override
|
||||
@SuppressWarnings("unchecked")
|
||||
protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
|
||||
{
|
||||
// check that all arguments have the same vector dimensions
|
||||
VectorType<Float> firstArgType = (VectorType<Float>) argTypes.get(0);
|
||||
int dimensions = firstArgType.dimension;
|
||||
if (!argTypes.stream().allMatch(t -> ((VectorType<?>) t).dimension == dimensions))
|
||||
throw new InvalidRequestException("All arguments must have the same vector dimensions");
|
||||
return createSimilarityFunction(name.name, firstArgType, luceneFunction, supportsZeroVectors);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private static NativeFunction createSimilarityFunction(String name,
|
||||
VectorType<Float> type,
|
||||
VectorSimilarityFunction f,
|
||||
boolean supportsZeroVectors)
|
||||
{
|
||||
return new NativeScalarFunction(name, FloatType.instance, type, type)
|
||||
{
|
||||
@Override
|
||||
public Arguments newArguments(ProtocolVersion version)
|
||||
{
|
||||
return new FunctionArguments(version,
|
||||
(v, b) -> type.composeAsFloat(b),
|
||||
(v, b) -> type.composeAsFloat(b));
|
||||
}
|
||||
|
||||
@Override
|
||||
public ByteBuffer execute(Arguments arguments) throws InvalidRequestException
|
||||
{
|
||||
if (arguments.containsNulls())
|
||||
return null;
|
||||
|
||||
float[] v1 = arguments.get(0);
|
||||
float[] v2 = arguments.get(1);
|
||||
|
||||
if (!supportsZeroVectors)
|
||||
{
|
||||
if (isAllZero(v1) || isAllZero(v2))
|
||||
throw new InvalidRequestException("Function " + name + " doesn't support all-zero vectors.");
|
||||
}
|
||||
|
||||
return FloatType.instance.decompose(f.compare(v1, v2));
|
||||
}
|
||||
|
||||
private boolean isAllZero(float[] v)
|
||||
{
|
||||
for (float f : v)
|
||||
if (f != 0)
|
||||
return false;
|
||||
return true;
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
@ -62,7 +62,9 @@ public class ReplaceMaskingFunction extends MaskingFunction
|
|||
/** @return a {@link FunctionFactory} to build new {@link ReplaceMaskingFunction}s. */
|
||||
public static FunctionFactory factory()
|
||||
{
|
||||
return new MaskingFunction.Factory(NAME, FunctionParameter.anyType(true), FunctionParameter.sameAsFirst())
|
||||
return new MaskingFunction.Factory(NAME,
|
||||
FunctionParameter.anyType(true),
|
||||
FunctionParameter.sameAs(0, true, FunctionParameter.anyType(true)))
|
||||
{
|
||||
@Override
|
||||
protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ public class VectorType extends DataType
|
|||
VectorType(DataType subtype, int dimensions)
|
||||
{
|
||||
super(Name.VECTOR);
|
||||
assert dimensions > 0 : "vectors may only have positive dimentions; given " + dimensions;
|
||||
assert dimensions > 0 : "vectors may only have positive dimensions; given " + dimensions;
|
||||
this.subtype = subtype;
|
||||
this.dimensions = dimensions;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -726,7 +726,7 @@ public interface Selectable extends AssignmentTestable
|
|||
/**
|
||||
* The list elements
|
||||
*/
|
||||
private final List<Selectable> selectables;
|
||||
protected final List<Selectable> selectables;
|
||||
|
||||
public WithArrayLiteral(List<Selectable> selectables)
|
||||
{
|
||||
|
|
@ -787,6 +787,11 @@ public interface Selectable extends AssignmentTestable
|
|||
return Lists.listToString(selectables);
|
||||
}
|
||||
|
||||
public int getSize()
|
||||
{
|
||||
return selectables.size();
|
||||
}
|
||||
|
||||
public static class Raw implements Selectable.Raw
|
||||
{
|
||||
private final List<Selectable.Raw> raws;
|
||||
|
|
@ -804,16 +809,11 @@ public interface Selectable extends AssignmentTestable
|
|||
}
|
||||
}
|
||||
|
||||
public static class WithList implements Selectable
|
||||
public static class WithList extends WithArrayLiteral
|
||||
{
|
||||
/**
|
||||
* The list elements
|
||||
*/
|
||||
private final List<Selectable> selectables;
|
||||
|
||||
public WithList(List<Selectable> selectables)
|
||||
{
|
||||
this.selectables = selectables;
|
||||
super(selectables);
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -876,16 +876,11 @@ public interface Selectable extends AssignmentTestable
|
|||
}
|
||||
}
|
||||
|
||||
public static class WithVector implements Selectable
|
||||
public static class WithVector extends WithArrayLiteral
|
||||
{
|
||||
/**
|
||||
* The vector elements
|
||||
*/
|
||||
private final List<Selectable> selectables;
|
||||
|
||||
public WithVector(List<Selectable> selectables)
|
||||
{
|
||||
this.selectables = selectables;
|
||||
super(selectables);
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
|
|||
|
|
@ -84,7 +84,7 @@ public final class VectorType<T> extends AbstractType<List<T>>
|
|||
{
|
||||
super(ComparisonType.CUSTOM);
|
||||
if (dimension <= 0)
|
||||
throw new InvalidRequestException(String.format("vectors may only have positive dimentions; given %d", dimension));
|
||||
throw new InvalidRequestException(String.format("vectors may only have positive dimensions; given %d", dimension));
|
||||
this.elementType = elementType;
|
||||
this.dimension = dimension;
|
||||
this.elementSerializer = elementType.getSerializer();
|
||||
|
|
|
|||
|
|
@ -0,0 +1,179 @@
|
|||
/*
|
||||
* 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;
|
||||
|
||||
import java.util.Arrays;
|
||||
import java.util.Collection;
|
||||
|
||||
import org.apache.commons.lang3.ArrayUtils;
|
||||
import org.junit.Test;
|
||||
import org.junit.runner.RunWith;
|
||||
import org.junit.runners.Parameterized;
|
||||
|
||||
import org.apache.cassandra.cql3.CQLTester;
|
||||
import org.apache.lucene.index.VectorSimilarityFunction;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
@RunWith(Parameterized.class)
|
||||
public class VectorFctsTest extends CQLTester
|
||||
{
|
||||
@Parameterized.Parameter
|
||||
public String function;
|
||||
|
||||
@Parameterized.Parameter(1)
|
||||
public VectorSimilarityFunction luceneFunction;
|
||||
|
||||
@Parameterized.Parameters(name = "{index}: function={0}")
|
||||
public static Collection<Object[]> data()
|
||||
{
|
||||
return Arrays.asList(new Object[][]{
|
||||
{ "system.similarity_cosine", VectorSimilarityFunction.COSINE },
|
||||
{ "system.similarity_euclidean", VectorSimilarityFunction.EUCLIDEAN },
|
||||
{ "system.similarity_dot_product", VectorSimilarityFunction.DOT_PRODUCT }
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testVectorSimilarityFunction()
|
||||
{
|
||||
createTable(KEYSPACE, "CREATE TABLE %s (pk int PRIMARY KEY, value vector<float, 2>, " +
|
||||
"l list<float>, " + // lists shouldn't be accepted by the functions
|
||||
"fl frozen<list<float>>, " + // frozen lists shouldn't be accepted by the functions
|
||||
"v1 vector<float, 1>, " + // 1-dimension vector to test missmatching dimensions
|
||||
"v_int vector<int, 2>, " + // int vectors shouldn't be accepted by the functions
|
||||
"v_double vector<double, 2>)");// double vectors shouldn't be accepted by the functions
|
||||
|
||||
float[] values = new float[]{ 1f, 2f };
|
||||
CQLTester.Vector<Float> vector = vector(ArrayUtils.toObject(values));
|
||||
Object[] similarity = row(luceneFunction.compare(values, values));
|
||||
|
||||
// basic functionality
|
||||
execute("INSERT INTO %s (pk, value, l, fl, v1, v_int, v_double) VALUES (0, ?, ?, ?, ?, ?, ?)",
|
||||
vector, list(1f, 2f), list(1f, 2f), vector(1f), vector(1, 2), vector(1d, 2d));
|
||||
assertRows(execute("SELECT " + function + "(value, value) FROM %s"), similarity);
|
||||
|
||||
// literals
|
||||
assertRows(execute("SELECT " + function + "(value, [1, 2]) FROM %s"), similarity);
|
||||
assertRows(execute("SELECT " + function + "([1, 2], value) FROM %s"), similarity);
|
||||
assertRows(execute("SELECT " + function + "([1, 2], [1, 2]) FROM %s"), similarity);
|
||||
|
||||
// bind markers
|
||||
assertRows(execute("SELECT " + function + "(value, ?) FROM %s", vector), similarity);
|
||||
assertRows(execute("SELECT " + function + "(?, value) FROM %s", vector), similarity);
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(?, ?) FROM %s", vector, vector))
|
||||
.hasMessageContaining("Cannot infer type of argument ?");
|
||||
|
||||
// bind markers with type hints
|
||||
assertRows(execute("SELECT " + function + "((vector<float, 2>) ?, ?) FROM %s", vector, vector), similarity);
|
||||
assertRows(execute("SELECT " + function + "(?, (vector<float, 2>) ?) FROM %s", vector, vector), similarity);
|
||||
assertRows(execute("SELECT " + function + "((vector<float, 2>) ?, (vector<float, 2>) ?) FROM %s", vector, vector), similarity);
|
||||
|
||||
// bind markers and literals
|
||||
assertRows(execute("SELECT " + function + "([1, 2], ?) FROM %s", vector), similarity);
|
||||
assertRows(execute("SELECT " + function + "(?, [1, 2]) FROM %s", vector), similarity);
|
||||
assertRows(execute("SELECT " + function + "([1, 2], ?) FROM %s", vector), similarity);
|
||||
|
||||
// wrong column types with columns
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(l, value) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument l of type list<float>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(fl, value) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument fl of type frozen<list<float>>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(value, l) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument l of type list<float>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(value, fl) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument fl of type frozen<list<float>>");
|
||||
|
||||
// wrong column types with columns and literals
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(l, [1, 2]) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument l of type list<float>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(fl, [1, 2]) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument fl of type frozen<list<float>>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "([1, 2], l) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument l of type list<float>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "([1, 2], fl) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument fl of type frozen<list<float>>");
|
||||
|
||||
// wrong column types with cast literals
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "((List<Float>)[1, 2], [3, 4]) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument (list<float>)[1, 2] of type frozen<list<float>>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "((List<Float>)[1, 2], (List<Float>)[3, 4]) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument (list<float>)[1, 2] of type frozen<list<float>>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "([1, 2], (List<Float>)[3, 4]) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument (list<float>)[3, 4] of type frozen<list<float>>");
|
||||
|
||||
// wrong non-float vectors
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(v_int, [1, 2]) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument v_int of type vector<int, 2>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(v_double, [1, 2]) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument v_double of type vector<double, 2>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "([1, 2], v_int) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument v_int of type vector<int, 2>");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "([1, 2], v_double) FROM %s"))
|
||||
.hasMessageContaining("Function " + function + " requires a float vector argument, but found argument v_double of type vector<double, 2>");
|
||||
|
||||
// mismatching dimensions with literals
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "([1, 2], [3]) FROM %s", vector(1)))
|
||||
.hasMessageContaining("All arguments must have the same vector dimensions");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(value, [1]) FROM %s", vector(1)))
|
||||
.hasMessageContaining("All arguments must have the same vector dimensions");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "([1], value) FROM %s", vector(1)))
|
||||
.hasMessageContaining("All arguments must have the same vector dimensions");
|
||||
|
||||
// mismatching dimensions with bind markers
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "((vector<float, 1>) ?, value) FROM %s", vector(1)))
|
||||
.hasMessageContaining("All arguments must have the same vector dimensions");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(value, (vector<float, 1>) ?) FROM %s", vector(1)))
|
||||
.hasMessageContaining("All arguments must have the same vector dimensions");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "((vector<float, 2>) ?, (vector<float, 1>) ?) FROM %s", vector(1, 2), vector(1)))
|
||||
.hasMessageContaining("All arguments must have the same vector dimensions");
|
||||
|
||||
// mismatching dimensions with columns
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(value, v1) FROM %s"))
|
||||
.hasMessageContaining("All arguments must have the same vector dimensions");
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(v1, value) FROM %s"))
|
||||
.hasMessageContaining("All arguments must have the same vector dimensions");
|
||||
|
||||
// null arguments with literals
|
||||
assertRows(execute("SELECT " + function + "(value, null) FROM %s"), row((Float) null));
|
||||
assertRows(execute("SELECT " + function + "(null, value) FROM %s"), row((Float) null));
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(null, null) FROM %s"))
|
||||
.hasMessageContaining("Cannot infer type of argument NULL in call to function " + function);
|
||||
|
||||
// null arguments with bind markers
|
||||
assertRows(execute("SELECT " + function + "(value, ?) FROM %s", (CQLTester.Vector<Float>) null), row((Float) null));
|
||||
assertRows(execute("SELECT " + function + "(?, value) FROM %s", (CQLTester.Vector<Float>) null), row((Float) null));
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(?, ?) FROM %s", null, null))
|
||||
.hasMessageContaining("Cannot infer type of argument ? in call to function " + function);
|
||||
|
||||
// test all-zero vectors, only cosine similarity should reject them
|
||||
if (luceneFunction == VectorSimilarityFunction.COSINE)
|
||||
{
|
||||
String expected = "Function " + function + " doesn't support all-zero vectors";
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "(value, [0, 0]) FROM %s")) .hasMessageContaining(expected);
|
||||
assertThatThrownBy(() -> execute("SELECT " + function + "([0, 0], value) FROM %s")).hasMessageContaining(expected);
|
||||
}
|
||||
else
|
||||
{
|
||||
float expected = luceneFunction.compare(values, new float[]{ 0, 0 });
|
||||
assertRows(execute("SELECT " + function + "(value, [0, 0]) FROM %s"), row(expected));
|
||||
assertRows(execute("SELECT " + function + "([0, 0], value) FROM %s"), row(expected));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -457,14 +457,14 @@ public class CQLVectorTest extends CQLTester.InMemory
|
|||
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",
|
||||
assertInvalidThrowMessage("vectors may only have positive dimensions; 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",
|
||||
assertInvalidThrowMessage("vectors may only have positive dimensions; given 0",
|
||||
InvalidRequestException.class,
|
||||
"CREATE FUNCTION %s (x vector<int, 2>) " +
|
||||
"CALLED ON NULL INPUT " +
|
||||
|
|
|
|||
|
|
@ -520,15 +520,15 @@ public final class AbstractTypeGenerators
|
|||
return vectorTypeGen(typeGen, SourceDSL.integers().between(1, 100));
|
||||
}
|
||||
|
||||
public static Gen<VectorType<?>> vectorTypeGen(Gen<AbstractType<?>> typeGen, Gen<Integer> dimentionGen)
|
||||
public static Gen<VectorType<?>> vectorTypeGen(Gen<AbstractType<?>> typeGen, Gen<Integer> dimensionGen)
|
||||
{
|
||||
return rnd -> {
|
||||
int dimention = dimentionGen.generate(rnd);
|
||||
int dimension = dimensionGen.generate(rnd);
|
||||
AbstractType<?> element = typeGen.generate(rnd);
|
||||
// empty type not supported
|
||||
while (element == EmptyType.instance)
|
||||
element = typeGen.generate(rnd);
|
||||
return VectorType.getInstance(element.freeze(), dimention);
|
||||
return VectorType.getInstance(element.freeze(), dimension);
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue