Add support for Storage-attached indexes via @SaiIndexed.

Closes #1505
This commit is contained in:
Mark Paluch
2024-08-15 15:31:47 +02:00
parent b53b9bb91f
commit 4bb860fd87
6 changed files with 268 additions and 36 deletions

View File

@@ -24,12 +24,14 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.data.cassandra.core.cql.keyspace.CreateIndexSpecification;
import org.springframework.data.cassandra.core.cql.keyspace.SpecificationBuilder;
import org.springframework.data.cassandra.core.mapping.CassandraPersistentProperty;
import org.springframework.data.cassandra.core.mapping.Indexed;
import org.springframework.data.cassandra.core.mapping.SAIIndexed;
import org.springframework.data.cassandra.core.mapping.SASI;
import org.springframework.data.cassandra.core.mapping.SASI.NonTokenizingAnalyzed;
import org.springframework.data.cassandra.core.mapping.SASI.Normalization;
@@ -50,6 +52,7 @@ import com.datastax.oss.driver.api.core.CqlIdentifier;
* @since 2.0
* @see Indexed
* @see SASI
* @see SAIIndexed
*/
@SuppressWarnings("unchecked")
class IndexSpecificationFactory {
@@ -92,45 +95,67 @@ class IndexSpecificationFactory {
}
if (property.isAnnotationPresent(SASI.class)) {
indexes.add(createIndexSpecification(keyspace, property.findAnnotation(SASI.class), property));
indexes.add(createIndexSpecification(keyspace, property.getRequiredAnnotation(SASI.class), property));
}
if (property.isAnnotationPresent(SAIIndexed.class)) {
CreateIndexSpecification index = createIndexSpecification(keyspace,
property.getRequiredAnnotation(SAIIndexed.class), property);
if (property.isMapLike()) {
index.entries();
}
indexes.add(index);
}
if (property.isMapLike()) {
AnnotatedType type = property.findAnnotatedType(Indexed.class);
if (type instanceof AnnotatedParameterizedType) {
AnnotatedParameterizedType parameterizedType = (AnnotatedParameterizedType) type;
AnnotatedType[] typeArgs = parameterizedType.getAnnotatedActualTypeArguments();
Indexed keyIndex = typeArgs.length == 2 ? AnnotatedElementUtils.getMergedAnnotation(typeArgs[0], Indexed.class)
: null;
Indexed valueIndex = typeArgs.length == 2
? AnnotatedElementUtils.getMergedAnnotation(typeArgs[1], Indexed.class)
: null;
if ((!indexes.isEmpty() && (keyIndex != null || valueIndex != null))
|| (keyIndex != null && valueIndex != null)) {
throw new MappingException("Multiple index declarations for " + property
+ " found; A map index must be either declared for entries, keys or values");
}
if (keyIndex != null) {
indexes.add(createIndexSpecification(keyspace, keyIndex, property).keys());
}
if (valueIndex != null) {
indexes.add(createIndexSpecification(keyspace, valueIndex, property).values());
}
}
indexes.addAll(createTypeAnnotatedIndexes(Indexed.class, property,
indexed -> createIndexSpecification(keyspace, indexed, property)));
indexes.addAll(createTypeAnnotatedIndexes(SAIIndexed.class, property,
indexed -> createIndexSpecification(keyspace, indexed, property)));
}
return indexes;
}
private static <T extends Annotation> List<CreateIndexSpecification> createTypeAnnotatedIndexes(
Class<T> annotationType, CassandraPersistentProperty property,
Function<T, CreateIndexSpecification> indexFunction) {
AnnotatedType type = property.findAnnotatedType(annotationType);
if (type instanceof AnnotatedParameterizedType parameterizedType) {
List<CreateIndexSpecification> indexes = new ArrayList<>(2);
AnnotatedType[] typeArgs = parameterizedType.getAnnotatedActualTypeArguments();
T keyIndex = typeArgs.length == 2 ? AnnotatedElementUtils.getMergedAnnotation(typeArgs[0], annotationType) : null;
T valueIndex = typeArgs.length == 2 ? AnnotatedElementUtils.getMergedAnnotation(typeArgs[1], annotationType)
: null;
if (keyIndex != null && valueIndex != null) {
throw new MappingException("Multiple index declarations for " + property
+ " found; A map index must be either declared for entries, keys or values");
}
if (keyIndex != null) {
indexes.add(indexFunction.apply(keyIndex).keys());
}
if (valueIndex != null) {
indexes.add(indexFunction.apply(valueIndex).values());
}
return indexes;
}
return Collections.emptyList();
}
static CreateIndexSpecification createIndexSpecification(@Nullable CqlIdentifier keyspace, Indexed annotation,
CassandraPersistentProperty property) {
@@ -181,6 +206,27 @@ class IndexSpecificationFactory {
return index;
}
private static CreateIndexSpecification createIndexSpecification(@Nullable CqlIdentifier keyspace,
SAIIndexed annotation, CassandraPersistentProperty property) {
CreateIndexSpecification index;
if (StringUtils.hasText(annotation.value())) {
index = SpecificationBuilder.createIndex(keyspace, CqlIdentifier.fromCql(annotation.value()));
} else {
index = SpecificationBuilder.createIndex(keyspace, null);
}
index.using("sai") //
.columnName(property.getRequiredColumnName())
.withOption("case_sensitive", Boolean.toString(annotation.caseSensitive()))
.withOption("normalize", Boolean.toString(annotation.normalize()))
.withOption("ascii", Boolean.toString(annotation.ascii()))
.withOption("similarity_function", annotation.similarityFunction().name());
return index;
}
interface CreateIndexConfigurer<T extends Annotation> extends BiConsumer<T, CreateIndexSpecification> {}
enum StandardAnalyzedConfigurer implements CreateIndexConfigurer<StandardAnalyzed> {

View File

@@ -61,7 +61,7 @@ public class CreateIndexCqlGenerator extends IndexNameCqlGenerator<CreateIndexSp
cql.append(")");
if (spec().isCustom()) {
if (StringUtils.hasText(spec().getUsing())) {
cql.append(" USING ").append("'").append(spec().getUsing()).append("'");
}

View File

@@ -196,7 +196,10 @@ public class CreateIndexSpecification extends IndexNameSpecification<CreateIndex
if (StringUtils.hasText(className)) {
this.using = className;
this.custom = true;
if (!"sai".equalsIgnoreCase(className)) {
this.custom = true;
}
} else {
this.using = null;
this.custom = false;

View File

@@ -0,0 +1,97 @@
/*
* Copyright 2024 the original author or authors.
*
* Licensed 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
*
* https://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.springframework.data.cassandra.core.mapping;
import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
/**
* Identifies a Storage-Attached Indexing (SAI) index. You can create multiple secondary indexes on the same database
* table, with each SAI index based on any column in the table. All column date types except the following are supported
* for SAI indexes:
* <ul>
* <li>counter</li>
* <li>non-frozen user-defined type (UDT)</li>
* </ul>
* <p>
* The following columns of a {@link Table} type can be annotated with {@link SAIIndexed}:
* <ul>
* <li>Scalar data types</li>
* <li>Frozen user-defined types</li>
* <li>Collection types</li>
* <li>Map type</li>
* </ul>
* <p>
* Map types distinguish between entry, key or value-level indexing.
*
* <pre class="code">
* &#64;Table
* class Person {
*
* Map&lt;@SAIIndexed String, String&gt; indexedKey; // allows CONTAINS KEY queries
* Map&lt;String, @SAIIndexed String&gt; indexedValue; // allows CONTAINS queries
* }
* </pre>
*
* @author Mark Paluch
*/
@Documented
@Retention(value = RetentionPolicy.RUNTIME)
@Target(value = { ElementType.FIELD, ElementType.METHOD, ElementType.ANNOTATION_TYPE, ElementType.TYPE_USE })
public @interface SAIIndexed {
/**
* The name of the index. If {@literal null} or empty, then the index name will be generated by Cassandra and will be
* unknown unless column metadata is used to discover the generated index name.
*/
String value() default "";
/**
* Enable case-sensitive matching.
*/
boolean caseSensitive() default true;
/**
* When set to true, perform Unicode normalization on indexed strings. SAI supports Normalization Form C (NFC)
* Unicode. When set to true, SAI normalizes the different versions of a given Unicode character to a single version,
* retaining all the marks and symbols in the index. For example, SAI would change the character {@code Å (U+212B)} to
* {@code Å (U+00C5)}.
*/
boolean normalize() default false;
/**
* When set to true, SAI converts alphabetic, numeric, and symbolic characters that are not in the Basic Latin Unicode
* block (the first 127 ASCII characters) to the ASCII equivalent, if one exists. For example, this option changes
* {@code à} to {@code a}.
*/
boolean ascii() default false;
/**
* Vector search relies on computing the similarity or distance between vectors to identify relevant matches. The
* similarity function is used to compute the similarity between two vectors.
*/
SimilarityFunction similarityFunction() default SimilarityFunction.COSINE;
/**
* Enumeration of similarity functions.
*/
enum SimilarityFunction {
COSINE, DOT_PRODUCT, EUCLIDEAN
}
}

View File

@@ -30,6 +30,7 @@ import org.springframework.data.cassandra.core.mapping.BasicCassandraPersistentE
import org.springframework.data.cassandra.core.mapping.CassandraMappingContext;
import org.springframework.data.cassandra.core.mapping.Indexed;
import org.springframework.data.cassandra.core.mapping.PrimaryKeyColumn;
import org.springframework.data.cassandra.core.mapping.SAIIndexed;
import org.springframework.data.cassandra.core.mapping.SASI;
import org.springframework.data.cassandra.core.mapping.SASI.NonTokenizingAnalyzed;
import org.springframework.data.cassandra.core.mapping.SASI.Normalization;
@@ -86,6 +87,45 @@ class IndexSpecificationFactoryUnitTests {
assertThat(entries.getColumnFunction()).isEqualTo(ColumnFunction.VALUES);
}
@Test // GH-1505
void createSaiIndexShouldCreateCreateIndexSpecification() {
CreateIndexSpecification simpleSasi = createIndexFor(IndexedType.class, "simpleSai");
assertThat(simpleSasi.getColumnName()).isEqualTo(CqlIdentifier.fromInternal("simplesai"));
assertThat(simpleSasi.getTableName()).isNull();
assertThat(simpleSasi.isCustom()).isFalse();
assertThat(simpleSasi.getUsing()).isEqualTo("sai");
assertThat(simpleSasi.getColumnFunction()).isEqualTo(ColumnFunction.NONE);
assertThat(simpleSasi.getOptions()).containsEntry("case_sensitive", "true").containsEntry("normalize", "false")
.containsEntry("ascii", "false").containsEntry("similarity_function", "COSINE");
}
@Test // GH-1505
void createSaiIndexShouldApplyIndexOptions() {
CreateIndexSpecification simpleSasi = createIndexFor(IndexedType.class, "customSai");
assertThat(simpleSasi.getName()).isEqualTo(CqlIdentifier.fromInternal("foo"));
assertThat(simpleSasi.getColumnName()).isEqualTo(CqlIdentifier.fromInternal("customsai"));
assertThat(simpleSasi.getTableName()).isNull();
assertThat(simpleSasi.isCustom()).isFalse();
assertThat(simpleSasi.getUsing()).isEqualTo("sai");
assertThat(simpleSasi.getColumnFunction()).isEqualTo(ColumnFunction.NONE);
assertThat(simpleSasi.getOptions()).containsEntry("case_sensitive", "false").containsEntry("normalize", "true")
.containsEntry("ascii", "true").containsEntry("similarity_function", "EUCLIDEAN");
}
@Test // GH-1505
void createSaiIndexMapKeyShouldCreateCreateIndexSpecification() {
CreateIndexSpecification simpleSasi = createIndexFor(SaiIndexedMapKeyProperty.class, "entries");
assertThat(simpleSasi.getColumnName()).isEqualTo(CqlIdentifier.fromInternal("entries"));
assertThat(simpleSasi.getUsing()).isEqualTo("sai");
assertThat(simpleSasi.getColumnFunction()).isEqualTo(ColumnFunction.KEYS);
}
@Test // DATACASS-306
void createIndexForSimpleSasiShouldApplyIndexOptions() {
@@ -155,12 +195,15 @@ class IndexSpecificationFactoryUnitTests {
}
private CreateIndexSpecification createIndexFor(Class<?> type, String property) {
return createIndexesFor(type, property).stream().findFirst().get();
}
private List<CreateIndexSpecification> createIndexesFor(Class<?> type, String property) {
BasicCassandraPersistentEntity<?> entity = mappingContext.getRequiredPersistentEntity(type);
return IndexSpecificationFactory
.createIndexSpecifications(entity.getKeyspace(), entity.getRequiredPersistentProperty(property)).stream()
.findFirst().get();
return IndexSpecificationFactory.createIndexSpecifications(entity.getKeyspace(),
entity.getRequiredPersistentProperty(property));
}
private static class IndexedType {
@@ -192,6 +235,11 @@ class IndexSpecificationFactoryUnitTests {
@SASI
@NonTokenizingAnalyzed(caseSensitive = false,
normalization = Normalization.LOWERCASE) String sasiNontokenizingLowercase;
@SAIIndexed String simpleSai;
@SAIIndexed(value = "foo", caseSensitive = false, normalize = true, ascii = true,
similarityFunction = SAIIndexed.SimilarityFunction.EUCLIDEAN) String customSai;
}
@AccessType(Type.PROPERTY)
@@ -204,6 +252,16 @@ class IndexSpecificationFactoryUnitTests {
public void setEntries(Map<String, String> entries) {}
}
@AccessType(Type.PROPERTY)
private static class SaiIndexedMapKeyProperty {
public Map<@SAIIndexed String, String> getEntries() {
return null;
}
public void setEntries(Map<String, String> entries) {}
}
@AccessType(Type.PROPERTY)
private static class MapValueIndexProperty {

View File

@@ -97,6 +97,34 @@ class CreateIndexCqlGeneratorUnitTests {
.tableName(CqlIdentifier.fromInternal("order"));
assertThat(CqlGenerator.toCql(spec)).isEqualTo("CREATE INDEX order_dob ON \"order\" (\"order\");");
}
@Test // GH-1505
void createSaiIndex() {
CreateIndexSpecification spec = SpecificationBuilder.createIndex("my_index")
.tableName("comments_vs").columnName("comment_vector").using("sai");
assertThat(CqlGenerator.toCql(spec)).isEqualTo("CREATE INDEX my_index ON comments_vs (comment_vector) USING 'sai';");
}
@Test // GH-1505
void createSaiIndexOnKeys() {
CreateIndexSpecification spec = SpecificationBuilder.createIndex("my_index")
.tableName("comments_vs").columnName("comment_vector").columnFunction(CreateIndexSpecification.ColumnFunction.KEYS).using("sai");
assertThat(CqlGenerator.toCql(spec)).isEqualTo("CREATE INDEX my_index ON comments_vs (KEYS(comment_vector)) USING 'sai';");
}
@Test // GH-1505
void createSaiIndexWithOptions() {
CreateIndexSpecification spec = SpecificationBuilder.createIndex("my_index")
.tableName("comments_vs").columnName("comment_vector").using("sai")
.withOption("similarity_function", "DOT_PRODUCT");
assertThat(CqlGenerator.toCql(spec)).isEqualTo("CREATE INDEX my_index ON comments_vs (comment_vector) USING 'sai' WITH OPTIONS = {'similarity_function': 'DOT_PRODUCT'};");
}
}